From a8b3beab42a2c4ad842b7264ba3d447345a748bd Mon Sep 17 00:00:00 2001 From: yuecideng Date: Wed, 25 Mar 2026 15:51:01 +0000 Subject: [PATCH] Fix KeyError when 'add' mode not present in observation_manager.active_functors The code previously assumed 'add' always existed in active_functors, which caused KeyError in configurations without observation functors in 'add' mode. Added a conditional check before iterating over 'add' functors to prevent this error. Co-Authored-By: Claude Opus 4.6 --- embodichain/lab/gym/envs/managers/datasets.py | 36 ++++++++++--------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/embodichain/lab/gym/envs/managers/datasets.py b/embodichain/lab/gym/envs/managers/datasets.py index 33aa0969..005eb699 100644 --- a/embodichain/lab/gym/envs/managers/datasets.py +++ b/embodichain/lab/gym/envs/managers/datasets.py @@ -416,23 +416,25 @@ def _modify_feature_names(self, features: dict[str, Any]) -> None: if feature["shape"] == (): features[key]["shape"] = (1,) - for functor_name in self._env.observation_manager.active_functors["add"]: - functor_cfg = self._env.observation_manager.get_functor_cfg( - functor_name=functor_name - ) - if functor_cfg.func == get_object_uid: - obs_key = functor_cfg.name - asset_uid = functor_cfg.params["entity_cfg"].uid - asset = self._env.sim.get_asset(asset_uid) - if isinstance(asset, RigidObject): - features[obs_key]["names"] = asset_uid - elif isinstance(asset, (Articulation, Robot)): - link_names = asset.link_names - features[obs_key]["names"] = link_names - else: - logger.log_warning( - f"Asset with UID '{asset_uid}' is not RigidObject, Articulation or Robot. Cannot assign feature names based on asset properties." - ) + # Add extra observation in `add` mode based on functor config + if "add" in self._env.observation_manager.active_functors: + for functor_name in self._env.observation_manager.active_functors["add"]: + functor_cfg = self._env.observation_manager.get_functor_cfg( + functor_name=functor_name + ) + if functor_cfg.func == get_object_uid: + obs_key = functor_cfg.name + asset_uid = functor_cfg.params["entity_cfg"].uid + asset = self._env.sim.get_asset(asset_uid) + if isinstance(asset, RigidObject): + features[obs_key]["names"] = asset_uid + elif isinstance(asset, (Articulation, Robot)): + link_names = asset.link_names + features[obs_key]["names"] = link_names + else: + logger.log_warning( + f"Asset with UID '{asset_uid}' is not RigidObject, Articulation or Robot. Cannot assign feature names based on asset properties." + ) def _convert_frame_to_lerobot( self, obs: TensorDict, action: TensorDict | torch.Tensor, task: str