mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-26 03:06:01 +00:00
feat(rollout): go back to initial position is now a config
This commit is contained in:
@@ -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"
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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():
|
||||||
|
|||||||
Reference in New Issue
Block a user