updated proces reward pipeline (58612539) · Commits · 张泽凯 / VAGEN

vagen/env/base/base_service.py

+4 −1

Original line number	Diff line number	Diff line
		@@ -20,7 +20,10 @@ class BaseService(ABC):
		ids2configs (Dict[Any, Any]):
		A dictionary where each key is an environment ID and the corresponding
		value is the configuration for that environment.

		id: a string
		config: {"env_name": env_name, "env_config": env_config}
		env_name: The name of the environment to create.
		env_config: A dictionary containing the configuration parameters for the environment.
		Returns:
		None

+1 −2

Original line number	Diff line number	Diff line
		@@ -141,8 +141,7 @@ Your answer should be in the format of <think>...</think><answer>...</answer>.
		e.g. <think>I can see there're multiple images with different view. I can see from the second view the object is on the target's left.I think the correct answer is A</think><answer>A</answer>"""

		def compute_reward(self) -> float:
		"""Return the total reward accumulated so far"""
		return self.total_reward
		return 0.0


		if __name__ == "__main__":

+4 −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 env_state_reward_wrapper
		from vagen.env.utils.state_reward_text_utils import env_state_reward_wrapper
		class FrozenLakeEnv(BaseEnv):
		"""
		FrozenLake Environment for training and evaluating language models as agents.
		@@ -217,7 +217,9 @@ class FrozenLakeEnv(BaseEnv):
		Returns:
		float: Total reward accumulated during the current episode
		"""
		return self.total_reward
		# Now we accumulate reward in each step in rollout_manager
		# Set it to non-zero only if you want to give a special trajectory reward
		return 0.0

		def close(self):
		self.gym_env.close()

+1 −0

Original line number	Diff line number	Diff line
		@@ -4,6 +4,7 @@ from typing import Optional, List, Union

		@dataclass
		class FrozenLakeEnvConfig(BaseEnvConfig):
		env_name: str = "frozenlake"
		desc: Optional[List[str]] = None # environment map
		is_slippery: bool = False
		size: int = 4

+2 −0

Original line number	Diff line number	Diff line
		@@ -6,6 +6,7 @@ from vagen.server.serial import serialize_observation
		from .env import FrozenLakeEnv
		from .env_config import FrozenLakeEnvConfig
		from ..base.base_service_config import BaseServiceConfig
		from vagen.env.utils.state_reward_text_utils import service_state_reward_wrapper

		class FrozenLakeService(BaseService):
		"""
		@@ -122,6 +123,7 @@ class FrozenLakeService(BaseService):

		return results

		@service_state_reward_wrapper
		def step_batch(self, ids2actions: Dict[Any, Any]) -> Dict[Any, Tuple[Dict, float, bool, Dict]]:
		"""
		Take a step in multiple FrozenLake environments in parallel.