feat(execution): add continuous mode for interval selection with keypoint marking.

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.

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."""

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."""

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

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):

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."""

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

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,

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,

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,

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.