feat(rollout): go back to initial position is now a config

This commit is contained in:
Steven Palma
2026-04-27 16:09:30 +02:00
parent bef67960d0
commit 7451813da6
7 changed files with 31 additions and 9 deletions
+6
View File
@@ -209,6 +209,12 @@ class RolloutConfig:
# Rename map for mapping robot/dataset observation keys to policy keys # Rename map for mapping robot/dataset observation keys to policy keys
rename_map: dict[str, str] = field(default_factory=dict) rename_map: dict[str, str] = field(default_factory=dict)
# Hardware teardown
# When True (default), smoothly interpolate the robot back to the joint
# positions captured at startup before disconnecting. Set to False to
# leave the robot in its final achieved pose at shutdown.
return_to_initial_position: bool = True
# Torch compile # Torch compile
use_torch_compile: bool = False use_torch_compile: bool = False
torch_compile_backend: str = "inductor" torch_compile_backend: str = "inductor"
+4 -1
View File
@@ -78,5 +78,8 @@ class BaseStrategy(RolloutStrategy):
def teardown(self, ctx: RolloutContext) -> None: def teardown(self, ctx: RolloutContext) -> None:
"""Disconnect hardware and stop inference.""" """Disconnect hardware and stop inference."""
self._teardown_hardware(ctx.hardware) self._teardown_hardware(
ctx.hardware,
return_to_initial_position=ctx.runtime.cfg.return_to_initial_position,
)
logger.info("Base strategy teardown complete") logger.info("Base strategy teardown complete")
+7 -3
View File
@@ -116,16 +116,20 @@ class RolloutStrategy(abc.ABC):
engine.resume() engine.resume()
return False return False
def _teardown_hardware(self, hw: HardwareContext) -> None: def _teardown_hardware(self, hw: HardwareContext, return_to_initial_position: bool = True) -> None:
"""Stop the inference engine, return robot to initial position, and disconnect hardware.""" """Stop the inference engine, optionally return robot to initial position, and disconnect hardware."""
if self._engine is not None: if self._engine is not None:
logger.info("Stopping inference engine...") logger.info("Stopping inference engine...")
self._engine.stop() self._engine.stop()
robot = hw.robot_wrapper.inner robot = hw.robot_wrapper.inner
if robot.is_connected: if robot.is_connected:
if hw.initial_position: if return_to_initial_position and hw.initial_position:
logger.info("Returning robot to initial position before shutdown...") logger.info("Returning robot to initial position before shutdown...")
self._return_to_initial_position(hw) self._return_to_initial_position(hw)
elif not return_to_initial_position:
logger.info(
"Skipping return-to-initial-position (disabled by config); leaving robot in final pose."
)
logger.info("Disconnecting robot...") logger.info("Disconnecting robot...")
robot.disconnect() robot.disconnect()
teleop = hw.teleop teleop = hw.teleop
+4 -1
View File
@@ -371,7 +371,10 @@ class DAggerStrategy(RolloutStrategy):
logger.info("Dataset uploaded to hub") logger.info("Dataset uploaded to hub")
log_say("Dataset uploaded to hub", play_sounds) log_say("Dataset uploaded to hub", play_sounds)
self._teardown_hardware(ctx.hardware) self._teardown_hardware(
ctx.hardware,
return_to_initial_position=ctx.runtime.cfg.return_to_initial_position,
)
logger.info("DAgger strategy teardown complete") logger.info("DAgger strategy teardown complete")
# ------------------------------------------------------------------ # ------------------------------------------------------------------
+4 -1
View File
@@ -227,7 +227,10 @@ class HighlightStrategy(RolloutStrategy):
logger.info("Dataset uploaded to hub") logger.info("Dataset uploaded to hub")
log_say("Dataset uploaded to hub", play_sounds) log_say("Dataset uploaded to hub", play_sounds)
self._teardown_hardware(ctx.hardware) self._teardown_hardware(
ctx.hardware,
return_to_initial_position=ctx.runtime.cfg.return_to_initial_position,
)
logger.info("Highlight strategy teardown complete") logger.info("Highlight strategy teardown complete")
def _setup_keyboard(self, shutdown_event: ThreadingEvent) -> None: def _setup_keyboard(self, shutdown_event: ThreadingEvent) -> None:
+4 -1
View File
@@ -196,7 +196,10 @@ class SentryStrategy(RolloutStrategy):
logger.info("Dataset uploaded to hub") logger.info("Dataset uploaded to hub")
log_say("Dataset uploaded to hub", play_sounds) log_say("Dataset uploaded to hub", play_sounds)
self._teardown_hardware(ctx.hardware) self._teardown_hardware(
ctx.hardware,
return_to_initial_position=ctx.runtime.cfg.return_to_initial_position,
)
logger.info("Sentry strategy teardown complete") logger.info("Sentry strategy teardown complete")
def _background_push(self, dataset, cfg) -> None: def _background_push(self, dataset, cfg) -> None:
+2 -2
View File
@@ -254,7 +254,7 @@ def test_create_inference_engine_sync():
def test_estimate_max_episode_seconds_no_video(): def test_estimate_max_episode_seconds_no_video():
from lerobot.rollout.strategies import estimate_max_episode_seconds from lerobot.rollout.strategies import estimate_max_episode_seconds
assert estimate_max_episode_seconds({}, fps=30.0) == 600.0 assert estimate_max_episode_seconds({}, fps=30.0) == 300.0
def test_estimate_max_episode_seconds_with_video(): def test_estimate_max_episode_seconds_with_video():
@@ -264,7 +264,7 @@ def test_estimate_max_episode_seconds_with_video():
result = estimate_max_episode_seconds(features, fps=30.0) result = estimate_max_episode_seconds(features, fps=30.0)
assert result > 0 assert result > 0
# With a real camera, duration should differ from the fallback # With a real camera, duration should differ from the fallback
assert result != 600.0 assert result != 300.0
def test_safe_push_to_hub(): def test_safe_push_to_hub():