From 5a18d94f6392a5e7566dfec92185e2e8817a9109 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Thu, 19 Feb 2026 09:22:48 +0000 Subject: [PATCH 1/4] wip --- embodichain/lab/gym/envs/embodied_env.py | 11 +++++-- embodichain/lab/gym/utils/gym_utils.py | 38 ++++++++++++++---------- 2 files changed, 32 insertions(+), 17 deletions(-) diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index f89621c7..b412babb 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -125,6 +125,13 @@ class EnvLightCfg: This is useful when we want to disable visual randomization for debug motion and physics issues. """ + filter_dataset_saving: bool = False + """Whether to filter out dataset saving + + This is useful when we want to disable dataset saving for debug motion and physics issues. + If no dataset manager is configured, this flag will have no effect. + """ + @register_env("EmbodiedEnv-v1") class EmbodiedEnv(BaseEnv): @@ -201,7 +208,7 @@ def _init_sim_state(self, **kwargs): if self.cfg.rewards: self.reward_manager = RewardManager(self.cfg.rewards, self) - if self.cfg.dataset: + if self.cfg.dataset and not self.cfg.filter_dataset_saving: self.dataset_manager = DatasetManager(self.cfg.dataset, self) def _apply_functor_filter(self) -> None: @@ -376,7 +383,7 @@ def _initialize_episode( env_ids_to_process = list(env_ids) # Save dataset before clearing buffers for environments that are being reset - if save_data and self.cfg.dataset: + if save_data and self.dataset_manager: if "save" in self.dataset_manager.available_modes: # Filter to only save successful episodes diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 72d5fb40..4482b2a8 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -704,15 +704,16 @@ def add_env_launcher_args_to_parser(parser: argparse.ArgumentParser) -> None: """Add common environment launcher arguments to an existing argparse parser. This function adds the following arguments to the provided parser: - - --num_envs: Number of environments to run in parallel (default: 1) - - --device: Device to run the environment on (default: 'cpu') - - --headless: Whether to perform the simulation in headless mode (default: False) - - --enable_rt: Whether to use RTX rendering backend for the simulation (default: False) - - --gpu_id: The GPU ID to use for the simulation (default: 0) - - --filter_visual_rand: Whether to filter out visual randomization (default: False) - - --gym_config: Path to gym config file (default: '') - - --action_config: Path to action config file (default: None) - - --preview: Whether to preview the environment after launching (default: False) + --num_envs: Number of environments to run in parallel (default: 1) + --device: Device to run the environment on (default: 'cpu') + --headless: Whether to perform the simulation in headless mode (default: False) + --enable_rt: Whether to use RTX rendering backend for the simulation (default: False) + --gpu_id: The GPU ID to use for the simulation (default: 0) + --gym_config: Path to gym config file (default: '') + --action_config: Path to action config file (default: None) + --preview: Whether to preview the environment after launching (default: False) + --filter_visual_rand: Whether to filter out visual randomization (default: False) + --filter_dataset_saving: Whether to filter out dataset saving (default: False) Note: 1. In preview mode, the environment will be launched and keep running in a loop for user interaction. @@ -750,12 +751,6 @@ def add_env_launcher_args_to_parser(parser: argparse.ArgumentParser) -> None: default=0, type=int, ) - parser.add_argument( - "--filter_visual_rand", - help="Whether to filter out visual randomization.", - default=False, - action="store_true", - ) parser.add_argument( "--gym_config", type=str, @@ -772,6 +767,18 @@ def add_env_launcher_args_to_parser(parser: argparse.ArgumentParser) -> None: default=False, action="store_true", ) + parser.add_argument( + "--filter_visual_rand", + help="Whether to filter out visual randomization.", + default=False, + action="store_true", + ) + parser.add_argument( + "--filter_dataset_saving", + help="Whether to filter out dataset saving.", + default=False, + action="store_true", + ) def build_env_cfg_from_args( @@ -795,6 +802,7 @@ def build_env_cfg_from_args( gym_config, manager_modules=DEFAULT_MANAGER_MODULES ) cfg.filter_visual_rand = args.filter_visual_rand + cfg.filter_dataset_saving = args.filter_dataset_saving action_config = {} if args.action_config is not None: From 78c649551f650bb6ab0e5f3c7edb99863cdcd325 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Fri, 20 Feb 2026 08:30:28 +0000 Subject: [PATCH 2/4] wip --- .../gym/envs/managers/randomization/spatial.py | 17 ++++++++++++----- embodichain/lab/gym/utils/gym_utils.py | 3 +++ 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/embodichain/lab/gym/envs/managers/randomization/spatial.py b/embodichain/lab/gym/envs/managers/randomization/spatial.py index dda5658d..408fa079 100644 --- a/embodichain/lab/gym/envs/managers/randomization/spatial.py +++ b/embodichain/lab/gym/envs/managers/randomization/spatial.py @@ -110,6 +110,7 @@ def randomize_rigid_object_pose( rotation_range: tuple[list[float], list[float]] | None = None, relative_position: bool = True, relative_rotation: bool = False, + physics_update_step: int = -1, ) -> None: """Randomize the pose of a rigid object in the environment. @@ -122,6 +123,7 @@ def randomize_rigid_object_pose( The rotation is represented as Euler angles (roll, pitch, yaw) in degree. relative_position (bool): Whether to randomize the position relative to the object's initial position. Default is True. relative_rotation (bool): Whether to randomize the rotation relative to the object's initial rotation. Default is False. + physics_update_step (int): The number of physics update steps to apply after randomization. Default is -1 (no update). """ if entity_cfg.uid not in env.sim.get_rigid_object_uid_list(): @@ -155,6 +157,9 @@ def randomize_rigid_object_pose( rigid_object.set_local_pose(pose, env_ids=env_ids) rigid_object.clear_dynamics() + if physics_update_step > 0: + env.sim.update(step=physics_update_step) + def randomize_robot_eef_pose( env: EmbodiedEnv, @@ -202,7 +207,7 @@ def set_random_eef_pose(joint_ids: List[int], robot: Robot) -> None: ) new_qpos[ret == False] = current_qpos[ret == False] - robot.set_qpos(new_qpos, env_ids=env_ids, joint_ids=joint_ids) + robot.set_qpos(new_qpos, env_ids=env_ids, joint_ids=joint_ids, target=False) robot = env.sim.get_robot(entity_cfg.uid) @@ -215,8 +220,8 @@ def set_random_eef_pose(joint_ids: List[int], robot: Robot) -> None: joint_ids = robot.get_joint_ids(part) set_random_eef_pose(joint_ids, robot) - # simulate 10 steps to let the robot reach the target pose. - env.sim.update(step=10) + # simulate 1 steps to let the robot reach the target pose. + env.sim.update(step=1) def randomize_robot_qpos( @@ -264,8 +269,10 @@ def randomize_robot_qpos( else: current_qpos = qpos - robot.set_qpos(qpos=current_qpos, env_ids=env_ids, joint_ids=joint_ids) - env.sim.update(step=100) + robot.set_qpos( + qpos=current_qpos, env_ids=env_ids, joint_ids=joint_ids, target=False + ) + env.sim.update(step=1) def randomize_target_pose( diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 4482b2a8..8cc93ae4 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -803,6 +803,9 @@ def build_env_cfg_from_args( ) cfg.filter_visual_rand = args.filter_visual_rand cfg.filter_dataset_saving = args.filter_dataset_saving + if args.preview: + # In preview mode, we typically don't want to save data + cfg.filter_dataset_saving = True action_config = {} if args.action_config is not None: From 536c17a4ec5371241f0a68d904aa6409a7f75448 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Fri, 20 Feb 2026 08:38:49 +0000 Subject: [PATCH 3/4] wip --- embodichain/lab/gym/utils/gym_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 4482b2a8..8cc93ae4 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -803,6 +803,9 @@ def build_env_cfg_from_args( ) cfg.filter_visual_rand = args.filter_visual_rand cfg.filter_dataset_saving = args.filter_dataset_saving + if args.preview: + # In preview mode, we typically don't want to save data + cfg.filter_dataset_saving = True action_config = {} if args.action_config is not None: From eff2f65313d60efa68772d23263d6d7ea5bf8b02 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Fri, 20 Feb 2026 16:47:49 +0800 Subject: [PATCH 4/4] wip --- embodichain/lab/gym/envs/managers/randomization/spatial.py | 2 ++ embodichain/lab/gym/utils/gym_utils.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/embodichain/lab/gym/envs/managers/randomization/spatial.py b/embodichain/lab/gym/envs/managers/randomization/spatial.py index 408fa079..503f10fa 100644 --- a/embodichain/lab/gym/envs/managers/randomization/spatial.py +++ b/embodichain/lab/gym/envs/managers/randomization/spatial.py @@ -208,6 +208,7 @@ def set_random_eef_pose(joint_ids: List[int], robot: Robot) -> None: new_qpos[ret == False] = current_qpos[ret == False] robot.set_qpos(new_qpos, env_ids=env_ids, joint_ids=joint_ids, target=False) + robot.set_qpos(new_qpos, env_ids=env_ids, joint_ids=joint_ids) robot = env.sim.get_robot(entity_cfg.uid) @@ -272,6 +273,7 @@ def randomize_robot_qpos( robot.set_qpos( qpos=current_qpos, env_ids=env_ids, joint_ids=joint_ids, target=False ) + robot.set_qpos(qpos=current_qpos, env_ids=env_ids, joint_ids=joint_ids) env.sim.update(step=1) diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 8cc93ae4..637afcb9 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -805,7 +805,7 @@ def build_env_cfg_from_args( cfg.filter_dataset_saving = args.filter_dataset_saving if args.preview: # In preview mode, we typically don't want to save data - cfg.filter_dataset_saving = True + cfg.filter_dataset_saving = True action_config = {} if args.action_config is not None: