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,