Reset rollout state after robot episode end

tests/policies/test_relative_actions.py

@@ -171,6 +171,23 @@ def test_full_pipeline_roundtrip(dataset, action_dim):
     torch.testing.assert_close(recovered_actions, original_actions, atol=1e-4, rtol=1e-4)
 
 
+def test_relative_actions_processor_reset_clears_cached_state():
+    relative_step = RelativeActionsProcessorStep(enabled=True)
+    transition = batch_to_transition(
+        {
+            OBS_STATE: torch.tensor([[1.0, 2.0]]),
+            ACTION: torch.tensor([[[1.5, 1.0]]]),
+        }
+    )
+
+    relative_step(transition)
+    assert relative_step.get_cached_state() is not None
+
+    relative_step.reset()
+
+    assert relative_step.get_cached_state() is None
+
+
 def test_normalized_relative_values_are_reasonable(dataset, action_dim):
     """With correct chunk stats, normalized relative actions should be in a reasonable range."""
     action_chunks, states = _build_action_chunks(dataset, CHUNK_SIZE)

tests/test_rollout.py

@@ -177,6 +177,26 @@ def test_thread_safe_robot_delegates():
     robot.disconnect()
 
 
+def test_thread_safe_robot_tracks_action_response_metadata():
+    from lerobot.rollout.robot_wrapper import ThreadSafeRobot
+
+    class MetadataRobot:
+        def get_observation(self):
+            return {}
+
+        def send_action(self, action):
+            self.last_action_response = {"action": action, "done": True, "success": False}
+            return action
+
+    robot = MetadataRobot()
+    wrapper = ThreadSafeRobot(robot)
+
+    action = {"motor_1.pos": 0.0, "motor_2.pos": 1.0, "motor_3.pos": 2.0}
+    assert wrapper.send_action(action) == action
+    assert wrapper.pop_last_action_response() == {"action": action, "done": True, "success": False}
+    assert wrapper.pop_last_action_response() is None
+
+
 def test_thread_safe_robot_properties():
     from lerobot.rollout.robot_wrapper import ThreadSafeRobot
     from tests.mocks.mock_robot import MockRobot, MockRobotConfig
@@ -194,6 +214,42 @@ def test_thread_safe_robot_properties():
     robot.disconnect()
 
 
+def test_base_strategy_resets_inference_when_robot_reports_episode_done():
+    from lerobot.rollout import BaseStrategy, BaseStrategyConfig
+
+    strategy = BaseStrategy(BaseStrategyConfig())
+    strategy._engine = MagicMock()
+    strategy._interpolator = MagicMock()
+    strategy._cached_obs_processed = {"stale": True}
+
+    ctx = MagicMock()
+    ctx.hardware.robot_wrapper.pop_last_action_response.return_value = {"done": True, "success": False}
+
+    strategy._reset_inference_after_robot_episode_done(ctx)
+
+    strategy._engine.reset.assert_called_once()
+    strategy._interpolator.reset.assert_called_once()
+    assert strategy._cached_obs_processed is None
+
+
+def test_base_strategy_ignores_action_responses_without_episode_done():
+    from lerobot.rollout import BaseStrategy, BaseStrategyConfig
+
+    strategy = BaseStrategy(BaseStrategyConfig())
+    strategy._engine = MagicMock()
+    strategy._interpolator = MagicMock()
+    strategy._cached_obs_processed = {"cached": True}
+
+    ctx = MagicMock()
+    ctx.hardware.robot_wrapper.pop_last_action_response.return_value = {"done": False}
+
+    strategy._reset_inference_after_robot_episode_done(ctx)
+
+    strategy._engine.reset.assert_not_called()
+    strategy._interpolator.reset.assert_not_called()
+    assert strategy._cached_obs_processed == {"cached": True}
+
+
 # ---------------------------------------------------------------------------
 # Strategy factory
 # ---------------------------------------------------------------------------