From 5a18d94f6392a5e7566dfec92185e2e8817a9109 Mon Sep 17 00:00:00 2001
From: yuecideng <dengyueci@qq.com>
Date: Thu, 19 Feb 2026 09:22:48 +0000
Subject: [PATCH] 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: