fix(rollout): robot_observation_processor and notify_observation at policy frequency instead of interpolator rate

Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>
This commit is contained in:
Steven Palma
2026-04-26 18:49:52 +02:00
parent eb9519eb91
commit 8c4d4c8208
5 changed files with 39 additions and 16 deletions
+1 -2
View File
@@ -60,8 +60,7 @@ class BaseStrategy(RolloutStrategy):
break break
obs = robot.get_observation() obs = robot.get_observation()
obs_processed = ctx.processors.robot_observation_processor(obs) obs_processed = self._process_observation_and_notify(ctx.processors, obs)
engine.notify_observation(obs_processed)
if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval): if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval):
continue continue
+26 -1
View File
@@ -32,7 +32,7 @@ from ..inference import InferenceEngine
if TYPE_CHECKING: if TYPE_CHECKING:
from ..configs import RolloutStrategyConfig from ..configs import RolloutStrategyConfig
from ..context import HardwareContext, RolloutContext, RuntimeContext from ..context import HardwareContext, ProcessorContext, RolloutContext, RuntimeContext
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -50,6 +50,7 @@ class RolloutStrategy(abc.ABC):
self._engine: InferenceEngine | None = None self._engine: InferenceEngine | None = None
self._interpolator: ActionInterpolator | None = None self._interpolator: ActionInterpolator | None = None
self._warmup_flushed: bool = False self._warmup_flushed: bool = False
self._cached_obs_processed: dict | None = None
def _init_engine(self, ctx: RolloutContext) -> None: def _init_engine(self, ctx: RolloutContext) -> None:
"""Attach the inference engine and action interpolator, then start the backend. """Attach the inference engine and action interpolator, then start the backend.
@@ -65,8 +66,32 @@ class RolloutStrategy(abc.ABC):
self._engine.reset() self._engine.reset()
self._engine.start() self._engine.start()
self._warmup_flushed = False self._warmup_flushed = False
self._cached_obs_processed = None
logger.info("Inference engine started") logger.info("Inference engine started")
def _process_observation_and_notify(self, processors: ProcessorContext, obs_raw: dict) -> dict:
"""Run the observation processor and notify the engine — throttled to policy ticks.
Callers are responsible for calling ``robot.get_observation()`` every loop
iteration so ``obs_raw`` stays fresh for the action post-processor. This
helper gates only the comparatively expensive bits — the processor pipeline
and ``engine.notify_observation`` — to fire when the interpolator signals
it needs a new action (once per ``interpolation_multiplier`` ticks). On
interpolated ticks the cached ``obs_processed`` is reused.
With ``interpolation_multiplier == 1`` this is equivalent to the unthrottled
path: ``needs_new_action()`` is True every tick.
The cache is implicitly invalidated whenever ``interpolator.reset()`` is
called (warmup completion, DAgger phase transitions back to AUTONOMOUS),
because reset makes ``needs_new_action()`` return True on the next call.
"""
if self._cached_obs_processed is None or self._interpolator.needs_new_action():
obs_processed = processors.robot_observation_processor(obs_raw)
self._engine.notify_observation(obs_processed)
self._cached_obs_processed = obs_processed
return self._cached_obs_processed
def _handle_warmup(self, use_torch_compile: bool, loop_start: float, control_interval: float) -> bool: def _handle_warmup(self, use_torch_compile: bool, loop_start: float, control_interval: float) -> bool:
"""Handle torch.compile warmup phase. """Handle torch.compile warmup phase.
+9 -8
View File
@@ -434,19 +434,19 @@ class DAggerStrategy(RolloutStrategy):
phase = events.phase phase = events.phase
obs = robot.get_observation() obs = robot.get_observation()
obs_processed = ctx.processors.robot_observation_processor(obs)
obs_frame = build_dataset_frame(features, obs_processed, prefix=OBS_STR)
# --- CORRECTING: human teleop control --- # --- CORRECTING: human teleop control ---
if phase == DAggerPhase.CORRECTING: if phase == DAggerPhase.CORRECTING:
obs_processed = ctx.processors.robot_observation_processor(obs)
teleop_action = teleop.get_action() teleop_action = teleop.get_action()
processed_teleop = ctx.processors.teleop_action_processor((teleop_action, obs)) processed_teleop = ctx.processors.teleop_action_processor((teleop_action, obs))
robot_action_to_send = ctx.processors.robot_action_processor((processed_teleop, obs)) robot_action_to_send = ctx.processors.robot_action_processor((processed_teleop, obs))
robot.send_action(robot_action_to_send) robot.send_action(robot_action_to_send)
last_action = robot_action_to_send last_action = robot_action_to_send
self._log_telemetry(obs_processed, processed_teleop, ctx.runtime) self._log_telemetry(obs_processed, processed_teleop, ctx.runtime)
action_frame = build_dataset_frame(features, processed_teleop, prefix=ACTION)
if record_tick % record_stride == 0: if record_tick % record_stride == 0:
obs_frame = build_dataset_frame(features, obs_processed, prefix=OBS_STR)
action_frame = build_dataset_frame(features, processed_teleop, prefix=ACTION)
frame = { frame = {
**obs_frame, **obs_frame,
**action_frame, **action_frame,
@@ -463,7 +463,7 @@ class DAggerStrategy(RolloutStrategy):
# --- AUTONOMOUS: policy control --- # --- AUTONOMOUS: policy control ---
else: else:
engine.notify_observation(obs_processed) obs_processed = self._process_observation_and_notify(ctx.processors, obs)
if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval): if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval):
continue continue
@@ -472,8 +472,9 @@ class DAggerStrategy(RolloutStrategy):
if action_dict is not None: if action_dict is not None:
self._log_telemetry(obs_processed, action_dict, ctx.runtime) self._log_telemetry(obs_processed, action_dict, ctx.runtime)
last_action = ctx.processors.robot_action_processor((action_dict, obs)) last_action = ctx.processors.robot_action_processor((action_dict, obs))
action_frame = build_dataset_frame(features, action_dict, prefix=ACTION)
if record_tick % record_stride == 0: if record_tick % record_stride == 0:
obs_frame = build_dataset_frame(features, obs_processed, prefix=OBS_STR)
action_frame = build_dataset_frame(features, action_dict, prefix=ACTION)
frame = { frame = {
**obs_frame, **obs_frame,
**action_frame, **action_frame,
@@ -608,10 +609,10 @@ class DAggerStrategy(RolloutStrategy):
phase = events.phase phase = events.phase
obs = robot.get_observation() obs = robot.get_observation()
obs_processed = ctx.processors.robot_observation_processor(obs)
# --- CORRECTING: human teleop control + recording --- # --- CORRECTING: human teleop control + recording ---
if phase == DAggerPhase.CORRECTING: if phase == DAggerPhase.CORRECTING:
obs_processed = ctx.processors.robot_observation_processor(obs)
teleop_action = teleop.get_action() teleop_action = teleop.get_action()
processed_teleop = ctx.processors.teleop_action_processor((teleop_action, obs)) processed_teleop = ctx.processors.teleop_action_processor((teleop_action, obs))
robot_action_to_send = ctx.processors.robot_action_processor((processed_teleop, obs)) robot_action_to_send = ctx.processors.robot_action_processor((processed_teleop, obs))
@@ -619,9 +620,9 @@ class DAggerStrategy(RolloutStrategy):
last_action = robot_action_to_send last_action = robot_action_to_send
self._log_telemetry(obs_processed, processed_teleop, ctx.runtime) self._log_telemetry(obs_processed, processed_teleop, ctx.runtime)
if record_tick % record_stride == 0:
obs_frame = build_dataset_frame(features, obs_processed, prefix=OBS_STR) obs_frame = build_dataset_frame(features, obs_processed, prefix=OBS_STR)
action_frame = build_dataset_frame(features, processed_teleop, prefix=ACTION) action_frame = build_dataset_frame(features, processed_teleop, prefix=ACTION)
if record_tick % record_stride == 0:
dataset.add_frame( dataset.add_frame(
{ {
**obs_frame, **obs_frame,
@@ -639,7 +640,7 @@ class DAggerStrategy(RolloutStrategy):
# --- AUTONOMOUS: policy control (no recording) --- # --- AUTONOMOUS: policy control (no recording) ---
else: else:
engine.notify_observation(obs_processed) obs_processed = self._process_observation_and_notify(ctx.processors, obs)
if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval): if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval):
continue continue
+1 -2
View File
@@ -135,8 +135,7 @@ class HighlightStrategy(RolloutStrategy):
break break
obs = robot.get_observation() obs = robot.get_observation()
obs_processed = ctx.processors.robot_observation_processor(obs) obs_processed = self._process_observation_and_notify(ctx.processors, obs)
engine.notify_observation(obs_processed)
if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval): if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval):
continue continue
+1 -2
View File
@@ -111,8 +111,7 @@ class SentryStrategy(RolloutStrategy):
break break
obs = robot.get_observation() obs = robot.get_observation()
obs_processed = ctx.processors.robot_observation_processor(obs) obs_processed = self._process_observation_and_notify(ctx.processors, obs)
engine.notify_observation(obs_processed)
if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval): if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval):
continue continue