From 9b25b7fe0ae4a226a19e22a37f00131358e5dc45 Mon Sep 17 00:00:00 2001 From: Xingdong Zuo Date: Tue, 28 Jul 2026 19:52:02 +0900 Subject: [PATCH] feat(lekiwi): support LeKiwi in the rollout/eval CLI (#3742) * feat(lekiwi): support LeKiwi in the rollout/eval CLI Register the lekiwi robot in lerobot_rollout.py so policies can be rolled out on a LeKiwi, and keep base-velocity (.vel) features in build_rollout_context. LeKiwi's observation.state and action are 9-dim (6 arm .pos + x/y/theta.vel) and the policy is normalized on all 9. The old filter kept only .pos features, so it fed a 6-dim vector into a 9-dim normalizer (RuntimeError, size 6 vs 9) and silently dropped the base velocities from the action, leaving the base unable to move. Keeping both .pos and .vel fixes both. Pure-arm robots have no .vel keys, so this is a no-op for them. * style: format LeKiwi rollout action features --------- Co-authored-by: Steven Palma Co-authored-by: Steven Palma --- src/lerobot/rollout/context.py | 14 ++++++++++++-- src/lerobot/scripts/lerobot_rollout.py | 1 + 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/src/lerobot/rollout/context.py b/src/lerobot/rollout/context.py index 2bada502e..2b39ab871 100644 --- a/src/lerobot/rollout/context.py +++ b/src/lerobot/rollout/context.py @@ -294,12 +294,22 @@ def build_rollout_context( # ``observation_features`` values are either a tuple (camera shape) or the # ``float`` type itself used as a sentinel for scalar motor features — # see ``dict[str, type | tuple]`` annotation on ``Robot.observation_features``. + # Keep cameras (tuple) plus both joint-position (.pos) and base-velocity (.vel) + # scalar state features. LeKiwi's observation.state is 9-dim (6 arm .pos + + # x/y/theta.vel) and the policy was trained/normalized on all 9; the old .pos-only + # filter fed a 6-dim state into a 9-dim normalizer → RuntimeError (size 6 vs 9). + # Pure-arm robots have no .vel state keys, so this is a no-op for them. observation_features_hw = { k: v for k, v in all_obs_features.items() - if isinstance(v, tuple) or (v is float and k.endswith(".pos")) + if isinstance(v, tuple) or (v is float and k.endswith((".pos", ".vel"))) } - action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith(".pos")} + # Keep both joint-position (.pos) and base-velocity (.vel) action features so + # mobile manipulators command the base too (e.g. LeKiwi: 6 arm .pos + + # x/y/theta.vel = 9-dim action). Pure-arm robots have no .vel keys, so this is + # a no-op for them. Without the .vel keys the base velocities are silently + # dropped from dataset_features[ACTION]/ordered_action_keys and the base never moves. + action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith((".pos", ".vel"))} # The action side is always needed: sync inference reads action names from # ``dataset_features[ACTION]`` to map policy tensors back to robot actions. diff --git a/src/lerobot/scripts/lerobot_rollout.py b/src/lerobot/scripts/lerobot_rollout.py index 879070721..5456a463e 100644 --- a/src/lerobot/scripts/lerobot_rollout.py +++ b/src/lerobot/scripts/lerobot_rollout.py @@ -165,6 +165,7 @@ from lerobot.robots import ( # noqa: F401 earthrover_mini_plus, hope_jr, koch_follower, + lekiwi, omx_follower, openarm_follower, reachy2,