updated env setting (bcab2cd8) · Commits · 张泽凯 / VAGEN

vagen/env_new/init.py

+7 −3

Original line number	Diff line number	Diff line
		from vagen.env.register import REGISTERED_ENVS, register

		from vagen.env.sokoban.env import SokobanInterface
		No newline at end of file
		from .sokoban import SokobanEnv,SokobanConfig
		REGISTERED_ENV = {
		"sokoban": {
		"env": SokobanEnv,
		"config": SokobanConfig,
		}
		}
		No newline at end of file

vagen/env_new/base_config.py

0 → 100644

+20 −0

Original line number	Diff line number	Diff line
		from dataclasses import dataclass, field
		from abc import ABC, abstractmethod
		@dataclass
		class BaseConfig(ABC):

		@abstractmethod
		def config_id(self) -> str: # config identifier, wandb and mllm rollout manager use this to identify the config
		pass

		def __init__(self, **kwargs):
		pass

		def get(self, key, default=None):
		"""
		Get the value of a config key.
		Args:
		key: Key to get
		default: Default value if key is not found
		"""
		return getattr(self, key, default)
		No newline at end of file

vagen/env_new/base_env.py

+42 −10

Original line number	Diff line number	Diff line
		from abc import ABC, abstractmethod
		from typing import Optional, List, Tuple, Any, Dict
		from typing import Optional, List, Tuple, Dict

		class BaseEnv(ABC):
		def __init__(self, config):
		self.config = config


		@abstractmethod
		def step(self, action) -> Tuple[Any, float, bool, Dict]:
		def step(self, llm_raw_response) -> Tuple[Dict, float, bool, Dict]:
		"""
		Execute one step in the environment.
		NOTE should also handle predefined invalid action (0)
		action is llm raw response
		Args:
		action: Action to take, must be in action space, or default invalid action
		action: Action to take, assume it's raw llm action

		Returns:
		obs, reward, done, info

		obs: {
		'obs_str': "This is the obs template, you see <image> and <image>, you heard <audio> and <audio>",
		'multi_modal_inputs':{
		'<image>':[list of images],
		'<audio>':[list of audios],
		}
		# num of <image> and <audio> in the obs_str should match len(multi_modal_inputs['<image>']) and len(multi_modal_inputs['<audio>'])
		}
		info: {
		"metrics": {
		'success': success,
		'action_is_effective': action_is_effective,
		'action_is_valid': action_is_valid,
		} # metrics you want to log in wandb
		"llm_raw_response": llm_raw_response,
		"llm_response": llm_response, # for update
		}
		"""
		pass

		@@ -25,7 +38,7 @@ class BaseEnv(ABC):
		pass

		@abstractmethod
		def reset(self, seed: Optional[Any] = None) -> Tuple[Any, Dict]:
		def reset(self, seed= None) -> Tuple[Dict, Dict]:
		"""
		Reset the environment.
		NOTE: the environment should be same for the same seed
		@@ -34,5 +47,24 @@ class BaseEnv(ABC):

		Returns:
		obs,info

		format should be same as step
		"""
		pass

		@abstractmethod
		def system_prompt(self) -> str:
		"""
		Get the system prompt for the environment.

		Returns:
		System prompt string.
		"""
		pass

		@abstractmethod
		def compute_reward(self) -> float:
		"""
		give final reward
		"""
		pass
		No newline at end of file

vagen/env_new/base_interface.py

deleted100644 → 0

+0 −49

Original line number	Diff line number	Diff line
		from abc import ABC, abstractmethod
		import re
		from typing import Optional, List, Tuple, Any, Dict
		from copy import deepcopy
		from transformers import AutoTokenizer
		import torch
		from PIL import Image
		import numpy as np
		from dataclasses import dataclass, field
		from .utils.io_utils import validate_reset_io,validate_step_io

		class BaseInterface(ABC):
		image_placeholder="<image>"

		def __init__(self, config: Dict):
		self.config = config

		@classmethod
		@abstractmethod
		def config_repr(cls, config) -> str:
		"""convert config to str"""
		pass

		@abstractmethod
		def close(self):
		"""Close the environment."""
		pass

		@abstractmethod
		def get_task_instruction(self) -> str:
		"""Get the task instruction."""
		pass

		@abstractmethod
		@validate_step_io
		def step(self, action: str) -> Tuple[Dict, float, bool, Dict]:
		pass

		@abstractmethod
		@validate_reset_io
		def reset(self, seed: int) -> Tuple[Dict, Dict]:
		"""Reset the environment."""
		pass

		@abstractmethod
		def get_traj_reward(self) -> float:
		"""Get the reward of the environment."""

vagen/env_new/register.py

deleted100644 → 0

+0 −37

Original line number	Diff line number	Diff line

		REGISTERED_ENVS = {}

		def register(cls=None, *, name=None):
		"""
		A decorator to register environment classes in the REGISTERED_ENVS dictionary.

		Args:
		cls: The class to register
		name: Optional custom name for the environment. If not provided,
		the class name will be used

		Usage:
		@register
		class SokobanEnv(BaseEnv):
		pass

		@register(name="custom_sokoban")
		class CustomSokobanEnv(BaseEnv):
		pass
		"""
		def _register(cls):
		# Use provided name or class name as the registry key
		key = name if name is not None else cls.__name__

		# Register the class
		REGISTERED_ENVS[key] = cls

		# Return the class unchanged so it can be used normally
		return cls

		# Handle case when decorator is used with no arguments: @register
		if cls is not None:
		return _register(cls)

		# Handle case when decorator is used with arguments: @register(name="...")
		return _register
		No newline at end of file

Admin message