fix pi052 runtime and training safety

tests/runtime/test_language_runtime.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from lerobot.runtime import (
-    LanguageConditionedRuntime,
-)
+import threading
+import time
+
+from lerobot.runtime import LanguageConditionedRuntime, Tick
 
 
 class FakeAdapter:
@@ -69,3 +70,31 @@ def test_runtime_handles_user_interjection():
 
     assert "please say ok" in adapter.interjections
     assert runtime.state.language_context["plan"] == "new plan"
+
+
+def test_prompt_change_discards_in_flight_action_chunk():
+    started = threading.Event()
+    release = threading.Event()
+
+    class BlockingAdapter(FakeAdapter):
+        def select_action(self, observation, state):
+            started.set()
+            assert release.wait(timeout=2)
+            return ["stale"]
+
+    runtime = LanguageConditionedRuntime(
+        policy_adapter=BlockingAdapter(),
+        observation_provider=lambda: {"observation.state": 1},
+    )
+    runtime.set_task("old task")
+    runtime.state.tick = Tick(index=1, monotonic_seconds=time.monotonic())
+    inference = threading.Thread(target=runtime.maybe_enqueue_action_chunk, kwargs={"force": True})
+    inference.start()
+    assert started.wait(timeout=2)
+
+    runtime.set_task("new task")
+    release.set()
+    inference.join(timeout=2)
+
+    assert not inference.is_alive()
+    assert list(runtime.state.action_queue) == []

tests/runtime/test_sim_robocasa.py

@@ -17,6 +17,7 @@ from types import SimpleNamespace
 
 import numpy as np
 
+from lerobot.runtime.sim_robocasa import RoboCasaSimBackend
 from lerobot.utils.video_annotation import annotate_frame
 
 
@@ -59,3 +60,23 @@ def test_overlay_draws_each_label_once(monkeypatch):
     assert all(color == (255, 255, 255) and thickness == 1 for _, color, thickness in put_text_calls)
     assert len(rectangle_calls) == 1
     assert not np.shares_memory(annotated, frame)
+
+
+def test_capture_updates_live_frame_when_recording_is_disabled(monkeypatch):
+    backend = object.__new__(RoboCasaSimBackend)
+    frame = np.full((8, 8, 3), 42, dtype=np.uint8)
+    written = []
+    backend.record = False
+    backend.runtime_state = None
+    backend._multiview_frame = lambda: frame
+    backend._current_task = lambda: "task"
+    backend._subtask_getter = None
+    backend._memory_getter = None
+    backend._latest_frame = None
+    backend._write_live_frame = written.append
+    monkeypatch.setattr("lerobot.runtime.sim_robocasa.annotate_frame", lambda image, labels: image)
+
+    backend._capture_frame()
+
+    assert backend._latest_frame is frame
+    assert written == [frame]