diff --git a/egomimic/hydra_configs/robot/yam_rl2_earbuds_hpt300_rollout.yaml b/egomimic/hydra_configs/robot/yam_rl2_earbuds_hpt300_rollout.yaml index 1308e2db2..8e07d6fd1 100644 --- a/egomimic/hydra_configs/robot/yam_rl2_earbuds_hpt300_rollout.yaml +++ b/egomimic/hydra_configs/robot/yam_rl2_earbuds_hpt300_rollout.yaml @@ -108,6 +108,70 @@ policy: device: cuda:0 use_ema: false max_valid_samples: 8 + # Station start values for the model's own controls (skipped when not exposed): + # sampling steps at the trained/validated 50 (ARC contracts open at 10). + inference_defaults: + inference_steps: 50 + # Station-owned rollout knobs, shown on the dashboard for any model that does + # not declare a control of the same name; decoder knobs only when the loaded + # decoder has the attribute (ARC first-stream decoder; the time retimer + # declares its own speed controls). None of them changes what was trained. + station_controls: + replan_every: + label: Actions before replanning + description: Execute this prefix before predicting again (used when fastest-stream termination is 0). + type: integer + min: 1 + max: 100 + step: 1 + default: 30 + target: + kind: policy_attribute + attribute_path: replan_every + multistream_fastest_stream: + label: Fastest-stream termination + description: 1 ends the chunk when the first moving stream reaches its waypoint cap; 0 executes the fixed prefix of Actions before replanning. + type: integer + min: 0 + max: 1 + step: 1 + default: 1 + target: + kind: decoder_attribute + attribute_path: first_stream + execute_waypoint_percent: + label: Execute first % of waypoints + description: Waypoint cap while fastest-stream termination is on; must be a whole prefix of the token's waypoints. + type: integer + min: 2 + max: 100 + step: 1 + default: 50 + target: + kind: decoder_attribute + attribute_path: execute_percent + arc_speed_percent: + label: Replay speed (%) + description: Decoder replay tempo of the moving phases, not a learned speed condition. + type: integer + min: 25 + max: 400 + step: 5 + default: 100 + target: + kind: decoder_attribute + attribute_path: speed_percent + arc_hold_speed_percent: + label: Hold speed (%) + description: Replay tempo of holds (grasp dwell, placement); 100 preserves the demonstrated hold time while the moving phases change tempo. + type: integer + min: 25 + max: 400 + step: 5 + default: 100 + target: + kind: decoder_attribute + attribute_path: hold_speed_percent adapter: _target_: egomimic.robot.graph_policy.CartesianGraphAdapter embodiment_id: 7 diff --git a/egomimic/hydra_configs/robot/yam_rl2_hptflow_rollout.yaml b/egomimic/hydra_configs/robot/yam_rl2_hptflow_rollout.yaml index 0f8b58c0d..4c49f3fa6 100644 --- a/egomimic/hydra_configs/robot/yam_rl2_hptflow_rollout.yaml +++ b/egomimic/hydra_configs/robot/yam_rl2_hptflow_rollout.yaml @@ -115,8 +115,73 @@ policy: # HPT-Flow is stochastic. Reject an invalid whole plan and resample rather # than clamping an individual gripper; abort before robot commands if all fail. max_valid_samples: 8 + # Station start values for the model's own controls (skipped when not exposed): + # sampling steps at the trained/validated 50 (ARC contracts open at 10). + inference_defaults: + inference_steps: 50 + # Station-owned rollout knobs, shown on the dashboard for any model that does + # not declare a control of the same name; decoder knobs only when the loaded + # decoder has the attribute (ARC first-stream decoder; the time retimer + # declares its own speed controls). None of them changes what was trained. + station_controls: + replan_every: + label: Actions before replanning + description: Execute this prefix before predicting again (used when fastest-stream termination is 0). + type: integer + min: 1 + max: 100 + step: 1 + default: 30 + target: + kind: policy_attribute + attribute_path: replan_every + multistream_fastest_stream: + label: Fastest-stream termination + description: 1 ends the chunk when the first moving stream reaches its waypoint cap; 0 executes the fixed prefix of Actions before replanning. + type: integer + min: 0 + max: 1 + step: 1 + default: 1 + target: + kind: decoder_attribute + attribute_path: first_stream + execute_waypoint_percent: + label: Execute first % of waypoints + description: Waypoint cap while fastest-stream termination is on; must be a whole prefix of the token's waypoints. + type: integer + min: 2 + max: 100 + step: 1 + default: 50 + target: + kind: decoder_attribute + attribute_path: execute_percent + arc_speed_percent: + label: Replay speed (%) + description: Decoder replay tempo of the moving phases, not a learned speed condition. + type: integer + min: 25 + max: 400 + step: 5 + default: 100 + target: + kind: decoder_attribute + attribute_path: speed_percent + arc_hold_speed_percent: + label: Hold speed (%) + description: Replay tempo of holds (grasp dwell, placement); 100 preserves the demonstrated hold time while the moving phases change tempo. + type: integer + min: 25 + max: 400 + step: 5 + default: 100 + target: + kind: decoder_attribute + attribute_path: hold_speed_percent # History, shapes, decoder and typed sampler/replanning controls belong to - # the selected checkpoint's model-owned, bound inference-config.yaml. + # the selected checkpoint's model-owned, bound inference-config.yaml; + # station_controls above only fill in rollout knobs a model does not declare. adapter: _target_: egomimic.robot.graph_policy.CartesianGraphAdapter embodiment_id: 7 diff --git a/egomimic/hydra_configs/robot/yam_rl2_towels_hpt300_rollout.yaml b/egomimic/hydra_configs/robot/yam_rl2_towels_hpt300_rollout.yaml index 608ec5682..5183ca457 100644 --- a/egomimic/hydra_configs/robot/yam_rl2_towels_hpt300_rollout.yaml +++ b/egomimic/hydra_configs/robot/yam_rl2_towels_hpt300_rollout.yaml @@ -108,6 +108,70 @@ policy: device: cuda:0 use_ema: false max_valid_samples: 8 + # Station start values for the model's own controls (skipped when not exposed): + # sampling steps at the trained/validated 50 (ARC contracts open at 10). + inference_defaults: + inference_steps: 50 + # Station-owned rollout knobs, shown on the dashboard for any model that does + # not declare a control of the same name; decoder knobs only when the loaded + # decoder has the attribute (ARC first-stream decoder; the time retimer + # declares its own speed controls). None of them changes what was trained. + station_controls: + replan_every: + label: Actions before replanning + description: Execute this prefix before predicting again (used when fastest-stream termination is 0). + type: integer + min: 1 + max: 100 + step: 1 + default: 30 + target: + kind: policy_attribute + attribute_path: replan_every + multistream_fastest_stream: + label: Fastest-stream termination + description: 1 ends the chunk when the first moving stream reaches its waypoint cap; 0 executes the fixed prefix of Actions before replanning. + type: integer + min: 0 + max: 1 + step: 1 + default: 1 + target: + kind: decoder_attribute + attribute_path: first_stream + execute_waypoint_percent: + label: Execute first % of waypoints + description: Waypoint cap while fastest-stream termination is on; must be a whole prefix of the token's waypoints. + type: integer + min: 2 + max: 100 + step: 1 + default: 50 + target: + kind: decoder_attribute + attribute_path: execute_percent + arc_speed_percent: + label: Replay speed (%) + description: Decoder replay tempo of the moving phases, not a learned speed condition. + type: integer + min: 25 + max: 400 + step: 5 + default: 100 + target: + kind: decoder_attribute + attribute_path: speed_percent + arc_hold_speed_percent: + label: Hold speed (%) + description: Replay tempo of holds (grasp dwell, placement); 100 preserves the demonstrated hold time while the moving phases change tempo. + type: integer + min: 25 + max: 400 + step: 5 + default: 100 + target: + kind: decoder_attribute + attribute_path: hold_speed_percent adapter: _target_: egomimic.robot.graph_policy.CartesianGraphAdapter embodiment_id: 7 diff --git a/egomimic/pipeline/inference_controls.py b/egomimic/pipeline/inference_controls.py index ad2403b8b..f949e3a97 100644 --- a/egomimic/pipeline/inference_controls.py +++ b/egomimic/pipeline/inference_controls.py @@ -128,12 +128,29 @@ def public(self): } -def configure_profile_controls(graph, training, inference_profiles, *, decoder=None): - """Bind declared settings by stable stage ID, with atomic preflight.""" +def configure_profile_controls( + graph, training, inference_profiles, *, decoder=None, station_controls=None +): + """Bind declared settings by stable stage ID, with atomic preflight. + + ``station_controls`` are rollout-only settings declared by the station YAML. + Each is added only when the model declares no control of that name and, for + a decoder target, the bound decoder has that attribute; the checkpoint-bound + model declaration itself is never changed. + """ _, profile, _ = resolve_inference_profile(training, inference_profiles) stage_ids = OmegaConf.select(training, "model.pipeline.stage_ids", default={}) + overrides = dict(profile.get("overrides", {})) + for name, spec in (station_controls or {}).items(): + target = spec.get("target", {}) if isinstance(spec, Mapping) else {} + if name in overrides or ( + target.get("kind") == "decoder_attribute" + and not hasattr(decoder, str(target.get("attribute_path"))) + ): + continue + overrides[name] = spec bindings = [] - for name, spec in profile.get("overrides", {}).items(): + for name, spec in overrides.items(): validate_control(name, spec, stage_ids=stage_ids) target = spec["target"] kind, path = target["kind"], target["attribute_path"] diff --git a/egomimic/rldb/zarr/cartesian_arc.py b/egomimic/rldb/zarr/cartesian_arc.py index 59a60c548..2ed465483 100644 --- a/egomimic/rldb/zarr/cartesian_arc.py +++ b/egomimic/rldb/zarr/cartesian_arc.py @@ -441,6 +441,20 @@ def control_steps( ): if max_steps is not None and (int(max_steps) != max_steps or max_steps < 1): raise ValueError("max_steps must be a positive integer") + duration = self.control_duration( + token, initial_action=initial_action, clocks_in_seconds=clocks_in_seconds + ) + if duration is None: + if max_steps is None: + raise ValueError( + "ARC token timing does not define a finite replan boundary" + ) + return int(max_steps) + steps = max(1, int(np.ceil(duration / self.tokenizer.config.dt - 1e-9))) + return min(steps, int(max_steps)) if max_steps is not None else steps + + def control_duration(self, token, *, initial_action, clocks_in_seconds=False): + """Seconds until the first moving stream ends; None if that is undefined.""" value = self._native(token, initial_action, clocks_in_seconds) moving_ends, static_ends = [], [] for index, name in enumerate(self._stream_names): @@ -451,28 +465,29 @@ def control_steps( moving_ends.append(end) elif self.clock_mode == "duration": static_ends.append(end) - if moving_ends: - duration = min(moving_ends) - elif static_ends: - duration = min(static_ends) - else: - if max_steps is None: - raise ValueError( - "A fully static velocity token has no duration; supply max_steps" - ) - return int(max_steps) - if not np.isfinite(duration): - if max_steps is None: - raise ValueError( - "ARC token timing does not define a finite replan boundary" - ) - return int(max_steps) - steps = max(1, int(np.ceil(duration / self.tokenizer.config.dt - 1e-9))) - return min(steps, int(max_steps)) if max_steps is not None else steps + ends = moving_ends or static_ends + duration = min(ends) if ends else np.inf + return float(duration) if np.isfinite(duration) else None - def _clock_indices(self, durations, action_horizon): + def stream_clocks(self, token, *, initial_action, clocks_in_seconds=False): + """Decode-only: per arm, (arrival seconds, cumulative distance) of the + primary stream (translation if present), each of length M + 1.""" + value = self._native(token, initial_action, clocks_in_seconds) + anchor = _initial_actions(initial_action) + names = self._stream_names + name = "translation" if "translation" in names else names[0] + index = names.index(name) + out = [] + for arm, off in enumerate((0, 7)): + distance = self._distances(value, off, name, anchor) + durations = self._intervals(distance, value[:, 14 + index * 2 + arm]) + out.append( + (np.r_[0.0, np.cumsum(durations)], np.r_[0.0, np.cumsum(distance)]) + ) + return out + + def _clock_indices(self, durations, query): elapsed = np.r_[0.0, np.cumsum(durations)] - query = self.tokenizer.config.dt * np.arange(1, action_horizon + 1) upper = np.clip( np.searchsorted(elapsed, query, side="right"), 1, len(elapsed) - 1 ) @@ -488,8 +503,16 @@ def _clock_indices(self, durations, action_horizon): return lower, upper, np.clip(alpha, 0.0, 1.0), elapsed def detokenize( - self, arc_actions, action_horizon, *, initial_action, clocks_in_seconds=False + self, + arc_actions, + action_horizon, + *, + initial_action, + clocks_in_seconds=False, + times=None, ): + """``times`` (decode-only): token-clock seconds to read the H rows at, + replacing the uniform next-frame grid (dt, 2 dt, ..., H dt).""" value = self._native(arc_actions, initial_action, clocks_in_seconds) h = int(action_horizon) if h < 1: @@ -497,6 +520,10 @@ def detokenize( anchor = _initial_actions(initial_action) geometry = np.concatenate((anchor[None], value[:, :14]), axis=0) query = np.arange(1, h + 1) * self.tokenizer.config.dt + if times is not None: + query = np.asarray(times, dtype=np.float64) + if query.shape != (h,) or not np.isfinite(query).all(): + raise ValueError(f"times must be {h} finite clock readings") arms = [] for arm, off in enumerate((0, 7)): indices, clocks = {}, {} @@ -505,7 +532,7 @@ def detokenize( self._distances(value, off, name, anchor), value[:, 14 + index * 2 + arm], ) - lower, upper, alpha, elapsed = self._clock_indices(durations, h) + lower, upper, alpha, elapsed = self._clock_indices(durations, query) indices[name] = (lower, upper, alpha) clocks[name] = elapsed primary = indices.get("translation", indices.get("rotation")) diff --git a/egomimic/robot/AGENTS.md b/egomimic/robot/AGENTS.md index 124ba06cf..e2084ede7 100644 --- a/egomimic/robot/AGENTS.md +++ b/egomimic/robot/AGENTS.md @@ -70,6 +70,11 @@ for Eva/Yam setup, Quest app build requirements, data format and operator comman Two, four or six timing channels follow all 14 pose channels; they are not interleaved. Decoder settings belong to explicit model YAML `decoder_attribute` controls; never append controls based on a model class or inferred token width. The + station YAML's `policy.station_controls` adds rollout-only knobs (replan, + fastest stream, waypoint %, replay/hold speed) where the model declares none + of that name and the loaded decoder has the target attribute, and + `policy.inference_defaults` sets start values of exposed controls; neither + edits the checkpoint-bound declaration. The decoders use the live `egomimic.rldb.zarr` tokenizers only; no checkpoint-era codec copies are carried. Anything the codec needs for deployment (``detokenize(times=)``, `stream_clocks`, `stream_ends`) is decode-only; never diff --git a/egomimic/robot/graph_policy.py b/egomimic/robot/graph_policy.py index e50a4254b..1630ad4c1 100644 --- a/egomimic/robot/graph_policy.py +++ b/egomimic/robot/graph_policy.py @@ -541,10 +541,14 @@ def load_graph_policy(config): inference_controls = () if inference_profiles is not None: inference_controls = configure_profile_controls( - graph, training, inference_profiles, decoder=adapter.decoder + graph, + training, + inference_profiles, + decoder=adapter.decoder, + station_controls=config.get("station_controls"), ) graph.nets.eval() - return GraphRobotPolicy( + policy = GraphRobotPolicy( graph, context.normalizer, adapter, @@ -552,3 +556,13 @@ def load_graph_policy(config): inference_graph=inference_graph, inference_controls=inference_controls, ) + # Station start values for exposed controls (e.g. sampling steps at the + # trained 50 where a model declaration opens at 10); unexposed names are skipped. + defaults = { + name: value + for name, value in (config.get("inference_defaults") or {}).items() + if name in policy.inference_controls() + } + if defaults: + policy.apply_inference_overrides(defaults) + return policy diff --git a/egomimic/robot/stream_decoder.py b/egomimic/robot/stream_decoder.py index 6a8bb20a8..ddc738e96 100644 --- a/egomimic/robot/stream_decoder.py +++ b/egomimic/robot/stream_decoder.py @@ -63,9 +63,8 @@ def execution_steps(self): class ReplayTempo: - """speed / hold_speed state and the time-chunk retimer's clock warp. - - ARC decoders do not inherit this optional baseline tempo control. + """speed / hold_speed state and the clock warp shared by the time-chunk + retimer and the first-stream ARC decoder (rollout-only, never trained). Subclasses set dt, action_horizon, hold_threshold and ``holds`` (whether hold_speed can act on their output).""" @@ -112,15 +111,16 @@ def hold_speed_percent(self): def hold_speed_percent(self, value): self.set_speed(self.speed, value / 100) - def _warp(self, clocks, cums, end=np.inf): + def _warp(self, clocks, cums, end=np.inf, start=0): """Token-clock time to read each of the H control ticks at, and the share of the decoded span that was a hold. rate = d(clock) / d(wall). ``end`` is the clock time at which the token's path runs out: nothing after it is a hold, - it is just the endpoint being repeated, so it is left out of the share.""" - wall = self.dt * np.arange(self.action_horizon, dtype=np.float64) + it is just the endpoint being repeated, so it is left out of the share. + ``start`` is the first tick (1 for next-frame ARC rows).""" + wall = self.dt * np.arange(start, start + self.action_horizon, dtype=np.float64) top = max(self.speed, self.hold_speed) n = ( - int(np.ceil(self.action_horizon * top * _WARP_OVERSAMPLE)) + int(np.ceil((start + self.action_horizon) * top * _WARP_OVERSAMPLE)) + _WARP_OVERSAMPLE + 1 ) @@ -291,25 +291,62 @@ def _decode(self, row, initial_action): ) -class FirstStreamArcDecoder(BimanualIntervalArcDecoder, ChunkTermination): - """Keep a waypoint prefix; replan when the first moving stream ends.""" +class FirstStreamArcDecoder(BimanualIntervalArcDecoder, ChunkTermination, ReplayTempo): + """Keep a waypoint prefix; replan when the first moving stream ends. + + speed / hold_speed replay the token's translation clock faster or slower + (moving vs held, as TimeChunkRetimer); 1.0 / 1.0 is the plain decode.""" + + holds = True - def __init__(self, *, execute_percent=30, first_stream=1, **kwargs): + def __init__( + self, + *, + execute_percent=30, + first_stream=1, + speed=1.0, + hold_speed=1.0, + hold_threshold=0.05, + **kwargs, + ): super().__init__(**kwargs) self.execute_percent, self.first_stream = execute_percent, first_stream + self._init_tempo(speed, hold_speed, hold_threshold) def _decode(self, row, initial_action): token = ( row[: self.M * self.execute_percent // 100] if self.first_stream else row ) - steps = ( - self.codec.control_steps( - token, initial_action=initial_action, max_steps=self.action_horizon + h, times, hold_fraction = self.action_horizon, None, None + if self.speed != 1.0 or self.hold_speed != 1.0: + clocks, cums = [], [] + for clock, cum in self.codec.stream_clocks( + token, initial_action=initial_action + ): + finite = np.isfinite(clock) # a stalled velocity stream ends there + clocks.append(clock[finite]) + cums.append(cum[finite]) + end = max(clock[-1] for clock in clocks) + times, hold_fraction = self._warp(clocks, cums, end, start=1) + if not self.first_stream: + steps = h + elif times is None: + steps = self.codec.control_steps( + token, initial_action=initial_action, max_steps=h + ) + else: + end = self.codec.control_duration(token, initial_action=initial_action) + steps = ( + h + if end is None + else min(h, int(np.searchsorted(times, end - 1e-9)) + 1) ) - if self.first_stream - else self.action_horizon + decoded = self.codec.detokenize( + token, + steps, + initial_action=initial_action, + times=None if times is None else times[:steps], ) - decoded = self.codec.detokenize(token, steps, initial_action=initial_action) if steps < self.action_horizon: decoded = np.concatenate( (decoded, np.repeat(decoded[-1:], self.action_horizon - steps, axis=0)) @@ -318,5 +355,8 @@ def _decode(self, row, initial_action): "replan_steps": steps if self.first_stream else None, # 0: replan_every "horizon": self.action_horizon, "execute_waypoints": len(token), + "speed": self.speed, + "hold_speed": self.hold_speed, + "hold_fraction": hold_fraction, } return decoded