Commit 739c11cb authored by jameskrw's avatar jameskrw
Browse files

updated config and env for process reward

parent 02560f92

vagen/env/frozenlake/env.py

+2 −2
Original line number Diff line number Diff line
@@ -10,7 +10,7 @@ from vagen.env.utils.parse_utils import PARSE_FUNC_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

from vagen.env.utils.state_reward_utils import state_reward_wrapper

from vagen.env.utils.state_reward_utils import env_state_reward_wrapper

class FrozenLakeEnv(BaseEnv):

    """

    FrozenLake Environment for training and evaluating language models as agents.
@@ -103,7 +103,7 @@ class FrozenLakeEnv(BaseEnv): 
        self.total_reward = 0

        return self._render(init_obs=True), {}



    @state_reward_wrapper

    @env_state_reward_wrapper

    def step(self, action_str: str):

        """

        Take a step in the environment based on the agent's action.

vagen/env/frozenlake/env_config.py

+4 −0
Original line number Diff line number Diff line
@@ -15,7 +15,11 @@ class FrozenLakeEnvConfig(BaseEnvConfig): 
    # "free_think", "no_think", "grounding", "worldmodeling", "grounding_worldmodeling"

    # "grounding_symbolic", "worldmodeling_symbolic", "grounding_worldmodeling_symbolic"

    # "grounding_structured", "worldmodeling_structured", "grounding_worldmodeling_structured"

    

    # configs for process reward for grounding and world modeling

    use_state_reward: bool = False

    grounding_reward_weight: float = 0.5

    worldmodeling_reward_weight: float = 0.5

    

    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/navigation/env.py

+2 −2
Original line number Diff line number Diff line
@@ -8,7 +8,7 @@ from vagen.env.utils.context_utils import convert_numpy_to_PIL 
from vagen.env.utils.parse_utils import PARSE_FUNC_MAP

from .env_config import NavigationEnvConfig

from .prompt import system_prompt,init_observation_template, action_template, format_prompt

from vagen.env.utils.state_reward_utils import state_reward_wrapper

from vagen.env.utils.state_reward_utils import env_state_reward_wrapper



class NavigationEnv(BaseEnv):

    """Navigation environment from embodied bench. """
@@ -184,7 +184,7 @@ class NavigationEnv(BaseEnv): 
        

        return self._render(init_obs=True), {}

    

    @state_reward_wrapper

    @env_state_reward_wrapper

    def step(self, action_str: str):

        """Execute an action in the environment.

vagen/env/navigation/env_config.py

+6 −1
Original line number Diff line number Diff line
@@ -16,8 +16,13 @@ class NavigationEnvConfig(BaseEnvConfig): 
    gpu_device: int = 0

    prompt_format: str = "free_think" 

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

    use_state_reward: bool = False

    

    

    # configs for process reward for grounding and world modeling

    max_objects_in_state: int = 10

    use_state_reward: bool = False

    grounding_reward_weight: float = 0.5

    worldmodeling_reward_weight: float = 0.5



    def config_id(self) -> str:

        """Generate a unique identifier for this configuration."""

vagen/env/primitive_skill/env.py

+2 −2
Original line number Diff line number Diff line
@@ -10,7 +10,7 @@ from .maniskill.utils import build_env, handle_info, get_workspace_limits 
from .prompt import system_prompt, init_observation_template, action_template, format_prompt

import vagen.env.primitive_skill.maniskill.env

import random

from vagen.env.utils.state_reward_utils import state_reward_wrapper

from vagen.env.utils.state_reward_utils import env_state_reward_wrapper

class PrimitiveSkillEnv(BaseEnv):

    def __init__(self, config: PrimitiveSkillEnvConfig):

        """
@@ -55,7 +55,7 @@ class PrimitiveSkillEnv(BaseEnv): 
        self.steps = 0

        return obs, {}

    

    @state_reward_wrapper

    @env_state_reward_wrapper

    def step(self, action_str):

        """

        Take a step in the environment based on the agent's action.

CRA Git | Maintained and supported by SUSTech CRA and CCSE