Reset rollout state after robot episode end

src/lerobot/processor/relative_action_processor.py

@@ -146,6 +146,9 @@ class RelativeActionsProcessorStep(ProcessorStep):
         """Return the cached ``observation.state`` used as the reference point for relative/absolute action conversions."""
         return self._last_state
 
+    def reset(self) -> None:
+        self._last_state = None
+
     def get_config(self) -> dict[str, Any]:
         return {
             "enabled": self.enabled,

src/lerobot/rollout/robot_wrapper.py

@@ -36,6 +36,7 @@ class ThreadSafeRobot:
     def __init__(self, robot: Robot) -> None:
         self._robot = robot
         self._lock = Lock()
+        self._last_action_response: Any | None = None
 
     # -- Lock-protected I/O --------------------------------------------------
 
@@ -45,7 +46,15 @@ class ThreadSafeRobot:
 
     def send_action(self, action: dict[str, Any] | Any) -> Any:
         with self._lock:
-            return self._robot.send_action(action)
+            response = self._robot.send_action(action)
+            self._last_action_response = getattr(self._robot, "last_action_response", response)
+            return response
+
+    def pop_last_action_response(self) -> Any | None:
+        with self._lock:
+            response = self._last_action_response
+            self._last_action_response = None
+            return response
 
     # -- Read-only proxies (no lock needed) -----------------------------------
 

src/lerobot/rollout/strategies/base.py

@@ -67,6 +67,8 @@ class BaseStrategy(RolloutStrategy):
 
             action_dict = send_next_action(obs_processed, obs, ctx, interpolator)
             self._log_telemetry(obs_processed, action_dict, ctx.runtime)
+            if action_dict is not None:
+                self._reset_inference_after_robot_episode_done(ctx)
 
             dt = time.perf_counter() - loop_start
             if (sleep_t := control_interval - dt) > 0:

src/lerobot/rollout/strategies/core.py

@@ -116,6 +116,22 @@ class RolloutStrategy(abc.ABC):
             engine.resume()
         return False
 
+    def _reset_inference_after_robot_episode_done(self, ctx: RolloutContext) -> None:
+        """Reset rollout-side episode state when a robot backend reports an environment reset."""
+        response = ctx.hardware.robot_wrapper.pop_last_action_response()
+        if not isinstance(response, dict) or not response.get("done"):
+            return
+
+        logger.info(
+            "Robot reported episode done (success=%s); resetting rollout inference state",
+            response.get("success"),
+        )
+        if self._engine is not None:
+            self._engine.reset()
+        if self._interpolator is not None:
+            self._interpolator.reset()
+        self._cached_obs_processed = None
+
     def _teardown_hardware(self, hw: HardwareContext, return_to_initial_position: bool = True) -> None:
         """Stop the inference engine, optionally return robot to initial position, and disconnect hardware."""
         if self._engine is not None: