diff --git a/changelog.d/3496-lerobot-async-rename-map-renames-the-camera-it-declares.md b/changelog.d/3496-lerobot-async-rename-map-renames-the-camera-it-declares.md new file mode 100644 index 000000000..cf33a141d --- /dev/null +++ b/changelog.d/3496-lerobot-async-rename-map-renames-the-camera-it-declares.md @@ -0,0 +1,3 @@ +### Fixed: `lerobot_async` `rename_map` renames the camera it declares + +Passing `rename_map={"observation.images.front": "observation.images.laptop"}` to `create_policy("lerobot_async", ...)` did not reach a stock lerobot `PolicyServer` as a rename: the client declared the camera under the robot's name and forwarded the map, but the server resizes every declared image by the checkpoint's own image features before its rename step runs, so the renamed camera was a `KeyError` on the server and every `get_actions` raised "server returned no actions". The client now applies an image rename itself, declaring and sending the camera under the model's name, so a checkpoint trained with `observation.images.laptop` is reachable from a robot whose camera is called `front`, as the docs describe. diff --git a/docs/policies/lerobot-async.md b/docs/policies/lerobot-async.md index f6feed6b9..d94fe5054 100644 --- a/docs/policies/lerobot-async.md +++ b/docs/policies/lerobot-async.md @@ -87,7 +87,7 @@ sim.run_policy( | `actions_per_step` | `actions_per_chunk` | Actions executed from one chunk before re-querying (the re-query interval). Positive `int`, or `None` for the default | | `connect_timeout` | `10.0` | Seconds to wait for the gRPC `Ready` handshake | | `request_timeout` | `60.0` | Seconds to wait for each observation/action RPC | -| `rename_map` | `{}` | `{robot_obs_key: model_feature_key}` forwarded to the server; renames observation keys before the policy sees them (async analog of `lerobot_local`'s `obs_rename`) | +| `rename_map` | `{}` | `{robot_obs_key: model_feature_key}` map. Camera entries (`observation.images.*`) are applied client-side (the server resizes images before its rename step); state entries are forwarded to the server. Async analog of `lerobot_local`'s `obs_rename` | ## Notes diff --git a/strands_robots/policies/lerobot_async/policy.py b/strands_robots/policies/lerobot_async/policy.py index 330675f21..86491c8e4 100644 --- a/strands_robots/policies/lerobot_async/policy.py +++ b/strands_robots/policies/lerobot_async/policy.py @@ -148,13 +148,15 @@ class LerobotAsyncPolicy(Policy): request_timeout: Seconds to wait for each observation/action RPC. Same domain; it also bounds ``Ready`` on :meth:`reset`, where a failure is logged rather than raised. - rename_map: Optional ``{robot_obs_key: model_feature_key}`` map forwarded - to the server's ``RemotePolicyConfig.rename_map``. The server applies - it as a ``RenameObservationsProcessorStep`` (renaming each matching - observation key to its mapped name) before the policy sees the - observation - the async analog of the ``lerobot_local`` provider's - ``obs_rename``. Use it when the checkpoint expects camera/state keys - that differ from the ones the robot exposes (e.g. + rename_map: Optional ``{robot_obs_key: model_feature_key}`` map. + Camera entries (``observation.images.*``) are applied **client-side**: + the handshake declares and the raw observation carries the image under + the model's feature name, because the server resizes every declared + image by ``policy_image_features`` before its + ``RenameObservationsProcessorStep`` runs (lerobot >= 0.6.1). State + entries are forwarded to the server's ``RemotePolicyConfig.rename_map`` + and applied there as usual. Use it when the checkpoint expects + camera/state keys that differ from the ones the robot exposes (e.g. ``{"observation.images.front": "observation.images.laptop"}``); keys not present in the map pass through unchanged. @@ -405,14 +407,25 @@ def close(self) -> None: # -- Observation / action wire conversion --------------------------------- def _camera_items(self, observation_dict: dict[str, Any]) -> list[tuple[str, np.ndarray]]: - """Return ``(key, HWC array)`` pairs for RGB/depth camera entries.""" + """Return ``(wire key, HWC array)`` pairs for RGB/depth camera entries. + + A camera whose ``observation.images.`` feature is renamed by + ``rename_map`` is declared and sent under the model's name. The server + resizes every declared image by the checkpoint's own image features + (``prepare_raw_observation``) BEFORE its rename step runs, so a camera + declared under the robot's name is a ``KeyError`` there, not a rename. + """ + from lerobot.utils.constants import OBS_IMAGES + cams: list[tuple[str, np.ndarray]] = [] for key, value in observation_dict.items(): if key in self.robot_state_keys or key == "task": continue arr = np.asarray(value) if arr.ndim == 3 and arr.shape[2] in (1, 3): - cams.append((key, arr)) + target = self.rename_map.get(f"{OBS_IMAGES}.{key}", "") + wire_key = target.removeprefix(f"{OBS_IMAGES}.") if target.startswith(f"{OBS_IMAGES}.") else key + cams.append((wire_key, arr)) return cams def _build_lerobot_features(self, observation_dict: dict[str, Any]) -> dict[str, Any]: diff --git a/tests/policies/lerobot_async/test_rename_map_renames_the_camera_the_server_resizes.py b/tests/policies/lerobot_async/test_rename_map_renames_the_camera_the_server_resizes.py new file mode 100644 index 000000000..f93604920 --- /dev/null +++ b/tests/policies/lerobot_async/test_rename_map_renames_the_camera_the_server_resizes.py @@ -0,0 +1,55 @@ +"""A ``rename_map`` camera entry is applied before the server resizes the image. + +lerobot's ``PolicyServer`` runs ``prepare_raw_observation`` on every declared +``observation.images.`` and looks the key up in the checkpoint's own image +features to pick the resize target - BEFORE the ``RenameObservationsProcessorStep`` +that ``rename_map`` configures. A camera declared under the robot's name is a +``KeyError`` there (``Error in StreamActions: 'observation.images.front'`` on a +stock server), and the client raises "server returned no actions". + +So the client applies an image rename itself: the handshake declares the +model's feature name and the raw observation carries the image under the +matching wire key. Fails on pre-fix code, which declared and sent ``front``. +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("lerobot") + +import numpy as np + +from strands_robots.policies.lerobot_async import LerobotAsyncPolicy + +STATE_KEYS = ["j0", "j1"] + + +def _policy() -> LerobotAsyncPolicy: + policy = LerobotAsyncPolicy( + server_address="h:1", + policy_type="act", + pretrained_name_or_path="x/y", + rename_map={"observation.images.front": "observation.images.laptop"}, + ) + policy.set_robot_state_keys(STATE_KEYS) + return policy + + +def _observation() -> dict[str, object]: + obs: dict[str, object] = {k: 0.0 for k in STATE_KEYS} + obs["front"] = np.zeros((8, 8, 3), dtype=np.uint8) + obs["wrist"] = np.zeros((8, 8, 3), dtype=np.uint8) + return obs + + +def test_handshake_declares_the_model_camera_name() -> None: + features = _policy()._build_lerobot_features(_observation()) + image_keys = sorted(k for k in features if k.startswith("observation.images.")) + assert image_keys == ["observation.images.laptop", "observation.images.wrist"] + + +def test_raw_observation_carries_the_image_under_the_model_name() -> None: + raw = _policy()._to_raw_observation(_observation(), "") + assert "laptop" in raw and "front" not in raw + assert "wrist" in raw