Commit fd51b132 authored by jameskrw's avatar jameskrw
Browse files

updated thinking format

parent d3e4742c

vagen/env/frozenlake/env.py

+12 −12
Original line number Diff line number Diff line
@@ -5,8 +5,9 @@ from typing import Dict, List, Optional, Tuple, Any 
from gymnasium.utils import seeding

from gymnasium.envs.toy_text.frozen_lake import FrozenLakeEnv as GymFrozenLakeEnv

from vagen.env.utils.env_utils import NoLoggerWarnings, set_seed

from vagen.env.utils.context_utils import parse_llm_raw_response, convert_numpy_to_PIL

from .prompt import system_prompt_text, system_prompt_vision, init_observation_template, action_template

from vagen.env.utils.context_utils import convert_numpy_to_PIL

from vagen.env.utils.parse_utils import parse_function_map

from .prompt import system_prompt, init_observation_template, action_template,format_prompt

from .env_config import FrozenLakeEnvConfig

from .utils import generate_random_map, is_valid
@@ -73,6 +74,9 @@ class FrozenLakeEnv(BaseEnv): 
        self.total_reward = 0

        self.valid_actions = []

        self.reward = 0

        self.format_prompt = format_prompt[self.config.prompt_format].format(

            max_actions_per_step=self.config.max_actions_per_step,action_sep=self.config.action_sep)

        self.parse_func= parse_function_map[self.config.prompt_format.rstrip("_symbol")]



    def reset(self, seed=None):

        """
@@ -121,7 +125,7 @@ class FrozenLakeEnv(BaseEnv): 
                - info: Dictionary containing metrics and parsed action data

        """

        # Parse the LLM's raw response to extract actions

        rst = parse_llm_raw_response(

        rst = self.parse_func(

            response=action_str,

            special_token_list=self.config.special_token_list,

            action_sep=self.config.action_sep,
@@ -171,7 +175,7 @@ class FrozenLakeEnv(BaseEnv): 
                break

        

        # Add format reward if actions were valid

        if metrics["turn_metrics"]['action_is_valid']:

        if metrics["turn_metrics"]['action_is_valid'] and rst["format_correct"]:

            self.reward += self.config.format_reward

        

        # Add metrics to info dictionary
@@ -196,10 +200,8 @@ class FrozenLakeEnv(BaseEnv): 
        Returns:

            str: System prompt string with environment description and instructions

        """

        if self.config.render_mode == 'vision':

            return system_prompt_vision.format(max_actions_per_step=self.config.max_actions_per_step,action_sep=self.config.action_sep)

        else:

            return system_prompt_text.format(max_actions_per_step=self.config.max_actions_per_step,action_sep=self.config.action_sep)

        

        return system_prompt+'\n'+self.format_prompt



    def compute_reward(self):

        """
@@ -252,15 +254,13 @@ class FrozenLakeEnv(BaseEnv): 
        # Format the observation string using the appropriate template

        if init_obs:

            # Initial observation doesn't include action results

            obs_str = init_observation_template.format(observation=img_str)

            obs_str = init_observation_template.format(observation=img_str)+"\n"+ self.format_prompt

        else:

            # Subsequent observations include action results

            obs_str = action_template.format(

                valid_action=self.valid_actions,

                observation=img_str,

                reward=self.reward,

                done=self._finished(),

            )

            )+"\n"+ self.format_prompt

        

        # Return observation dictionary with appropriate fields

        if multi_modal_data is not None:

vagen/env/frozenlake/env_config.py

+4 −0
Original line number Diff line number Diff line
@@ -11,6 +11,10 @@ class FrozenLakeEnvConfig(BaseEnvConfig): 
    render_mode: str = "vision"  # "text" or "vision"

    max_actions_per_step: int = 3

    min_actions_to_succeed: int = 5

    prompt_format: str = "free_think" 

    # "free_think", "no_think", "grounding", "worldmodeling", "grounding_worldmodeling"

    # "grounding_symbol", "worldmodeling_symbol", "grounding_worldmodeling_symbol"

    use_accuracy_reward: bool = False

    

    def config_id(self) -> str:

        id_fields=["is_slippery", "size", "p", "render_mode", "max_actions_per_step", "min_actions_to_succeed","format_reward"]

vagen/env/frozenlake/prompt.py

+60 −45
Original line number Diff line number Diff line 
system_prompt_text = """You are a FrozenLake solver.

system_prompt= """You are a FrozenLake solver.



FrozenLake Quick Guide

Goal: Reach the goal (G).



Symbols:

Symbols (If image is provided there are no symbols):

_ Frozen | O Hole | G Goal | P Player | X Player fell into hole | √ Player on goal



Rules:

1. Avoid falling into holes (O).

1. Avoid falling into holes.

2. Frozen tiles are slippery, you may move perpendicular to your intended direction.



Actions you can take: Left, Down, Right, Up. 

You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

Left: move left to the cell to the left.

Down: move down to the cell below.

Right: move right to the cell to the right.

Up: move up to the cell above.



Rewards:

Fall into hole: 0

Reach goal: +10.0

Format correct: +0.5

"""



Please think step by step and provide the actions you want to take.

Your response should be in the format of <think>...</think><answer>...</answer>

init_observation_template = """[Initial Observation]:

{observation}

Decide your next action(s).

"""



system_prompt_vision = """You are a FrozenLake solver.

action_template = """After your answer, the extracted valid action is {valid_action}.

After that, the observation is:

{observation}

Decide your next action(s).

"""



FrozenLake Quick Guide

Goal: Reach the goal (G).

free_think_format_prompt= """You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

<think>...</think><answer>...</answer>

e.g. <think>I can see the target is on my down left, I should go down then left</think><answer>Down{action_sep}Left</answer>

"""



Symbols:

Light blue: Frozen surface | Black: Hole | Green: Goal | Red: Player

no_think_format_prompt= """You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

<answer>...</answer>

e.g. <answer>Down{action_sep}Left</answer>

"""



Rules:

1. Avoid falling into holes.

2. Frozen tiles are slippery, you may move perpendicular to your intended direction.

grounding_format_prompt= """You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

<current_state>...</current_state><think>...</think><answer>...</answer>

e.g. <current_state>I'm in the row 2 col 3. The target is in the row 3 col 2.</current_state><think>I should go down then left to reach the target</think><answer>Down{action_sep}Left</answer>

"""



Actions you can take: Left, Down, Right, Up. 

You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

worldmodeling_format_prompt= """You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

<think>...</think><answer>...</answer><next_state>...</next_state>

e.g. <think>I can see the target is on my down left, I should go down then left</think><answer>Down{action_sep}Left</answer><next_state>I'm in the row 3 col 2. The target is in the row 3 col 2.</next_state>

"""



Rewards:

Fall into hole: 0

Reach goal: +10.0

Format correct: +0.5

grounding_worldmodeling_format_prompt= """You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

<current_state>...</current_state><think>...</think><answer>...</answer><next_state>...</next_state>

e.g. <current_state>I'm in the row 2 col 3. The target is in the row 3 col 2.</current_state><think>I should go down then left to reach the target</think><answer>Down{action_sep}Left</answer><next_state>I'm in the row 3 col 2. The target is in the row 3 col 2.</next_state>

"""



Please think step by step and provide the actions you want to take.

Your response should be in the format of <think>...</think><answer>...</answer>

grounding_symbol_format_prompt= """You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

<current_state>...</current_state><think>...</think><answer>...</answer>

e.g. <current_state>_P__\nG___\n_OO_\n____</current_state><think>I should go down then left to reach the target</think><answer>Down{action_sep}Left</answer>

"""



init_observation_template = """

[Initial Observation]:

{observation}

Decide your next action(s).

Please think step by step and provide the actions you want to take.

Your response should be in the format of <think>...</think><answer>...</answer>

worldmodeling_format_symbol= """You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

<think>...</think><answer>...</answer><next_state>...</next_state>

e.g. <think>I can see the target is on my down left, I should go down then left</think><answer>Down{action_sep}Left</answer><next_state>____\n√___\n_OO_\n____</next_state>

"""



action_template = """After your answer, the extracted valid action is {valid_action}.

After that, the observation is:

{observation}

reward: {reward}

done: {done}

Decide your next action(s).

Please think step by step and provide the actions you want to take.

Your response should be in the format of <think>...</think><answer>...</answer>

grounding_worldmodeling_symbol_format_prompt= """You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

<current_state>...</current_state><think>...</think><answer>...</answer><next_state>...</next_state>

e.g. <current_state>_P__\nG___\n_OO_\n____</current_state><think>I should go down then left to reach the target</think><answer>Down{action_sep}Left</answer><next_state>____\n√___\n_OO_\n____</next_state>

"""



format_prompt={

    "free_think": free_think_format_prompt,

    "no_think": no_think_format_prompt,

    "grounding": grounding_format_prompt,

    "worldmodeling": worldmodeling_format_prompt,

    "grounding_worldmodeling": grounding_worldmodeling_format_prompt,

    "grounding_symbol": grounding_symbol_format_prompt,

    "worldmodeling_symbol": worldmodeling_format_symbol,

    "grounding_worldmodeling_symbol": grounding_worldmodeling_symbol_format_prompt

}
 
 No newline at end of file

vagen/env/sokoban/env.py

+11 −12
Original line number Diff line number Diff line
@@ -4,9 +4,10 @@ from gym_sokoban.envs.sokoban_env import SokobanEnv as GymSokobanEnv 
from .utils import generate_room

from typing import Dict

from vagen.env.utils.env_utils import NoLoggerWarnings, set_seed

from vagen.env.utils.context_utils import parse_llm_raw_response,convert_numpy_to_PIL

from vagen.env.utils.context_utils import convert_numpy_to_PIL

import numpy as np

from .prompt import system_prompt_text, system_prompt_vision, init_observation_template, action_template

from vagen.env.utils.parse_utils import parse_function_map

from .prompt import system_prompt, init_observation_template, action_template,format_prompt

from .env_config import SokobanEnvConfig

class SokobanEnv(BaseEnv):

    GRID_LOOKUP = {
@@ -35,6 +36,9 @@ class SokobanEnv(BaseEnv): 
            max_steps=self.config.get('max_steps', 100),

            num_boxes=self.config.get('num_boxes', 3),

        )

        self.format_prompt = format_prompt[self.config.prompt_format].format(

            max_actions_per_step=self.config.max_actions_per_step,action_sep=self.config.action_sep)

        self.parse_func= parse_function_map[self.config.prompt_format.rstrip("_symbol")]

        

    def reset(self, seed=None):

        with NoLoggerWarnings():
@@ -58,7 +62,7 @@ class SokobanEnv(BaseEnv): 
        return self._render(init_obs=True), {}

    

    def step(self, action_str: str):

        rst=parse_llm_raw_response(

        rst=self.parse_func(

            response=action_str,

            special_token_list=self.config.get('special_token_list', None),

            action_sep=self.config.get('action_sep', ','),
@@ -96,7 +100,7 @@ class SokobanEnv(BaseEnv): 
            else:

                metrics['turn_metrics']['action_is_valid'] = False

                break

        if metrics['turn_metrics']['action_is_valid']:

        if metrics['turn_metrics']['action_is_valid'] and rst["format_correct"]:

            self.reward += self.config.format_reward

        info["metrics"] = metrics

        metrics['turn_metrics']['action_is_effective'] = not np.array_equal(prev_player_position, self.env.player_position)
@@ -105,10 +109,7 @@ class SokobanEnv(BaseEnv): 
        return self._render(init_obs=False), self.reward, done, info

    

    def system_prompt(self):

        if self.config.render_mode == 'vision':

            return system_prompt_vision.format(max_actions_per_step=self.config.max_actions_per_step, action_sep=self.config.action_sep)

        else:

            return system_prompt_text.format(max_actions_per_step=self.config.max_actions_per_step,action_sep=self.config.action_sep)

        return system_prompt+"\n"+self.format_prompt

    

    

    def compute_reward(self):
@@ -133,15 +134,13 @@ class SokobanEnv(BaseEnv): 
            img_str = "\n".join("".join(lookup(cell) for cell in row) for row in room_state)

        

        if init_obs:

            obs_str = init_observation_template.format(observation=img_str)

            obs_str = init_observation_template.format(observation=img_str)+"\n"+self.format_prompt

            

        else:

            obs_str = action_template.format(

                valid_action=self.valid_actions,

                observation=img_str,

                reward=self.reward,

                done=self._success(),

            )

            )+"\n"+self.format_prompt

        

        if multi_modal_data is not None:

            return {

vagen/env/sokoban/env_config.py

+4 −0
Original line number Diff line number Diff line
@@ -9,6 +9,10 @@ class SokobanEnvConfig(BaseEnvConfig): 
    render_mode: str = "vision" # "vision" or "text"

    min_actions_to_succeed: int = 5

    max_actions_per_step: int = 3

    prompt_format: str = "free_think" 

    # "free_think", "no_think", "grounding", "worldmodeling", "grounding_worldmodeling"

    # "grounding_symbol", "worldmodeling_symbol", "grounding_worldmodeling_symbol"

    use_accuracy_reward: bool = False

    

    def config_id(self) -> str:

        id_fields = ["dim_room", "max_steps", "num_boxes", "render_mode", "min_actions_to_succeed", "max_actions_per_step"]

CRA Git | Maintained and supported by SUSTech CRA and CCSE