diff --git a/src/lerobot/envs/utils.py b/src/lerobot/envs/utils.py index 6e6f352e9..8b9c4f94b 100644 --- a/src/lerobot/envs/utils.py +++ b/src/lerobot/envs/utils.py @@ -126,6 +126,26 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten if "camera_obs" in observations: return_observations[f"{OBS_STR}.camera_obs"] = observations["camera_obs"] + # Pass through any remaining ndarray/tensor keys not already handled above, + # so env plugins can expose extra observation keys via get_env_processors(). + _handled = {"pixels", "environment_state", "agent_pos", "robot_state", "policy", "camera_obs"} + for key, value in observations.items(): + if key in _handled: + continue + target = f"{OBS_STR}.{key}" + if target in return_observations: + continue + if isinstance(value, np.ndarray): + val = torch.from_numpy(value).float() + if val.dim() == 1: + val = val.unsqueeze(0) + return_observations[target] = val + elif isinstance(value, Tensor): + val = value.float() + if val.dim() == 1: + val = val.unsqueeze(0) + return_observations[target] = val + return return_observations diff --git a/src/lerobot/utils/import_utils.py b/src/lerobot/utils/import_utils.py index 5dbce2c5b..b0d894c04 100644 --- a/src/lerobot/utils/import_utils.py +++ b/src/lerobot/utils/import_utils.py @@ -216,9 +216,15 @@ def register_third_party_plugins() -> None: This function uses `importlib.metadata` to find packages installed in the environment (including editable installs) starting with 'lerobot_robot_', 'lerobot_camera_', - 'lerobot_teleoperator_', or 'lerobot_policy_' and imports them. + 'lerobot_teleoperator_', 'lerobot_policy_', or 'lerobot_env_' and imports them. """ - prefixes = ("lerobot_robot_", "lerobot_camera_", "lerobot_teleoperator_", "lerobot_policy_") + prefixes = ( + "lerobot_robot_", + "lerobot_camera_", + "lerobot_teleoperator_", + "lerobot_policy_", + "lerobot_env_", + ) imported: list[str] = [] failed: list[str] = []