From 18edd71c39707044c12aa7167079d41d78f9a212 Mon Sep 17 00:00:00 2001 From: Richie6667 <1052996586@qq.com> Date: Tue, 15 Sep 2026 12:49:16 +0800 Subject: [PATCH] feat(execution): add continuous mode for interval selection with keypoint marking --- aao_configs/pick_and_place.yaml | 2 + auto_atom/basis/mjc/gs_mujoco_env.py | 4 +- auto_atom/basis/mjc/mujoco_env.py | 33 ++- auto_atom/config/execution.py | 9 + auto_atom/config/task.py | 11 + auto_atom/contracts.py | 18 ++ auto_atom/execution_model.py | 2 + auto_atom/mock.py | 16 +- auto_atom/runtime.py | 69 ++++- auto_atom/stage_execution.py | 1 + docs/getting-started/cli_reference.md | 13 + .../stages_and_waypoints.md | 54 ++++ docs/task-configuration/task_file_schema.md | 3 + docs/tools/external_data_collection.md | 52 ++++ tests/test_task_runner_interval_selection.py | 280 ++++++++++++++++++ 15 files changed, 563 insertions(+), 4 deletions(-) diff --git a/aao_configs/pick_and_place.yaml b/aao_configs/pick_and_place.yaml index dc28c48a..5bc39f52 100644 --- a/aao_configs/pick_and_place.yaml +++ b/aao_configs/pick_and_place.yaml @@ -23,6 +23,8 @@ env: # start: { stage: pick_source, phase: post_move, waypoint: 0, side: before } # stop: { stage: place_source, phase: post_move, waypoint: 0, side: after } # max_fast_forward_updates: 10000 +# continuous: false # true: no-skip dense collection with per-step keypoint +# # marks; requires update_boundary: control_tick task: # Pose randomization applied at each reset. diff --git a/auto_atom/basis/mjc/gs_mujoco_env.py b/auto_atom/basis/mjc/gs_mujoco_env.py index 7823d233..0b9fe25d 100644 --- a/auto_atom/basis/mjc/gs_mujoco_env.py +++ b/auto_atom/basis/mjc/gs_mujoco_env.py @@ -2479,7 +2479,9 @@ def capture_observation(self) -> dict[str, dict[str, Any]]: # The canonical adapter stacks replicated observations or broadcasts a # single shared-physics observation exactly once, before noise. obs = self._capture_observation_raw() - return self._apply_camera_noise(obs) + obs = self._apply_camera_noise(obs) + self._merge_keypoint_mark(obs) + return obs def _capture_observation_raw(self) -> dict[str, dict[str, Any]]: """Capture the logical batch and inject GS streams before noise.""" diff --git a/auto_atom/basis/mjc/mujoco_env.py b/auto_atom/basis/mjc/mujoco_env.py index c73f0081..4b767e36 100644 --- a/auto_atom/basis/mjc/mujoco_env.py +++ b/auto_atom/basis/mjc/mujoco_env.py @@ -2103,6 +2103,7 @@ def __init__(self, config: Optional[EnvConfig] = None, **kwargs): if config.name: ComponentRegistry.register_env(config.name, self) self._key_creator = KeyCreator(self.config.structured) + self._keypoint_mark: Optional[list[dict[str, Any]]] = None self._camera_noise_processor = CameraNoiseProcessor( {spec.name: spec for spec in self.config.cameras} ) @@ -2666,13 +2667,43 @@ def capture_observation(self) -> dict[str, dict[str, Any]]: observation = self._capture_observation_raw() processor = getattr(self, "_camera_noise_processor", None) if processor is None: + self._merge_keypoint_mark(observation) return observation - return processor.process_batched_observation( + observation = processor.process_batched_observation( observation, self._key_creator, structured=self.config.structured, batch_size=self.batch_size, ) + self._merge_keypoint_mark(observation) + return observation + + def set_keypoint_mark( + self, + mark: Optional[list[dict[str, Any]]], + ) -> None: + """Store per-step keypoint marks published by TaskRunner. + + The stored rows are merged into :meth:`capture_observation` output + under the ``task/keypoint`` key (fully prefixed, e.g. + ``/robot/task/keypoint`` in structured mode). + """ + self._keypoint_mark = None if mark is None else [dict(row) for row in mark] + + def _merge_keypoint_mark(self, observation: dict[str, dict[str, Any]]) -> None: + mark = self._keypoint_mark + if mark is None: + return + times = [ + int(env.data.time * 1e9) if self.config.stamp_ns else float(env.data.time) + for env in self.envs + ] + if len(times) == 1 and self.batch_size > 1: + times = times * self.batch_size + observation[self._key_creator.apply_prefix("task/keypoint")] = { + "data": list(mark), + "t": times, + } def _capture_observation_raw(self) -> dict[str, dict[str, Any]]: """Capture a logical batch before RGB/depth sensor noise is applied.""" diff --git a/auto_atom/config/execution.py b/auto_atom/config/execution.py index 27f7c424..48fd94c7 100644 --- a/auto_atom/config/execution.py +++ b/auto_atom/config/execution.py @@ -115,6 +115,15 @@ class IntervalSelectionConfig(BaseModel, frozen=True): max_fast_forward_updates: PositiveInt = 10_000 """Maximum controller updates per environment while ``reset()`` advances to ``start``.""" + continuous: bool = False + """Non-transition collection mode. When true, ``reset()`` still + fast-forwards to the ``start`` boundary, but every public update returns + at ``control_tick`` granularity without skipping segments between + keypoints, and each ``TaskUpdate`` marks whether the current state is a + configured keypoint boundary (identity plus ``before``/``after`` side). + The marks let a dense collection pass be filtered into the same + boundary-only transition data afterwards. Requires + ``execution.update_boundary: control_tick``.""" @model_validator(mode="before") @classmethod diff --git a/auto_atom/config/task.py b/auto_atom/config/task.py index 2be9cb34..a9e527c2 100644 --- a/auto_atom/config/task.py +++ b/auto_atom/config/task.py @@ -21,6 +21,7 @@ KeypointSide, TaskKeypointConfig, TaskPhase, + UpdateBoundary, ) from auto_atom.config.motion import StageConfig from auto_atom.config.operations import Operation @@ -335,6 +336,16 @@ def _validate_interval_selection(self) -> "TaskFileConfig": selection = self.execution.interval_selection if selection is None: return self + if ( + selection.continuous + and self.execution.update_boundary != UpdateBoundary.CONTROL_TICK + ): + raise ValueError( + "execution.interval_selection.continuous requires " + "execution.update_boundary=control_tick; continuous mode " + "already steps through every keypoint boundary per tick, " + "while a coarser update_boundary skips them by design" + ) stages_by_name: Dict[str, List[Tuple[int, StageConfig]]] = {} for index, stage in enumerate(self.task.stages): diff --git a/auto_atom/contracts.py b/auto_atom/contracts.py index 3d4c4afa..b8777ee0 100644 --- a/auto_atom/contracts.py +++ b/auto_atom/contracts.py @@ -31,6 +31,7 @@ List, Optional, Protocol, + Sequence, TypeVar, cast, runtime_checkable, @@ -224,6 +225,23 @@ class ObservationEnvProtocol(EnvProtocol, Protocol): def capture_observation(self) -> Dict[str, Dict[str, Any]]: ... +@runtime_checkable +class KeypointMarkEnvProtocol(ObservationEnvProtocol, Protocol): + """Environment capability for publishing per-step keypoint marks. + + ``set_keypoint_mark`` receives one row per environment, where each row + carries ``is_keypoint`` and the marked keypoint's identity + (``stage_index`` / ``stage_name`` / ``phase`` / ``waypoint`` / ``side``). + Implementations merge the stored rows into :meth:`capture_observation` + output under the ``task/keypoint`` key; passing ``None`` clears them. + """ + + def set_keypoint_mark( + self, + mark: Optional[Sequence[Mapping[str, Any]]], + ) -> None: ... + + @runtime_checkable class JointActionEnvProtocol(EnvProtocol, Protocol): """Environment capability for directly applying operator joint actions.""" diff --git a/auto_atom/execution_model.py b/auto_atom/execution_model.py index 53a6949b..0a4e3ff8 100644 --- a/auto_atom/execution_model.py +++ b/auto_atom/execution_model.py @@ -17,6 +17,7 @@ from auto_atom.config.execution import ( IntervalSelectionConfig, KeypointSelector, + KeypointSide, TaskKeypointConfig, TaskPhase, UpdateBoundary, @@ -154,6 +155,7 @@ class _EnvRuntimeState: latest_status: StageExecutionStatus = StageExecutionStatus.PENDING latest_details: Dict[str, Any] = field(default_factory=dict) reported_keypoint: Optional[_ResolvedTaskKeypoint] = None + keypoint_mark_side: Optional[KeypointSide] = None @dataclass diff --git a/auto_atom/mock.py b/auto_atom/mock.py index 5e001128..2367ca0b 100644 --- a/auto_atom/mock.py +++ b/auto_atom/mock.py @@ -153,12 +153,26 @@ class MockEnv: """Minimal env stub that satisfies the ``SceneBackend.env`` contract.""" batch_size: int = 1 + keypoint_mark: Optional[List[Dict[str, Any]]] = None def step(self, action: np.ndarray, env_mask: np.ndarray | None = None) -> None: pass + def set_keypoint_mark( + self, + mark: Optional[List[Dict[str, Any]]], + ) -> None: + """Store per-step keypoint marks published by TaskRunner.""" + self.keypoint_mark = None if mark is None else list(mark) + def capture_observation(self) -> Dict[str, Dict[str, Any]]: - return {} + observation: Dict[str, Dict[str, Any]] = {} + if self.keypoint_mark is not None: + observation["task/keypoint"] = { + "data": list(self.keypoint_mark), + "t": [0.0 for _ in self.keypoint_mark], + } + return observation def apply_joint_action( self, diff --git a/auto_atom/runtime.py b/auto_atom/runtime.py index cff73093..91dd61f8 100644 --- a/auto_atom/runtime.py +++ b/auto_atom/runtime.py @@ -678,6 +678,7 @@ def _update_impl( pending[env_index] = False continue state.reported_keypoint = None + state.keypoint_mark_side = None while bool(np.any(pending)): for env_index_value in np.flatnonzero(pending): @@ -731,6 +732,16 @@ def _update_impl( pending[env_index] = False continue + if ( + selection is not None + and selection.continuous + and completed_keypoint is not None + ): + # Continuous (non-transition) collection marks every tick + # whose post-tick state is a configured keypoint boundary. + state.reported_keypoint = completed_keypoint + state.keypoint_mark_side = KeypointSide.AFTER + boundary_event = self._require_timeline().reached_update_boundary(event) if boundary_event is None: if state.done: @@ -812,6 +823,7 @@ def _fail_internal_update_limit( state.phase = None state.phase_step = None state.reported_keypoint = None + state.keypoint_mark_side = None def _fast_forward_to_interval_start( self, @@ -900,6 +912,7 @@ def _fast_forward_to_interval_start( else "interval_fast_forward_failed", "start": selection.start.model_dump(mode="json"), "stop": selection.stop.model_dump(mode="json"), + "continuous": bool(selection.continuous), "fast_forward_updates": int(ticks[index]), "max_fast_forward_updates": max_updates, } @@ -933,6 +946,7 @@ def _fast_forward_to_interval_start( "interval_selection": interval_details, } state.reported_keypoint = start_keypoint + state.keypoint_mark_side = selection.start.side state.phase = start_keypoint.phase.value state.phase_step = start_keypoint.waypoint if start_state_index == stop_state_index: @@ -981,6 +995,7 @@ def _fail_interval_fast_forward( state.phase = None state.phase_step = None state.reported_keypoint = None + state.keypoint_mark_side = None @staticmethod def _finish_interval( @@ -1013,6 +1028,7 @@ def _finish_interval( state.phase = keypoint.phase.value state.phase_step = keypoint.waypoint state.reported_keypoint = keypoint + state.keypoint_mark_side = selection.stop.side def get_env(self) -> EnvProtocol: """Return the underlying environment object managed by this runner.""" @@ -2465,6 +2481,7 @@ def _build_task_update(self) -> TaskUpdate: interval_details.setdefault( "stop", selection.stop.model_dump(mode="json") ) + interval_details.setdefault("continuous", bool(selection.continuous)) interval_details.setdefault( "max_fast_forward_updates", int(selection.max_fast_forward_updates), @@ -2502,7 +2519,7 @@ def _build_task_update(self) -> TaskUpdate: else: phase.append(state.phase) phase_step.append(-1 if state.phase_step is None else state.phase_step) - return TaskUpdate( + update = TaskUpdate( stage_index=np.asarray(stage_index, dtype=np.int64), stage_name=stage_name, status=np.asarray(status, dtype=object), @@ -2512,6 +2529,56 @@ def _build_task_update(self) -> TaskUpdate: phase=phase, phase_step=np.asarray(phase_step, dtype=np.int64), ) + self._publish_keypoint_marks() + return update + + def _publish_keypoint_marks(self) -> None: + """Push the current per-step keypoint marks into the environment. + + Environments exposing ``set_keypoint_mark`` merge the published rows + into ``capture_observation()`` under ``task/keypoint``, so + observation-only collection pipelines record the marks without + threading anything extra through a custom adapter. Envs without the + capability are skipped; ``None`` clears any stored rows. + """ + env = self.get_env() + setter = getattr(env, "set_keypoint_mark", None) + if setter is None: + return + selection = ( + self._timeline.interval_selection if self._timeline is not None else None + ) + continuous = selection is not None and bool(selection.continuous) + if not continuous: + setter(None) + return + rows = [] + for state in self._env_states: + keypoint = state.reported_keypoint + side = state.keypoint_mark_side + if keypoint is not None and side is not None: + rows.append( + { + "is_keypoint": True, + "stage_index": int(keypoint.stage_index), + "stage_name": keypoint.stage_name, + "phase": keypoint.phase.value, + "waypoint": int(keypoint.waypoint), + "side": side.value, + } + ) + else: + rows.append( + { + "is_keypoint": False, + "stage_index": -1, + "stage_name": "", + "phase": None, + "waypoint": -1, + "side": None, + } + ) + setter(rows) def _collect_reset_details( self, diff --git a/auto_atom/stage_execution.py b/auto_atom/stage_execution.py index 127d2bb4..bf8d62f3 100644 --- a/auto_atom/stage_execution.py +++ b/auto_atom/stage_execution.py @@ -1358,6 +1358,7 @@ def _set_failed(state: _EnvRuntimeState, details: Dict[str, Any]) -> None: state.phase = None state.phase_step = None state.reported_keypoint = None + state.keypoint_mark_side = None def _set_succeeded( self, diff --git a/docs/getting-started/cli_reference.md b/docs/getting-started/cli_reference.md index 0d8a1e39..8e1335c8 100644 --- a/docs/getting-started/cli_reference.md +++ b/docs/getting-started/cli_reference.md @@ -141,6 +141,7 @@ To discover which configs are runnable tasks, use [`aao-info`](#aao-info). | `[+]execution.interval_selection...` | mapping | unset | Run between states immediately before or after configured `stage` / `phase` / `waypoint` keypoints | | `[+]execution.interval_selection.{start,stop}.side=...` | enum | `before` / `after` | Endpoint side relative to its keypoint; the start default is `before`, while the stop default is `after` | | `[+]execution.interval_selection.max_fast_forward_updates=N` | int | 10000 | Per-environment controller-update limit while `reset()` advances to the interval start boundary | +| `[+]execution.interval_selection.continuous=true` | bool | false | Non-transition collection: still fast-forward to `start`, then return one control tick per public update without skipping segments, marking each step in `capture_observation()` under `task/keypoint` | | `[+]execution.keypoint_selection=[...]` | list | unset | Ordered keypoints to execute: task-wide ordinals or `{stage, phase?, waypoint?}` entries, with negative indexes counting from the end; unlisted keypoints are skipped. Mutually exclusive with `interval_selection` | Any key present in the YAML config can be overridden on the command line following Hydra syntax: @@ -225,6 +226,18 @@ update and reset fast-forward limits are independent; both default to `10000`. With `execution.render_internal_updates=false`, all of those internal updates still run, but their viewer refreshes and `step_delay` calls are coalesced into one delay-free refresh at the public boundary. + +To collect the interval densely instead of stepping keypoint by keypoint, add +`+execution.interval_selection.continuous=true` (requires the default +`update_boundary=control_tick`). Every public update then advances one control +tick without skipping the segments between keypoints, and each +`capture_observation()` result carries a `task/keypoint` entry (prefixed, +e.g. `/robot/task/keypoint`, in structured mode) whose per-environment rows +mark `is_keypoint` plus the marked keypoint's stage, phase, waypoint, and +`before`/`after` side. Filtering the dense pass to marked samples reproduces +the boundary-only data, so no separate keypoint-boundary collection run is +needed. See +[Stages & Waypoints](../task-configuration/stages_and_waypoints.md#continuous-keypoint-marked-collection). See [Stages & Waypoints](../task-configuration/stages_and_waypoints.md#task-interval-boundary-selection) for endpoint semantics and reporting. diff --git a/docs/task-configuration/stages_and_waypoints.md b/docs/task-configuration/stages_and_waypoints.md index 3dbe091c..e4d17597 100644 --- a/docs/task-configuration/stages_and_waypoints.md +++ b/docs/task-configuration/stages_and_waypoints.md @@ -45,6 +45,7 @@ execution: waypoint: 0 side: after max_fast_forward_updates: 10000 + continuous: false task: stages: @@ -60,6 +61,12 @@ Each endpoint contains: | `waypoint` | Zero-based index in that phase; `eef` is a singleton and only accepts `0` | | `side` | `before` or `after` the referenced keypoint; defaults to `before` for `start` and `after` for `stop` | +The selection also accepts one mode switch: + +| Field | Meaning | +| --- | --- | +| `continuous` | `false` (default) keeps the boundary-oriented behavior described below; `true` enables non-transition collection described in [Continuous keypoint-marked collection](#continuous-keypoint-marked-collection) | + Both endpoints use the same `TaskKeypointConfig` schema. Its standalone `side` default is `None`; `IntervalSelectionConfig` resolves that adaptive value by endpoint role and exposes the concrete `before` / `after` value when @@ -266,6 +273,53 @@ condition unless the operator already holds the target object. entries and the current selection event. Omitting `execution.keypoint_selection` preserves full-task execution. +### Continuous keypoint-marked collection + +`execution.interval_selection.continuous: true` turns the selected interval +into a non-transition (no-skip) collection mode: `reset()` still +fast-forwards to the `start` boundary exactly as before, but every public +`update()` afterwards returns at `control_tick` granularity and executes +every segment between the interval's keypoints. Nothing is skipped; the +runner simply stamps a keypoint mark on each step whose state is a +configured keypoint boundary, so one dense collection pass can later be +filtered into the same boundary-only transition data — no second collection +pass is needed. + +Because continuous mode is inherently dense, it requires +`execution.update_boundary: control_tick`; combining it with `primitive`, +`keypoint`, or `stage` is rejected during config validation. + +The marks are published into the environment, so observation-only collection +pipelines record them without any extra plumbing. After each `reset()` / +`update()`, `TaskRunner` pushes one row per environment into the env via the +optional `set_keypoint_mark` capability, and the env's +`capture_observation()` output carries them under the `task/keypoint` key — +fully prefixed, e.g. `/robot/task/keypoint`, in structured mode. The entry +follows the normal observation shape: `"data"` is the list of per-environment +rows and `"t"` holds one timestamp per environment. Each row contains: + +| Field | Meaning | +| --- | --- | +| `is_keypoint` | `bool`; `true` when that environment's post-step state is a configured keypoint boundary | +| `stage_index` / `stage_name` | The marked keypoint's stage identity; `-1` / `""` when unmarked | +| `phase` / `waypoint` | The marked keypoint's phase and zero-based waypoint index; `None` / `-1` when unmarked | +| `side` | `before` when the state precedes the keypoint (the reset state, or a `stop.side: before` terminal state) and `after` when the keypoint has just completed | + +The reset step marks the `start` boundary with `start.side`; each step whose +tick completes a keypoint marks it with `after`; the terminal step marks the +`stop` boundary with `stop.side`. In-between steps report +`is_keypoint: false`. `TaskUpdate.details[env_index]["interval_selection"]` +additionally reports `continuous`. Environments without the +`set_keypoint_mark` capability are skipped, and `None` clears any stored +rows, so non-continuous runs never emit the key. + +> [!NOTE] +> The state after completing keypoint *K* is the same physical state as the +> state before the next keypoint. The marks use the completed keypoint's +> `after` identity for that state; the only explicit `before` marks are the +> interval's `start` and `stop` boundaries, where the interval semantics +> define them. + ## Stage reference site By default, waypoints with `reference: object_world` or `reference: object` diff --git a/docs/task-configuration/task_file_schema.md b/docs/task-configuration/task_file_schema.md index 9d030daa..c30b52e6 100644 --- a/docs/task-configuration/task_file_schema.md +++ b/docs/task-configuration/task_file_schema.md @@ -331,6 +331,9 @@ execution: `update_boundary` can be `control_tick`, `primitive`, `keypoint`, or `stage`. `interval_selection` can restrict execution to the interval between two +keypoint boundaries; its `continuous: true` mode additionally collects the +interval without skipping segments between keypoints and marks each +configured keypoint boundary per step. See keypoint boundaries, while `keypoint_selection` can run an ordered subset of keypoints and skip the rest. Both accept only `TaskRunner` / `aao-demo` and are mutually exclusive. See diff --git a/docs/tools/external_data_collection.md b/docs/tools/external_data_collection.md index 6d1e02b0..0eeaf517 100644 --- a/docs/tools/external_data_collection.md +++ b/docs/tools/external_data_collection.md @@ -150,6 +150,58 @@ different: each held-object waypoint is written once. Use the default `control_tick` for dense physical trajectories; use `primitive`, `keypoint`, or `stage` only when boundary-only samples are intentional. +### Continuous keypoint-marked collection + +Collecting boundary-only transition data normally requires a dedicated pass +with `execution.update_boundary=keypoint`. To obtain the same keypoint data +from one dense pass instead, add +`+execution.interval_selection.continuous=true` (and keep the default +`update_boundary=control_tick`): + +```python +overrides=[ + "+execution.interval_selection.start.stage=pick_source", + "+execution.interval_selection.start.phase=post_move", + "+execution.interval_selection.start.waypoint=0", + "+execution.interval_selection.start.side=after", + "+execution.interval_selection.stop.stage=place_source", + "+execution.interval_selection.stop.phase=post_move", + "+execution.interval_selection.stop.waypoint=0", + "+execution.interval_selection.stop.side=after", + "+execution.interval_selection.continuous=true", +] +``` + +`runner.reset()` still fast-forwards to the start boundary, then every +public `runner.update()` advances exactly one controller update without +skipping the segments between keypoints. The marks are injected into the +observation itself, so pipelines that only persist `capture_observation()` +output record them automatically. After each `reset()` / `update()`, the +runner publishes one row per environment to the env, and +`capture_observation()` returns it under the `task/keypoint` key +(`/robot/task/keypoint` in structured mode): + +```python +runner.reset() +obs = env.capture_observation() +mark = obs["task/keypoint"] # {"data": [per-env rows], "t": [...]} +row = mark["data"][env_index] # is_keypoint, stage_index, + # stage_name, phase, waypoint, side +``` + +Each row carries `is_keypoint` plus the marked keypoint's +`stage_index` / `stage_name` / `phase` / `waypoint` / `side`. The reset step +marks the `start` boundary with `start.side`; a step that completes a +keypoint marks it with `after`; the terminal step marks the `stop` boundary +with `stop.side`. Environments without the `set_keypoint_mark` capability +and non-continuous runs omit the key entirely. A writer that persists +observations verbatim therefore keeps the marks automatically, and filtering +the dense trajectory to samples whose row has `is_keypoint: true` +reproduces the boundary-only transition data for the selected interval, +without a second collection pass. See +[Stages & Waypoints](../task-configuration/stages_and_waypoints.md#continuous-keypoint-marked-collection) +for the exact field semantics. + `execution.render_internal_updates: false` only coalesces passive-viewer refreshes and skips viewer `step_delay`; it does not change which observations the host captures or how many physics ticks run. diff --git a/tests/test_task_runner_interval_selection.py b/tests/test_task_runner_interval_selection.py index a8dc016f..7c75722a 100644 --- a/tests/test_task_runner_interval_selection.py +++ b/tests/test_task_runner_interval_selection.py @@ -101,6 +101,11 @@ def _operator_positions(runner: TaskRunner) -> np.ndarray: ) +def _obs_mark(runner: TaskRunner) -> dict: + """Return the env-0 keypoint mark row injected into the observation.""" + return runner.get_env().capture_observation()["task/keypoint"]["data"][0] + + @pytest.fixture(autouse=True) def _clear_component_registry(): ComponentRegistry.clear() @@ -199,6 +204,7 @@ def test_interval_endpoint_sides_have_role_specific_defaults() -> None: "side": "after", }, "max_fast_forward_updates": 10_000, + "continuous": False, } @@ -988,3 +994,277 @@ def test_primitive_action_preserves_existing_positional_constructor_order() -> N assert action.arc_cumulative_angle == pytest.approx(0.25) assert action.phase is None assert action.waypoint is None + + +def test_continuous_mode_marks_keypoint_boundaries_without_skipping_ticks() -> None: + interval = { + "start": _keypoint("selected", "pre_move", 0), + "stop": _keypoint("selected", "post_move", 1), + "continuous": True, + } + runner = TaskRunner().from_config( + TaskFileConfig.model_validate(_task_payload(interval=interval)) + ) + try: + reset_update = runner.reset() + + assert reset_update.done.tolist() == [False] + assert _obs_mark(runner) == { + "is_keypoint": True, + "stage_index": 0, + "stage_name": "selected", + "phase": "pre_move", + "waypoint": 0, + "side": "before", + } + assert reset_update.details[0]["interval_selection"]["continuous"] is True + assert ( + reset_update.details[0]["interval_selection"]["fast_forward_updates"] == 0 + ) + + # Each mock pose primitive takes two controller ticks. The four + # in-interval keypoints therefore need eight dense updates, and only + # the completion tick of each keypoint is marked. + expected_marks = [ + ("pre_move", 0, "after"), + ("pre_move", 1, "after"), + ("post_move", 0, "after"), + ("post_move", 1, "after"), + ] + positions: list[float] = [] + marks: list[tuple[str, int, str] | None] = [] + for _ in range(16): + update = runner.update() + positions.append(float(_operator_positions(runner)[0, 0])) + row = _obs_mark(runner) + if bool(row["is_keypoint"]): + marks.append((row["phase"], int(row["waypoint"]), row["side"])) + else: + marks.append(None) + if bool(update.done[0]): + break + + assert positions == pytest.approx([0.2, 0.1, 0.1, 0.2, 0.2, 0.3, 0.3, 0.4]) + assert [mark for mark in marks if mark is not None] == expected_marks + assert marks[1] == ("pre_move", 0, "after") + assert marks[3] == ("pre_move", 1, "after") + assert marks[5] == ("post_move", 0, "after") + assert marks[7] == ("post_move", 1, "after") + assert update.success.tolist() == [True] + finally: + runner.close() + + +def test_continuous_run_ticks_match_plain_control_tick_interval() -> None: + interval = { + "start": _keypoint("selected", "pre_move", 0), + "stop": _keypoint("selected", "post_move", 1), + } + positions: dict[bool, list[float]] = {} + for continuous in (False, True): + payload_interval = dict(interval) + if continuous: + payload_interval["continuous"] = True + runner = TaskRunner().from_config( + TaskFileConfig.model_validate(_task_payload(interval=payload_interval)) + ) + try: + runner.reset() + sequence: list[float] = [] + for _ in range(64): + update = runner.update() + sequence.append(float(_operator_positions(runner)[0, 0])) + if bool(update.done[0]): + break + positions[continuous] = sequence + finally: + runner.close() + + assert len(positions[True]) == 8 + assert positions[True] == positions[False] + + +def test_continuous_reset_after_start_marks_after_side() -> None: + interval = { + "start": _keypoint("selected", "post_move", 0, side="after"), + "stop": _keypoint("selected", "post_move", 1), + "continuous": True, + } + runner = TaskRunner().from_config( + TaskFileConfig.model_validate(_task_payload(interval=interval)) + ) + try: + reset_update = runner.reset() + + assert _operator_positions(runner)[0, 0] == pytest.approx(0.3) + assert _obs_mark(runner) == { + "is_keypoint": True, + "stage_index": 0, + "stage_name": "selected", + "phase": "post_move", + "waypoint": 0, + "side": "after", + } + assert ( + reset_update.details[0]["interval_selection"]["fast_forward_updates"] == 6 + ) + + running_update = runner.update() + assert _obs_mark(runner)["is_keypoint"] is False + + final_update = runner.update() + assert final_update.done.tolist() == [True] + assert final_update.success.tolist() == [True] + assert _obs_mark(runner) == { + "is_keypoint": True, + "stage_index": 0, + "stage_name": "selected", + "phase": "post_move", + "waypoint": 1, + "side": "after", + } + finally: + runner.close() + + +def test_continuous_stop_before_marks_the_before_side() -> None: + interval = { + "start": _keypoint("selected", "pre_move", 0), + "stop": _keypoint("selected", "pre_move", 1, side="before"), + "continuous": True, + } + runner = TaskRunner().from_config( + TaskFileConfig.model_validate(_task_payload(interval=interval)) + ) + try: + runner.reset() + assert _obs_mark(runner)["side"] == "before" + + runner.update() + assert _obs_mark(runner)["is_keypoint"] is False + + final_update = runner.update() + assert final_update.done.tolist() == [True] + assert final_update.success.tolist() == [True] + assert _obs_mark(runner) == { + "is_keypoint": True, + "stage_index": 0, + "stage_name": "selected", + "phase": "pre_move", + "waypoint": 1, + "side": "before", + } + assert _operator_positions(runner)[0, 0] == pytest.approx(0.1) + finally: + runner.close() + + +def test_continuous_partial_reset_mask_publishes_only_selected_envs() -> None: + interval = { + "start": _keypoint("selected", "pre_move", 0), + "stop": _keypoint("selected", "post_move", 1), + "continuous": True, + } + runner = TaskRunner().from_config( + TaskFileConfig.model_validate(_task_payload(batch_size=2, interval=interval)) + ) + try: + runner.reset(np.asarray([True, False], dtype=bool)) + rows = runner.get_env().capture_observation()["task/keypoint"]["data"] + + assert rows[0] == { + "is_keypoint": True, + "stage_index": 0, + "stage_name": "selected", + "phase": "pre_move", + "waypoint": 0, + "side": "before", + } + assert rows[1] == { + "is_keypoint": False, + "stage_index": -1, + "stage_name": "", + "phase": None, + "waypoint": -1, + "side": None, + } + finally: + runner.close() + + +def test_continuous_mode_requires_control_tick_update_boundary() -> None: + payload = _task_payload( + interval={ + "start": _keypoint("selected", "pre_move", 0), + "stop": _keypoint("selected", "post_move", 1), + "continuous": True, + } + ) + payload["execution"]["update_boundary"] = "keypoint" + + with pytest.raises( + ValidationError, + match="requires execution.update_boundary=control_tick", + ): + TaskFileConfig.model_validate(payload) + + +def test_continuous_mode_publishes_marks_into_observation() -> None: + interval = { + "start": _keypoint("selected", "pre_move", 0), + "stop": _keypoint("selected", "post_move", 1), + "continuous": True, + } + runner = TaskRunner().from_config( + TaskFileConfig.model_validate(_task_payload(interval=interval)) + ) + try: + runner.reset() + obs = runner.get_env().capture_observation() + + assert obs["task/keypoint"]["data"] == [ + { + "is_keypoint": True, + "stage_index": 0, + "stage_name": "selected", + "phase": "pre_move", + "waypoint": 0, + "side": "before", + } + ] + assert obs["task/keypoint"]["t"] == [0.0] + + runner.update() + obs = runner.get_env().capture_observation() + assert obs["task/keypoint"]["data"][0]["is_keypoint"] is False + assert obs["task/keypoint"]["data"][0]["side"] is None + + runner.update() + obs = runner.get_env().capture_observation() + row = obs["task/keypoint"]["data"][0] + assert row["is_keypoint"] is True + assert row["phase"] == "pre_move" + assert row["waypoint"] == 0 + assert row["side"] == "after" + finally: + runner.close() + + +def test_observation_has_no_keypoint_mark_without_continuous_mode() -> None: + runner = TaskRunner().from_config( + TaskFileConfig.model_validate( + _task_payload( + interval={ + "start": _keypoint("selected", "pre_move", 0), + "stop": _keypoint("selected", "post_move", 1), + } + ) + ) + ) + try: + runner.reset() + obs = runner.get_env().capture_observation() + + assert "task/keypoint" not in obs + finally: + runner.close()