From 0d60a855be5fcb42c801c058ed21e14b9c6f6a9c Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Mon, 4 May 2026 19:33:45 +0200 Subject: [PATCH] fix(rl): record pre-step observation so (obs, action, next.reward) align in gym_manipulator dataset --- src/lerobot/rl/gym_manipulator.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/lerobot/rl/gym_manipulator.py b/src/lerobot/rl/gym_manipulator.py index bcf937e46..ab43aae65 100644 --- a/src/lerobot/rl/gym_manipulator.py +++ b/src/lerobot/rl/gym_manipulator.py @@ -698,6 +698,12 @@ def control_loop( if use_gripper: neutral_action = torch.cat([neutral_action, torch.tensor([1.0])]) # Gripper stay + observation = { + k: v.squeeze(0).cpu() + for k, v in transition[TransitionKey.OBSERVATION].items() + if isinstance(v, torch.Tensor) + } + transition = step_env_and_process_transition( env=env, transition=transition, @@ -709,16 +715,11 @@ def control_loop( truncated = transition.get(TransitionKey.TRUNCATED, False) if cfg.mode == "record": - observations = { - k: v.squeeze(0).cpu() - for k, v in transition[TransitionKey.OBSERVATION].items() - if isinstance(v, torch.Tensor) - } action_to_record = transition[TransitionKey.COMPLEMENTARY_DATA].get( "teleop_action", transition[TransitionKey.ACTION] ) frame = { - **observations, + **observation, ACTION: action_to_record.cpu(), REWARD: np.array([transition[TransitionKey.REWARD]], dtype=np.float32), DONE: np.array([terminated or truncated], dtype=bool),