# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------
"""Dataset manager for orchestrating dataset collection functors."""
from __future__ import annotations
import inspect
import threading
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
from collections.abc import Sequence
import torch
from prettytable import PrettyTable
from embodichain.utils import logger
from .manager_base import ManagerBase
from .cfg import DatasetFunctorCfg
__all__ = ["DatasetManager"]
if TYPE_CHECKING:
from embodichain.lab.gym.envs import EmbodiedEnv
[docs]
class DatasetManager(ManagerBase):
"""Manager for orchestrating dataset collection and saving using functors.
The dataset manager supports multiple dataset formats through a functor system:
- LeRobot format (via LeRobotRecorder)
- HDF5 format (via HDF5Recorder)
- Zarr format (via ZarrRecorder)
- Custom formats (via user-defined functors)
Each functor's step() method is called once per environment step and handles:
- Recording observation-action pairs
- Detecting episode completion (dones=True)
- Auto-saving completed episodes
Example configuration:
>>> from embodichain.lab.gym.envs.managers.cfg import DatasetFunctorCfg
>>> from embodichain.lab.gym.envs.managers.datasets import LeRobotRecorder
>>>
>>> @configclass
>>> class MyEnvCfg:
>>> dataset: dict = {
>>> "lerobot": DatasetFunctorCfg(
>>> func=LeRobotRecorder,
>>> save_failed_episodes=True,
>>> params={
>>> "robot_meta": {...},
>>> "instruction": {"lang": "pick and place"},
>>> "extra": {"scene_type": "kitchen"},
>>> "save_path": "/data/datasets"
>>> }
>>> )
>>> }
"""
_env: EmbodiedEnv
"""The environment instance."""
[docs]
def __init__(self, cfg: object, env: EmbodiedEnv):
"""Initialize the dataset manager.
Args:
cfg: Configuration object containing dataset functor configurations.
env: The environment instance.
"""
# Store functors by mode (similar to EventManager)
self._mode_functor_names: dict[str, list[str]] = {}
self._mode_functor_cfgs: dict[str, list[DatasetFunctorCfg]] = {}
self._mode_class_functor_cfgs: dict[str, list[DatasetFunctorCfg]] = {}
self._finalize_lock = threading.Lock()
self._finalized = False
self._finalize_result: Optional[str] = None
self._finalize_error: Optional[str] = None
# Call base class to parse functors
super().__init__(cfg, env)
## TODO: fix configurable_action.py to avoid getting env.metadata['dataset']
# Extract robot_meta and instruction from functor params or plain config and add to env.metadata for backward compatibility
# This allows legacy code (like action_bank) to access robot_meta via env.metadata["dataset"]["robot_meta"]
robot_meta_found = False
# First, try to extract from functor params
for mode_cfgs in self._mode_functor_cfgs.values():
for functor_cfg in mode_cfgs:
if (
"robot_meta" in functor_cfg.params
or "instruction" in functor_cfg.params
):
if not hasattr(env, "metadata"):
env.metadata = {}
if "dataset" not in env.metadata:
env.metadata["dataset"] = {}
if "robot_meta" in functor_cfg.params:
env.metadata["dataset"]["robot_meta"] = functor_cfg.params[
"robot_meta"
]
if "instruction" in functor_cfg.params:
env.metadata["dataset"]["instruction"] = functor_cfg.params[
"instruction"
]
logger.log_info(
"Added robot_meta and instruction to env.metadata for backward compatibility"
)
robot_meta_found = True
break
if robot_meta_found:
break
# If not found in functor params, try to extract from plain config
if not robot_meta_found:
# Check if config is dict or object
if isinstance(self.cfg, dict):
cfg_items = self.cfg.items()
else:
cfg_items = self.cfg.__dict__.items()
for config_name, config_value in cfg_items:
if config_name == "robot_meta" and isinstance(config_value, dict):
if not hasattr(env, "metadata"):
env.metadata = {}
if "dataset" not in env.metadata:
env.metadata["dataset"] = {}
env.metadata["dataset"]["robot_meta"] = config_value
logger.log_info(
"Added robot_meta to env.metadata for backward compatibility (from plain config)"
)
break
logger.log_info(
f"DatasetManager initialized with {sum(len(v) for v in self._mode_functor_names.values())} functors"
)
def __str__(self) -> str:
"""Returns: A string representation for dataset manager."""
msg = f"<DatasetManager> contains {len(self._functor_names)} active functors.\n"
table = PrettyTable()
table.title = "Active Dataset Functors"
table.field_names = ["Index", "Name", "Type"]
table.align["Name"] = "l"
for index, name in enumerate(self._functor_names):
functor_cfg = self._functor_cfgs[index]
functor_type = (
functor_cfg.func.__class__.__name__
if hasattr(functor_cfg.func, "__class__")
else str(functor_cfg.func)
)
table.add_row([index, name, functor_type])
msg += table.get_string()
msg += "\n"
return msg
"""
Properties.
"""
@property
def active_functors(self) -> dict[str, list[str]]:
"""Name of active dataset functors by mode.
The keys are the modes and the values are the names of the dataset functors.
"""
return self._mode_functor_names
@property
def available_modes(self) -> list[str]:
"""List of available modes for the dataset manager."""
return list(self._mode_functor_names.keys())
@property
def save_failed_episodes(self) -> bool:
"""Whether any configured dataset recorder should keep failed episodes."""
return any(
functor_cfg.save_failed_episodes
for functor_cfg in self._mode_functor_cfgs.get("save", [])
)
"""
Operations.
"""
[docs]
def reset(
self, env_ids: Union[Sequence[int], torch.Tensor, None] = None
) -> dict[str, float]:
"""Reset all dataset functors.
Args:
env_ids: The environment ids. Defaults to None.
Returns:
Empty dict (no logging info).
"""
# Call reset on all functor instances across all modes
for mode_cfgs in self._mode_functor_cfgs.values():
for functor_cfg in mode_cfgs:
if hasattr(functor_cfg.func, "reset"):
functor_cfg.func.reset(env_ids=env_ids)
return {}
[docs]
def apply(
self,
mode: str,
env_ids: Union[Sequence[int], torch.Tensor, None] = None,
) -> None:
"""Apply dataset functors for the specified mode.
This method saves completed episodes by reading data from the environment's
episode buffers. It should be called before clearing the buffers during reset.
Args:
mode: The mode to apply (currently only "save" is supported).
env_ids: The indices of the environments to apply the functor to.
Defaults to None, in which case the functor is applied to all environments.
"""
# check if mode is valid
if mode not in self._mode_functor_names:
logger.log_warning(
f"Dataset mode '{mode}' is not defined. Skipping dataset operation."
)
return
# iterate over all the dataset functors for this mode
for functor_cfg in self._mode_functor_cfgs[mode]:
functor_cfg.func(
self._env,
env_ids,
**functor_cfg.params,
)
[docs]
def finalize(self) -> Optional[str]:
"""Finalize every dataset functor exactly once.
Finalization is a storage barrier only; individual recorders must not
implicitly commit live episode buffers here. All functors are attempted
even when one fails, and their failures are reported together.
Returns:
Path to the first finalized dataset, or ``None`` if none was returned.
Raises:
RuntimeError: If one or more functors fail to finalize.
"""
with self._finalize_lock:
if self._finalized:
if self._finalize_error is not None:
raise RuntimeError(self._finalize_error)
return self._finalize_result
dataset_paths: list[str] = []
errors: list[str] = []
# Call every functor even when an earlier cleanup failed.
for mode, mode_cfgs in self._mode_functor_cfgs.items():
names = self._mode_functor_names.get(mode, [])
for index, functor_cfg in enumerate(mode_cfgs):
functor = functor_cfg.func
if not hasattr(functor, "finalize"):
continue
functor_name = (
names[index] if index < len(names) else type(functor).__name__
)
try:
path = functor.finalize()
if path:
dataset_paths.append(path)
except Exception as error: # noqa: BLE001 - aggregate cleanup
errors.append(f"{functor_name}: {error}")
self._finalize_result = dataset_paths[0] if dataset_paths else None
self._finalized = True
if errors:
self._finalize_error = (
f"Failed to finalize {len(errors)} dataset functor(s): "
+ "; ".join(errors)
)
raise RuntimeError(self._finalize_error)
if dataset_paths:
logger.log_info(f"Finalized {len(dataset_paths)} datasets")
return self._finalize_result
[docs]
def close(self) -> Optional[str]:
"""Finalize all dataset functors; repeated calls are safe."""
return self.finalize()
[docs]
def get_cached_data(self) -> list[Dict[str, Any]]:
"""Get cached data from all dataset functors (for online training).
Iterates through all functors and collects cached data from those
that support online training mode (have get_cached_data method).
Returns:
List of cached data dictionaries from all functors.
"""
all_cached_data = []
# Iterate through all modes and functors
for mode_cfgs in self._mode_functor_cfgs.values():
for functor_cfg in mode_cfgs:
if hasattr(functor_cfg.func, "get_cached_data"):
cached_data = functor_cfg.func.get_cached_data()
all_cached_data.extend(cached_data)
return all_cached_data
[docs]
def clear_cache(self) -> int:
"""Clear cached data from all dataset functors (for online training).
Iterates through all functors and clears their cache if they
support online training mode (have clear_cache method).
Returns:
Total number of cached items cleared across all functors.
"""
total_cleared = 0
# Iterate through all modes and functors
for mode_cfgs in self._mode_functor_cfgs.values():
for functor_cfg in mode_cfgs:
if hasattr(functor_cfg.func, "clear_cache"):
cleared = functor_cfg.func.clear_cache()
total_cleared += cleared
return total_cleared
[docs]
def get_functor_cfg(self, functor_name: str) -> DatasetFunctorCfg:
"""Gets the configuration for the specified functor.
Args:
functor_name: The name of the dataset functor.
Returns:
The configuration of the dataset functor.
Raises:
ValueError: If the functor name is not found.
"""
for mode, functors in self._mode_functor_names.items():
if functor_name in functors:
return self._mode_functor_cfgs[mode][functors.index(functor_name)]
logger.log_error(f"Dataset functor '{functor_name}' not found.")
def _prepare_functors(self):
"""Prepare dataset functors from configuration.
This method parses the configuration and initializes all dataset functors,
organizing them by mode (similar to EventManager).
"""
# Check if config is dict already
if isinstance(self.cfg, dict):
cfg_items = self.cfg.items()
else:
cfg_items = self.cfg.__dict__.items()
# Iterate over all the functors
for functor_name, functor_cfg in cfg_items:
# Check for non config
if functor_cfg is None:
continue
# Skip non-functor configurations (e.g., robot_meta which is a plain dict)
# Functor configurations must have a "func" field
if isinstance(functor_cfg, dict) and "func" not in functor_cfg:
# This is a plain configuration (not a functor), skip it
continue
# Convert dict to DatasetFunctorCfg if needed (for JSON configs)
if isinstance(functor_cfg, dict):
functor_cfg = DatasetFunctorCfg(**functor_cfg)
# Check for valid config type
if not isinstance(functor_cfg, DatasetFunctorCfg):
raise TypeError(
f"Configuration for '{functor_name}' is not of type DatasetFunctorCfg."
f" Received: '{type(functor_cfg)}'."
)
# Resolve common parameters
# min_argc=7 to skip: env, env_ids, obs, action, dones, terminateds, info
# These are runtime positional arguments, not config parameters
self._resolve_common_functor_cfg(functor_name, functor_cfg, min_argc=7)
# Check if mode is a new mode
if functor_cfg.mode not in self._mode_functor_names:
# add new mode
self._mode_functor_names[functor_cfg.mode] = []
self._mode_functor_cfgs[functor_cfg.mode] = []
self._mode_class_functor_cfgs[functor_cfg.mode] = []
# Add functor name and parameters
self._mode_functor_names[functor_cfg.mode].append(functor_name)
self._mode_functor_cfgs[functor_cfg.mode].append(functor_cfg)
# Check if the functor is a class
if inspect.isclass(functor_cfg.func):
self._mode_class_functor_cfgs[functor_cfg.mode].append(functor_cfg)