Commit 4fc648d2 authored by jameskrw's avatar jameskrw
Browse files

added prompt.py for sokoban env, minor debug

parent 3c405982
Loading
Loading
Loading
Loading
+7 −55
Original line number Diff line number Diff line
@@ -21,60 +21,12 @@ from vagen.env.utils import (
    preprocess,
    postprocess,
)
from vagen.env.sokoban.prompt import (
    init_observation_template,
    action_template,
    instruction_template,
)

system_prompt = """
You are a helpful assistant. You first think about the reasoning process in the mind and then provides the user with the answer.
"""

instruction_template = """You are a Sokoban solver.

Sokoban Quick Guide
Goal: Push all boxes (X) onto targets (O).

Symbols:
# Wall | _ Floor | O Target | X Box | P You | √ = Box on Target | S = You on Target
The observation is a 2D grid of the current state of the Sokoban game.

Rules:
1. Push boxes (can't pull).
2. Avoid walls (#).

Actions you can take: Up, Down, Left, Right. You can only take one action at a time.
Up: move up to the cell above (to the above row).
Down: move down to the cell below (to the below row).
Left: move left to the cell to the left (to the left column).
Right: move right to the cell to the right (to the right column).

Rewards:
Move: -0.1
Box on target: +1.0
All boxes placed: +10.0

Please think step by step and provide the action you want to take.
Your reponse should be in the format of <think>...</think><answer>...</answer>
E.g. <think> There's a box on the upper left of me, I need to push it upward. </think><answer> Right,Up,Up </answer>
"""
# Let's try to use a format reward and answer reward
# If the reponse provides a final answer and is corect, the model receives an accurtacy reward of +1
# is the response encloses its thinking in <think></think> and the final answer is <answer></answer> tags, the model receives a format reward of +1



init_observation_template = """
[Initial Observation]:
{observation}
Decide your next action.
Your reponse should be in the format of <think>...</think><answer>...</answer>
"""

action_template = """After you answer {answer}, the extracted valid action is {valid_action}.\
After that, the observation is:
{observation}
reward: {reward}
done: {done}
Decide your next action.
Your reponse should be in the format of <think>...</think><answer>...</answer>
"""


class SokobanEnv(BaseEnv, GymSokobanEnv):
@@ -195,10 +147,10 @@ class SokobanInterface(BaseInterface):

    INVALID_ACTION = 0
    FORMAT_REWARD = 0.5
    FORMAT_PENALTY = -0.5
    FORMAT_PENALTY = 0.0
    VALID_ACTION_REWARD = 0.5
    MAX_ACTION_PER_STEP = 1 # NOTE hard coded here
    MAX_ACTION_PENALTY = -0.1
    MAX_ACTION_PENALTY = 0.0
    ACTION_LOOKUP = {
        0: "None",
        1: "Up",
+53 −0
Original line number Diff line number Diff line
system_prompt = """
You are a helpful assistant. You first think about the reasoning process in the mind and then provides the user with the answer.
"""

instruction_template = """You are a Sokoban solver.

Sokoban Quick Guide
Goal: Push all boxes (X) onto targets (O).

Symbols:
# Wall | _ Floor | O Target | X Box | P You | √ = Box on Target | S = You on Target
The observation is a 2D grid of the current state of the Sokoban game.

Rules:
1. Push boxes (can't pull).
2. Avoid walls (#).

Actions you can take: Up, Down, Left, Right. You can only take one action at a time.
Up: move up to the cell above (to the above row).
Down: move down to the cell below (to the below row).
Left: move left to the cell to the left (to the left column).
Right: move right to the cell to the right (to the right column).

Rewards:
Box on target: +1.0
All boxes placed: +10.0
Format correct: +0.5

Please think step by step and provide the action you want to take.
Your reponse should be in the format of <think>...</think><answer>...</answer>
E.g. <think> There's a box on the upper right of me, the target is on the upper side of the box, I need to push the box it upward. </think><answer> Right,Up,Up </answer>
"""
# Let's try to use a format reward and answer reward
# If the reponse provides a final answer and is corect, the model receives an accurtacy reward of +1
# is the response encloses its thinking in <think></think> and the final answer is <answer></answer> tags, the model receives a format reward of +1



init_observation_template = """
[Initial Observation]:
{observation}
Decide your next action.
Your reponse should be in the format of <think>...</think><answer>...</answer>
"""

action_template = """After you answer {answer}, the extracted valid action is {valid_action}.\
After that, the observation is:
{observation}
reward: {reward}
done: {done}
Decide your next action.
Your reponse should be in the format of <think>...</think><answer>...</answer>
"""
 No newline at end of file
+2 −2
Original line number Diff line number Diff line
@@ -288,8 +288,8 @@ def compute_data_metrics(batch, use_critic=True):

    if "loss_mask" in batch.batch.keys():
        end_of_response_position_mask=batch.batch["end_of_response_position_mask"]
        response_length = (batch.batch['loss_mask'].sum(-1)/ end_of_response_position_mask.sum(-1)).mean().item()
        prompt_length = ((1-batch.batch['loss_mask']).sum(-1)/ end_of_response_position_mask.sum(-1)).mean().item()
        response_length = (batch.batch['loss_mask'].sum(-1)/ end_of_response_position_mask.sum(-1))
        prompt_length = ((1-batch.batch['loss_mask']).sum(-1)/ end_of_response_position_mask.sum(-1))
    else:
        response_info = _compute_response_info(batch)
        prompt_length = response_info['prompt_length']