From 79a2c4803d386565e30c79a36a7cd3b6d8f53b9c Mon Sep 17 00:00:00 2001 From: yuanhaonan Date: Wed, 4 Feb 2026 11:46:21 +0800 Subject: [PATCH 1/8] save success status before reset --- embodichain/lab/gym/envs/base_env.py | 17 +++++++++++++++++ embodichain/lab/gym/envs/embodied_env.py | 3 ++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/embodichain/lab/gym/envs/base_env.py b/embodichain/lab/gym/envs/base_env.py index fd0f873c..eb0f1c00 100644 --- a/embodichain/lab/gym/envs/base_env.py +++ b/embodichain/lab/gym/envs/base_env.py @@ -482,6 +482,19 @@ def get_reward( return rewards + def is_task_success(self, **kwargs) -> torch.Tensor: + """Check if the task is successful for each environment. + + This method should be overridden in subclasses to implement task-specific success criteria. + + Args: + **kwargs: Additional keyword arguments. + + Returns: + A boolean tensor indicating success for each environment. + """ + return torch.ones(self.num_envs, dtype=torch.bool, device=self.device) + def _preprocess_action(self, action: EnvAction) -> EnvAction: """Preprocess action before sending to robot. @@ -531,6 +544,10 @@ def reset( "reset_ids", torch.arange(self.num_envs, dtype=torch.int32, device=self.device), ) + + # Save task success status before resetting objects + self._task_success = self.is_task_success() + self.sim.reset_objects_state(env_ids=reset_ids) self._elapsed_steps[reset_ids] = 0 diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index 98355f62..3885444a 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -379,7 +379,8 @@ def _initialize_episode( if save_data and self.cfg.dataset: if "save" in self.dataset_manager.available_modes: - current_task_success = self.is_task_success() + # Use task success status saved before reset_objects_state + current_task_success = self._task_success # Filter to only save successful episodes successful_env_ids = [ From 923032bad08c3bbe2c870706bbc32431cff6fa2b Mon Sep 17 00:00:00 2001 From: yuanhaonan Date: Wed, 4 Feb 2026 14:49:43 +0800 Subject: [PATCH 2/8] wip --- embodichain/lab/gym/envs/base_env.py | 12 ++++++++---- embodichain/lab/gym/envs/embodied_env.py | 7 ++++++- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/embodichain/lab/gym/envs/base_env.py b/embodichain/lab/gym/envs/base_env.py index eb0f1c00..5075eeb1 100644 --- a/embodichain/lab/gym/envs/base_env.py +++ b/embodichain/lab/gym/envs/base_env.py @@ -133,6 +133,10 @@ def __init__( self._num_envs, dtype=torch.int32, device=self.sim_cfg.sim_device ) + self._task_success = torch.zeros( + self._num_envs, dtype=torch.bool, device=self.device + ) + self._init_sim_state(**kwargs) self._init_raw_obs: Dict = self.get_obs(**kwargs) @@ -484,12 +488,12 @@ def get_reward( def is_task_success(self, **kwargs) -> torch.Tensor: """Check if the task is successful for each environment. - + This method should be overridden in subclasses to implement task-specific success criteria. - + Args: **kwargs: Additional keyword arguments. - + Returns: A boolean tensor indicating success for each environment. """ @@ -544,7 +548,7 @@ def reset( "reset_ids", torch.arange(self.num_envs, dtype=torch.int32, device=self.device), ) - + # Save task success status before resetting objects self._task_success = self.is_task_success() diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index 3885444a..0c35c251 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -380,7 +380,12 @@ def _initialize_episode( if "save" in self.dataset_manager.available_modes: # Use task success status saved before reset_objects_state - current_task_success = self._task_success + current_task_success = getattr(self, "_task_success", None) + if current_task_success is None: + logger.log_warning("task_success is not defined, nothing to save.") + current_task_success = torch.zeros( + self.num_envs, dtype=torch.bool, device=self.device + ) # Filter to only save successful episodes successful_env_ids = [ From 1d34500fe915442271cec74beafd62d0540ec02f Mon Sep 17 00:00:00 2001 From: yuanhaonan Date: Thu, 5 Feb 2026 11:42:27 +0800 Subject: [PATCH 3/8] wip --- embodichain/lab/gym/envs/embodied_env.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index 0c35c251..249945f2 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -380,10 +380,9 @@ def _initialize_episode( if "save" in self.dataset_manager.available_modes: # Use task success status saved before reset_objects_state - current_task_success = getattr(self, "_task_success", None) - if current_task_success is None: + if self._task_success is None: logger.log_warning("task_success is not defined, nothing to save.") - current_task_success = torch.zeros( + self._task_success = torch.zeros( self.num_envs, dtype=torch.bool, device=self.device ) @@ -393,7 +392,7 @@ def _initialize_episode( for env_id in env_ids_to_process if ( self.episode_success_status.get(env_id, False) - or current_task_success[env_id].item() + or self._task_success[env_id].item() ) ] From f2b591ae1340d9edb7fb3f6a54b36b7def2a7693 Mon Sep 17 00:00:00 2001 From: yuanhaonan Date: Thu, 5 Feb 2026 15:39:17 +0800 Subject: [PATCH 4/8] install qpsolvers[osqp] --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 0b4624d7..e98107e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ dependencies = [ "pin-pink", "casadi", "qpsolvers==4.8.1", + "qpsolvers[osqp]", "py_opw_kinematics==0.1.6", "pytorch_kinematics==0.7.6", "polars==1.31.0", From da27060dc5cf762ee3e03ac9c926cb02b3442c01 Mon Sep 17 00:00:00 2001 From: yuanhaonan Date: Thu, 5 Feb 2026 16:46:52 +0800 Subject: [PATCH 5/8] config_to_cfg modules parsing --- embodichain/agents/rl/train.py | 13 ++++++- embodichain/lab/gym/utils/gym_utils.py | 52 +++++++++++++------------- embodichain/lab/scripts/preview_env.py | 13 ++++++- embodichain/lab/scripts/run_agent.py | 12 +++++- embodichain/lab/scripts/run_env.py | 13 ++++++- tests/gym/envs/test_embodied_env.py | 12 +++++- 6 files changed, 83 insertions(+), 32 deletions(-) diff --git a/embodichain/agents/rl/train.py b/embodichain/agents/rl/train.py index 0f766954..cbc9dfdc 100644 --- a/embodichain/agents/rl/train.py +++ b/embodichain/agents/rl/train.py @@ -133,7 +133,18 @@ def train_from_config(config_path: str): logger.log_info(f"Current working directory: {Path.cwd()}") gym_config_data = load_json(str(gym_config_path)) - gym_env_cfg = config_to_cfg(gym_config_data) + + # Define manager modules for config parsing + manager_modules = [ + "embodichain.lab.gym.envs.managers.datasets", + "embodichain.lab.gym.envs.managers.randomization", + "embodichain.lab.gym.envs.managers.record", + "embodichain.lab.gym.envs.managers.events", + "embodichain.lab.gym.envs.managers.observations", + "embodichain.lab.gym.envs.managers.rewards", + ] + + gym_env_cfg = config_to_cfg(gym_config_data, manager_modules=manager_modules) # Override num_envs from train config if provided if num_envs is not None: diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 1cc1ba1e..ad1250c6 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -323,11 +323,13 @@ def cat_tensor_with_ids( return out -def config_to_cfg(config: dict) -> "EmbodiedEnvCfg": +def config_to_cfg(config: dict, manager_modules: list = None) -> "EmbodiedEnvCfg": """Parser configuration file into cfgs for env initialization. Args: config (dict): The configuration dictionary containing robot, sensor, light, background, and interactive objects. + manager_modules (list): List of module paths for dataset, event, observation, and reward managers. + If not provided, uses default module paths. Returns: EmbodiedEnvCfg: A configuration object for initializing the environment. @@ -437,13 +439,26 @@ class ComponentCfg: env_cfg.sim_steps_per_control = config["env"].get("sim_steps_per_control", 4) env_cfg.extensions = deepcopy(config.get("env", {}).get("extensions", {})) + # Initialize manager_modules with defaults + default_manager_modules = [ + "embodichain.lab.gym.envs.managers.datasets", + "embodichain.lab.gym.envs.managers.randomization", + "embodichain.lab.gym.envs.managers.record", + "embodichain.lab.gym.envs.managers.events", + "embodichain.lab.gym.envs.managers.observations", + "embodichain.lab.gym.envs.managers.rewards", + ] + + # Extend with user-provided modules, skipping duplicates + if manager_modules is not None: + for module in manager_modules: + if module not in default_manager_modules: + default_manager_modules.append(module) + + manager_modules = default_manager_modules + env_cfg.dataset = ComponentCfg() if "dataset" in config["env"]: - # Define modules to search for dataset functions - dataset_modules = [ - "embodichain.lab.gym.envs.managers.datasets", - ] - for dataset_name, dataset_params in config["env"]["dataset"].items(): dataset_params_modified = deepcopy(dataset_params) @@ -457,7 +472,7 @@ class ComponentCfg: # Find the function from multiple modules using the utility function dataset_func = find_function_from_modules( func_name, - dataset_modules, + manager_modules, raise_if_not_found=True, ) @@ -476,13 +491,6 @@ class ComponentCfg: env_cfg.events = ComponentCfg() if "events" in config["env"]: - # Define modules to search for event functions - event_modules = [ - "embodichain.lab.gym.envs.managers.randomization", - "embodichain.lab.gym.envs.managers.record", - "embodichain.lab.gym.envs.managers.events", - ] - # parser env events config for event_name, event_params in config["env"]["events"].items(): event_params_modified = deepcopy(event_params) @@ -500,7 +508,7 @@ class ComponentCfg: # Find the function from multiple modules using the utility function event_func = find_function_from_modules( - event_params["func"], event_modules, raise_if_not_found=True + event_params["func"], manager_modules, raise_if_not_found=True ) interval_step = event_params_modified.get("interval_step", 10) @@ -514,11 +522,6 @@ class ComponentCfg: env_cfg.observations = ComponentCfg() if "observations" in config["env"]: - # Define modules to search for observation functions - observation_modules = [ - "embodichain.lab.gym.envs.managers.observations", - ] - for obs_name, obs_params in config["env"]["observations"].items(): obs_params_modified = deepcopy(obs_params) @@ -531,7 +534,7 @@ class ComponentCfg: # Find the function from multiple modules using the utility function obs_func = find_function_from_modules( obs_params["func"], - observation_modules, + manager_modules, raise_if_not_found=True, ) @@ -546,11 +549,6 @@ class ComponentCfg: env_cfg.rewards = ComponentCfg() if "rewards" in config["env"]: - # Define modules to search for reward functions - reward_modules = [ - "embodichain.lab.gym.envs.managers.rewards", - ] - for reward_name, reward_params in config["env"]["rewards"].items(): reward_params_modified = deepcopy(reward_params) @@ -573,7 +571,7 @@ class ComponentCfg: # Find the function from multiple modules using the utility function reward_func = find_function_from_modules( reward_params["func"], - reward_modules, + manager_modules, raise_if_not_found=True, ) diff --git a/embodichain/lab/scripts/preview_env.py b/embodichain/lab/scripts/preview_env.py index fde071bd..1d813da7 100644 --- a/embodichain/lab/scripts/preview_env.py +++ b/embodichain/lab/scripts/preview_env.py @@ -86,7 +86,18 @@ ############################################################################################## # load gym config gym_config = load_json(args.gym_config) - cfg: EmbodiedEnvCfg = config_to_cfg(gym_config) + + # Define manager modules for config parsing + manager_modules = [ + "embodichain.lab.gym.envs.managers.datasets", + "embodichain.lab.gym.envs.managers.randomization", + "embodichain.lab.gym.envs.managers.record", + "embodichain.lab.gym.envs.managers.events", + "embodichain.lab.gym.envs.managers.observations", + "embodichain.lab.gym.envs.managers.rewards", + ] + + cfg: EmbodiedEnvCfg = config_to_cfg(gym_config, manager_modules=manager_modules) cfg.filter_visual_rand = args.filter_visual_rand action_config = {} diff --git a/embodichain/lab/scripts/run_agent.py b/embodichain/lab/scripts/run_agent.py index e6546a59..213d9dc8 100644 --- a/embodichain/lab/scripts/run_agent.py +++ b/embodichain/lab/scripts/run_agent.py @@ -119,8 +119,18 @@ gym_config = load_json(args.gym_config) agent_config = load_json(args.agent_config) + # Define manager modules for config parsing + manager_modules = [ + "embodichain.lab.gym.envs.managers.datasets", + "embodichain.lab.gym.envs.managers.randomization", + "embodichain.lab.gym.envs.managers.record", + "embodichain.lab.gym.envs.managers.events", + "embodichain.lab.gym.envs.managers.observations", + "embodichain.lab.gym.envs.managers.rewards", + ] + # Build environment configuration - cfg: EmbodiedEnvCfg = config_to_cfg(gym_config) + cfg: EmbodiedEnvCfg = config_to_cfg(gym_config, manager_modules=manager_modules) cfg.filter_visual_rand = args.filter_visual_rand cfg.num_envs = args.num_envs cfg.sim_cfg = SimulationManagerCfg( diff --git a/embodichain/lab/scripts/run_env.py b/embodichain/lab/scripts/run_env.py index 74398233..ea55f127 100644 --- a/embodichain/lab/scripts/run_env.py +++ b/embodichain/lab/scripts/run_env.py @@ -191,7 +191,18 @@ def main(args, env, gym_config): # log_error(f"Currently only support num_envs=1, but got {args.num_envs}.") gym_config = load_json(args.gym_config) - cfg: EmbodiedEnvCfg = config_to_cfg(gym_config) + + # Define manager modules for config parsing + manager_modules = [ + "embodichain.lab.gym.envs.managers.datasets", + "embodichain.lab.gym.envs.managers.randomization", + "embodichain.lab.gym.envs.managers.record", + "embodichain.lab.gym.envs.managers.events", + "embodichain.lab.gym.envs.managers.observations", + "embodichain.lab.gym.envs.managers.rewards", + ] + + cfg: EmbodiedEnvCfg = config_to_cfg(gym_config, manager_modules=manager_modules) cfg.filter_visual_rand = args.filter_visual_rand action_config = {} diff --git a/tests/gym/envs/test_embodied_env.py b/tests/gym/envs/test_embodied_env.py index 574fd60c..5e567003 100644 --- a/tests/gym/envs/test_embodied_env.py +++ b/tests/gym/envs/test_embodied_env.py @@ -120,7 +120,17 @@ class EmbodiedEnvTest: """Shared test logic for CPU and CUDA.""" def setup_simulation(self, sim_device): - cfg: EmbodiedEnvCfg = config_to_cfg(METADATA) + # Define manager modules for config parsing + manager_modules = [ + "embodichain.lab.gym.envs.managers.datasets", + "embodichain.lab.gym.envs.managers.randomization", + "embodichain.lab.gym.envs.managers.record", + "embodichain.lab.gym.envs.managers.events", + "embodichain.lab.gym.envs.managers.observations", + "embodichain.lab.gym.envs.managers.rewards", + ] + + cfg: EmbodiedEnvCfg = config_to_cfg(METADATA, manager_modules=manager_modules) cfg.num_envs = NUM_ENVS cfg.sim_cfg = SimulationManagerCfg(headless=True, sim_device=sim_device) From 15578f1cf26c840f56994fa88bbb8bfe87f7d0c7 Mon Sep 17 00:00:00 2001 From: yuanhaonan Date: Thu, 5 Feb 2026 17:25:54 +0800 Subject: [PATCH 6/8] wip --- embodichain/agents/rl/train.py | 21 ++++----------------- embodichain/lab/gym/envs/base_env.py | 6 +++++- embodichain/lab/gym/utils/gym_utils.py | 19 +++++++++++-------- embodichain/lab/scripts/preview_env.py | 16 ++++------------ embodichain/lab/scripts/run_agent.py | 15 ++++----------- embodichain/lab/scripts/run_env.py | 16 ++++------------ pyproject.toml | 3 +-- tests/gym/envs/test_embodied_env.py | 16 ++++------------ 8 files changed, 37 insertions(+), 75 deletions(-) diff --git a/embodichain/agents/rl/train.py b/embodichain/agents/rl/train.py index cbc9dfdc..e1f8ff29 100644 --- a/embodichain/agents/rl/train.py +++ b/embodichain/agents/rl/train.py @@ -32,7 +32,7 @@ from embodichain.agents.rl.utils.trainer import Trainer from embodichain.utils import logger from embodichain.lab.gym.envs.tasks.rl import build_env -from embodichain.lab.gym.utils.gym_utils import config_to_cfg +from embodichain.lab.gym.utils.gym_utils import config_to_cfg, DEFAULT_MANAGER_MODULES from embodichain.utils.utility import load_json from embodichain.utils.module_utils import find_function_from_modules from embodichain.lab.sim import SimulationManagerCfg @@ -133,22 +133,9 @@ def train_from_config(config_path: str): logger.log_info(f"Current working directory: {Path.cwd()}") gym_config_data = load_json(str(gym_config_path)) - - # Define manager modules for config parsing - manager_modules = [ - "embodichain.lab.gym.envs.managers.datasets", - "embodichain.lab.gym.envs.managers.randomization", - "embodichain.lab.gym.envs.managers.record", - "embodichain.lab.gym.envs.managers.events", - "embodichain.lab.gym.envs.managers.observations", - "embodichain.lab.gym.envs.managers.rewards", - ] - - gym_env_cfg = config_to_cfg(gym_config_data, manager_modules=manager_modules) - - # Override num_envs from train config if provided - if num_envs is not None: - gym_env_cfg.num_envs = num_envs + gym_env_cfg = config_to_cfg( + gym_config_data, manager_modules=DEFAULT_MANAGER_MODULES + ) # Ensure sim configuration mirrors runtime overrides if gym_env_cfg.sim_cfg is None: diff --git a/embodichain/lab/gym/envs/base_env.py b/embodichain/lab/gym/envs/base_env.py index 5075eeb1..59e7c7a2 100644 --- a/embodichain/lab/gym/envs/base_env.py +++ b/embodichain/lab/gym/envs/base_env.py @@ -497,7 +497,11 @@ def is_task_success(self, **kwargs) -> torch.Tensor: Returns: A boolean tensor indicating success for each environment. """ - return torch.ones(self.num_envs, dtype=torch.bool, device=self.device) + + raise NotImplementedError( + "BaseEnv.is_task_success must be implemented in subclasses to define " + "task-specific success criteria." + ) def _preprocess_action(self, action: EnvAction) -> EnvAction: """Preprocess action before sending to robot. diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index ad1250c6..3555322c 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -28,6 +28,16 @@ from embodichain.utils.utility import get_class_instance from dexsim.utility import log_debug, log_error +# Default manager modules for config parsing +DEFAULT_MANAGER_MODULES = [ + "embodichain.lab.gym.envs.managers.datasets", + "embodichain.lab.gym.envs.managers.randomization", + "embodichain.lab.gym.envs.managers.record", + "embodichain.lab.gym.envs.managers.events", + "embodichain.lab.gym.envs.managers.observations", + "embodichain.lab.gym.envs.managers.rewards", +] + def get_dtype_bounds(dtype: np.dtype): """Gets the min and max values of a given numpy type""" @@ -440,14 +450,7 @@ class ComponentCfg: env_cfg.extensions = deepcopy(config.get("env", {}).get("extensions", {})) # Initialize manager_modules with defaults - default_manager_modules = [ - "embodichain.lab.gym.envs.managers.datasets", - "embodichain.lab.gym.envs.managers.randomization", - "embodichain.lab.gym.envs.managers.record", - "embodichain.lab.gym.envs.managers.events", - "embodichain.lab.gym.envs.managers.observations", - "embodichain.lab.gym.envs.managers.rewards", - ] + default_manager_modules = DEFAULT_MANAGER_MODULES.copy() # Extend with user-provided modules, skipping duplicates if manager_modules is not None: diff --git a/embodichain/lab/scripts/preview_env.py b/embodichain/lab/scripts/preview_env.py index 1d813da7..09f7cda3 100644 --- a/embodichain/lab/scripts/preview_env.py +++ b/embodichain/lab/scripts/preview_env.py @@ -23,6 +23,7 @@ from embodichain.lab.gym.envs import EmbodiedEnvCfg from embodichain.lab.gym.utils.gym_utils import ( config_to_cfg, + DEFAULT_MANAGER_MODULES, ) from embodichain.utils.utility import load_json from embodichain.utils import logger @@ -86,18 +87,9 @@ ############################################################################################## # load gym config gym_config = load_json(args.gym_config) - - # Define manager modules for config parsing - manager_modules = [ - "embodichain.lab.gym.envs.managers.datasets", - "embodichain.lab.gym.envs.managers.randomization", - "embodichain.lab.gym.envs.managers.record", - "embodichain.lab.gym.envs.managers.events", - "embodichain.lab.gym.envs.managers.observations", - "embodichain.lab.gym.envs.managers.rewards", - ] - - cfg: EmbodiedEnvCfg = config_to_cfg(gym_config, manager_modules=manager_modules) + cfg: EmbodiedEnvCfg = config_to_cfg( + gym_config, manager_modules=DEFAULT_MANAGER_MODULES + ) cfg.filter_visual_rand = args.filter_visual_rand action_config = {} diff --git a/embodichain/lab/scripts/run_agent.py b/embodichain/lab/scripts/run_agent.py index 213d9dc8..f8e248ec 100644 --- a/embodichain/lab/scripts/run_agent.py +++ b/embodichain/lab/scripts/run_agent.py @@ -24,6 +24,7 @@ from embodichain.lab.gym.envs import EmbodiedEnvCfg from embodichain.lab.gym.utils.gym_utils import ( config_to_cfg, + DEFAULT_MANAGER_MODULES, ) from embodichain.utils.logger import log_warning, log_info, log_error from .run_env import main @@ -119,18 +120,10 @@ gym_config = load_json(args.gym_config) agent_config = load_json(args.agent_config) - # Define manager modules for config parsing - manager_modules = [ - "embodichain.lab.gym.envs.managers.datasets", - "embodichain.lab.gym.envs.managers.randomization", - "embodichain.lab.gym.envs.managers.record", - "embodichain.lab.gym.envs.managers.events", - "embodichain.lab.gym.envs.managers.observations", - "embodichain.lab.gym.envs.managers.rewards", - ] - # Build environment configuration - cfg: EmbodiedEnvCfg = config_to_cfg(gym_config, manager_modules=manager_modules) + cfg: EmbodiedEnvCfg = config_to_cfg( + gym_config, manager_modules=DEFAULT_MANAGER_MODULES + ) cfg.filter_visual_rand = args.filter_visual_rand cfg.num_envs = args.num_envs cfg.sim_cfg = SimulationManagerCfg( diff --git a/embodichain/lab/scripts/run_env.py b/embodichain/lab/scripts/run_env.py index ea55f127..6b3844b1 100644 --- a/embodichain/lab/scripts/run_env.py +++ b/embodichain/lab/scripts/run_env.py @@ -28,6 +28,7 @@ from embodichain.lab.gym.envs import EmbodiedEnvCfg from embodichain.lab.gym.utils.gym_utils import ( config_to_cfg, + DEFAULT_MANAGER_MODULES, ) from embodichain.utils.logger import log_warning, log_info, log_error @@ -191,18 +192,9 @@ def main(args, env, gym_config): # log_error(f"Currently only support num_envs=1, but got {args.num_envs}.") gym_config = load_json(args.gym_config) - - # Define manager modules for config parsing - manager_modules = [ - "embodichain.lab.gym.envs.managers.datasets", - "embodichain.lab.gym.envs.managers.randomization", - "embodichain.lab.gym.envs.managers.record", - "embodichain.lab.gym.envs.managers.events", - "embodichain.lab.gym.envs.managers.observations", - "embodichain.lab.gym.envs.managers.rewards", - ] - - cfg: EmbodiedEnvCfg = config_to_cfg(gym_config, manager_modules=manager_modules) + cfg: EmbodiedEnvCfg = config_to_cfg( + gym_config, manager_modules=DEFAULT_MANAGER_MODULES + ) cfg.filter_visual_rand = args.filter_visual_rand action_config = {} diff --git a/pyproject.toml b/pyproject.toml index e98107e8..29a2b35d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,8 +35,7 @@ dependencies = [ "pin", "pin-pink", "casadi", - "qpsolvers==4.8.1", - "qpsolvers[osqp]", + "qpsolvers[osqp]==4.8.1", "py_opw_kinematics==0.1.6", "pytorch_kinematics==0.7.6", "polars==1.31.0", diff --git a/tests/gym/envs/test_embodied_env.py b/tests/gym/envs/test_embodied_env.py index 5e567003..baae2700 100644 --- a/tests/gym/envs/test_embodied_env.py +++ b/tests/gym/envs/test_embodied_env.py @@ -22,7 +22,7 @@ from embodichain.lab.gym.envs import EmbodiedEnvCfg from embodichain.lab.sim.objects import RigidObject, Robot -from embodichain.lab.gym.utils.gym_utils import config_to_cfg +from embodichain.lab.gym.utils.gym_utils import config_to_cfg, DEFAULT_MANAGER_MODULES from embodichain.lab.gym.utils.registration import register_env from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.data import get_data_path @@ -120,17 +120,9 @@ class EmbodiedEnvTest: """Shared test logic for CPU and CUDA.""" def setup_simulation(self, sim_device): - # Define manager modules for config parsing - manager_modules = [ - "embodichain.lab.gym.envs.managers.datasets", - "embodichain.lab.gym.envs.managers.randomization", - "embodichain.lab.gym.envs.managers.record", - "embodichain.lab.gym.envs.managers.events", - "embodichain.lab.gym.envs.managers.observations", - "embodichain.lab.gym.envs.managers.rewards", - ] - - cfg: EmbodiedEnvCfg = config_to_cfg(METADATA, manager_modules=manager_modules) + cfg: EmbodiedEnvCfg = config_to_cfg( + METADATA, manager_modules=DEFAULT_MANAGER_MODULES + ) cfg.num_envs = NUM_ENVS cfg.sim_cfg = SimulationManagerCfg(headless=True, sim_device=sim_device) From 5ac4aa6e4c7feaa24d0a50b9ebdd366c0e2c91c8 Mon Sep 17 00:00:00 2001 From: yuanhaonan Date: Thu, 5 Feb 2026 17:27:31 +0800 Subject: [PATCH 7/8] wip --- embodichain/lab/gym/envs/embodied_env.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index 249945f2..9910160f 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -379,13 +379,6 @@ def _initialize_episode( if save_data and self.cfg.dataset: if "save" in self.dataset_manager.available_modes: - # Use task success status saved before reset_objects_state - if self._task_success is None: - logger.log_warning("task_success is not defined, nothing to save.") - self._task_success = torch.zeros( - self.num_envs, dtype=torch.bool, device=self.device - ) - # Filter to only save successful episodes successful_env_ids = [ env_id From ca712a993606339562af1e572a7b05f33edf2222 Mon Sep 17 00:00:00 2001 From: yuanhaonan Date: Fri, 6 Feb 2026 10:37:28 +0800 Subject: [PATCH 8/8] is_task_success --- embodichain/lab/gym/envs/base_env.py | 15 ++++++--------- embodichain/lab/gym/envs/embodied_env.py | 14 -------------- 2 files changed, 6 insertions(+), 23 deletions(-) diff --git a/embodichain/lab/gym/envs/base_env.py b/embodichain/lab/gym/envs/base_env.py index 2c1f38ae..38f2c2d9 100644 --- a/embodichain/lab/gym/envs/base_env.py +++ b/embodichain/lab/gym/envs/base_env.py @@ -489,21 +489,18 @@ def get_reward( return rewards def is_task_success(self, **kwargs) -> torch.Tensor: - """Check if the task is successful for each environment. - - This method should be overridden in subclasses to implement task-specific success criteria. + """ + Determine if the task is successfully completed. This is mainly used in the data generation process + of the imitation learning. Args: - **kwargs: Additional keyword arguments. + **kwargs: Additional arguments for task-specific success criteria. Returns: - A boolean tensor indicating success for each environment. + torch.Tensor: A boolean tensor indicating success for each environment in the batch. """ - raise NotImplementedError( - "BaseEnv.is_task_success must be implemented in subclasses to define " - "task-specific success criteria." - ) + return torch.ones(self.num_envs, dtype=torch.bool, device=self.device) def _preprocess_action(self, action: EnvAction) -> EnvAction: """Preprocess action before sending to robot. diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index 9910160f..ef339123 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -587,20 +587,6 @@ def create_demo_action_list(self, *args, **kwargs) -> Sequence[EnvAction] | None "The method 'create_demo_action_list' must be implemented in subclasses." ) - def is_task_success(self, **kwargs) -> torch.Tensor: - """ - Determine if the task is successfully completed. This is mainly used in the data generation process - of the imitation learning. - - Args: - **kwargs: Additional arguments for task-specific success criteria. - - Returns: - torch.Tensor: A boolean tensor indicating success for each environment in the batch. - """ - - return torch.ones(self.num_envs, dtype=torch.bool, device=self.device) - def close(self) -> None: """Close the environment and release resources.""" # Finalize dataset if present