Skip to content

Add filter dataset saving for env #143

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuecideng merged 1 commit into from
Feb 19, 2026
Merged

Add filter dataset saving for env #143

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions embodichain/lab/gym/envs/embodied_env.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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)

Comment on lines +211 to 213
Copy link

Copilot AI Feb 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new filter_dataset_saving behavior changes whether DatasetManager is constructed and whether episodes are saved. There are existing pytest-based env tests (e.g., tests/gym/envs/test_embodied_env.py), but none assert that setting this flag prevents dataset manager creation/saving or that close() works when cfg.dataset is configured but saving is filtered. Add a focused unit/integration test that configures a dataset, sets filter_dataset_saving=True, and asserts dataset_manager stays None (and no finalize/save calls occur).

Copilot uses AI. Check for mistakes.
def _apply_functor_filter(self) -> None:
Expand Down Expand Up @@ -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
Expand Down
38 changes: 23 additions & 15 deletions embodichain/lab/gym/utils/gym_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand All @@ -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(
Expand All @@ -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:
Expand Down