feat: finish finite policy runs and support replay

docs/component-reference.md

@@ -129,3 +129,18 @@ Omitted/null `horizontalFov` preserves the aspect-derived horizontal view.
 For example, `fov: 58`, `horizontalFov: 87`, `fullWidth: 48`, `fullHeight: 27`
 calibrates the two axes independently, before applying `cropX`/`cropY` and
 `width`/`height`. Use a matching browser runtime for editor and published sites.
+
+
+## 有限时长演示与重播
+
+`PolicyController(..., end_after_steps=500)` 在提交第 500 个控制步（含物理子步）后结束。
+也可以写 `end_after_steps={"parameter": "motion_id", "steps": {0: 417, 1: 523}}`，
+映射必须覆盖该选项参数的所有值。步数是正安全整数；秒数由项目显式乘控制频率并向上取整。
+每次复位后的第一控制步读取所选步数，暂停和续播不重置预算。配合已有
+`reset_on_change={"parameter": "motion_id", "motions": ...}`，动作切换会复位并开始新一轮。
+
+`open_headless()` 返回值的 `completed` 表示已达到声明的结束条件；`steps(count)`
+到点提前返回，结束后 `step()` 报错，调用 `reset()` 才能重播。原生 `run()` 自动返回。
+编辑器和静态演示显示“已结束”，点击播放自动复位后重播。未声明时继续原有连续运行行为。
+此阶段只支持单个无 Python action modifier 的 `PolicyController`；不推断任务成功、失败或跌倒。
+发布包含该字段的工程时，BotLab 浏览器运行时与 Python SDK 必须一起更新到支持该字段的版本。

src/botlab/_headless.py

@@ -83,6 +83,11 @@ def control_hz(self) -> int:
     def physics_hz(self) -> int:
         return self.session._physics_hz
 
+    @property
+    def completed(self) -> bool:
+        """Whether the declared end condition was reached (cleared by reset)."""
+        return self.session.completed
+
     def set_parameter(self, name, value) -> None:
         """Override a declared parameter from the next control step on."""
         import botlab_host
@@ -130,9 +135,11 @@ def step(self) -> StepRecord:
         return StepRecord(frame, None, None, None)
 
     def steps(self, count, on_step=None) -> list[StepRecord]:
-        """Run up to ``count`` control steps, stopping early if physics asks."""
+        """Run up to ``count`` control steps, stopping on completion or if physics asks."""
         records = []
         for _ in range(count):
+            if self.completed:
+                break
             record = self.step()
             records.append(record)
             if on_step is not None:

src/botlab/_runtime/_host_runtime.py

@@ -84,6 +84,11 @@ def __init__(self, physics, display=None, models=None, registry=None,
         self._in_plan_step = False
         self._pending_plan_reset = None
         self._plan_step_offset = 0
+        self._completed_steps = 0
+        self.completed = False
+        self._active_step_limit = None
+        self._run_step_limit = None
+        self._run_revision = 0
         self._recorder = None
         self._substep_observer = None
         self._custom_step_active = False
@@ -148,6 +153,9 @@ def registerPolicyController(self, plan, modifier=None):
             if any(node.get("data", {}).get("resetOnChange") is not None
                    for entry in controllers for node in entry.get("nodes", [])):
                 raise ValueError("reset_on_change requires a single controller")
+            if any("endAfterSteps" in node.get("data", {})
+                   for entry in controllers for node in entry.get("nodes", [])):
+                raise ValueError("end_after_steps requires a single controller")
             self._register_controller_group(controllers, registry)
             return
 
@@ -168,6 +176,19 @@ def registerPolicyController(self, plan, modifier=None):
             if modifier is not None or len(reset_specs) != 1:
                 raise ValueError("reset_on_change requires one reset rule and no Python modifier")
             modifier = self._motion_reset_modifier(reset_specs[0])
+        end_specs = [node["data"]["endAfterSteps"] for node in nodes
+                     if "endAfterSteps" in node.get("data", {})]
+        self._run_step_limit = None
+        if end_specs:
+            import botlab_host
+
+            from ..sdk._run_end import parse_end_after_steps
+
+            if self._modifier is not None or len(end_specs) != 1:
+                raise ValueError("end_after_steps requires one rule and no Python modifier")
+            self._run_step_limit = parse_end_after_steps(
+                end_specs[0], botlab_host.project_state().manifest()["parameters"],
+            )
         if modifier is not None:
             self._attach_action_hook(modifier)
 
@@ -344,21 +365,19 @@ def _motion_reset_modifier(self, spec):
             motion_ids[value] = motion_id
         if spec.get("robotId") != self._plan_robot_id:
             raise ValueError("reset_on_change must reset its controller's robot")
-        active, previous_index = None, -1
+        active, active_revision = None, -1
 
         def modify(action, frame):
-            nonlocal active, previous_index
+            nonlocal active, active_revision
             selected = botlab_host.parameter(parameter)
             if selected not in motion_ids:
                 raise ValueError("reset_on_change received an unknown option value")
-            if selected != active or (frame.index == 0 and previous_index > 0):
-                active, previous_index = selected, -1
+            if selected != active or active_revision != self._run_revision:
+                active, active_revision = selected, self._run_revision
                 self.requestReset("Parameter changed", {
                     "motionId": motion_ids[selected], "robotId": spec["robotId"],
                     "frame": 0, "rootBody": 0,
                 })
-            else:
-                previous_index = frame.index
             return action
 
         return modify
@@ -616,7 +635,7 @@ def _run_plan(self, max_control_steps):
             )
             self._broadcast_display(frame)
             step_index += 1
-            if self._physics.should_stop():
+            if self.completed or self._physics.should_stop():
                 break
         return step_index
 
@@ -701,6 +720,10 @@ def reset_for_run(self):
 
         self._pending_plan_reset = None
         self._plan_step_offset = 0
+        self._completed_steps = 0
+        self.completed = False
+        self._active_step_limit = None
+        self._run_revision += 1
         if self._reset_callback is not None:
             self._call_maybe_async(self._reset_callback())
         clear_pose_velocity_states()
@@ -717,10 +740,14 @@ def plan_step(self, step_index, last_action):
         """
         import numpy as np
 
+        if self.completed:
+            raise RuntimeError("run is completed; reset before stepping again")
         dt = 1.0 / self._control_hz
         # A hook may request a scene reset after inference. Discard that plan's
         # output and rebuild the first action from reset state before physics.
         for attempt in range(2):
+            if self._active_step_limit is None and self._run_step_limit is not None:
+                self._active_step_limit = self._run_step_limit(self._fill_parameters)
             index = step_index - self._plan_step_offset
             frame = Frame(dt, index * dt, index)
             self._current_frame = frame
@@ -744,6 +771,8 @@ def plan_step(self, step_index, last_action):
             clear_pose_velocity_states()
             self._executor.reset()
             self._plan_step_offset = step_index - round(float(start_time or 0) / dt)
+            self._completed_steps = 0
+            self._active_step_limit = None
             last_action = None
         action = np.array(self._policy_action(), dtype=np.float32)
         ctrl, substep_pd = self._control_payload()
@@ -754,6 +783,9 @@ def plan_step(self, step_index, last_action):
         self._physics.after_control_step()
         if self._recorder is not None:
             self._recorder.record_step(self._physics)
+        self._completed_steps += 1
+        self.completed = (self._active_step_limit is not None
+                          and self._completed_steps >= self._active_step_limit)
         return frame, action, ctrl, substep_pd
 
     def custom_step(self, step_index):

src/botlab/sdk/_controllers.py

@@ -41,6 +41,7 @@ def __init__(
         kd=None,
         effort_limit=None,
         reset_on_change=None,
+        end_after_steps=None,
     ):
         if not isinstance(robot, Robot):
             raise TypeError("PolicyController robot must be a BotLab Robot")
@@ -138,6 +139,13 @@ def __init__(
                 "parameter": parameter, "motions": mappings, "robotId": robot.id,
             }
 
+        if end_after_steps is not None:
+            from ._run_end import encode_end_after_steps
+            if self.modify_action is not None:
+                raise ValueError("end_after_steps cannot be combined with an action modifier")
+            hook = next(node for node in self._plan_data["nodes"] if node["type"] == "proc-action-hook")
+            hook["data"]["endAfterSteps"] = encode_end_after_steps(end_after_steps)
+
     @property
     def plan(self):
         return _plan_value(self._plan_data, "PolicyController plan")

src/botlab/sdk/_run_end.py

@@ -0,0 +1,58 @@
+"""Optional control-step budget. Hosts stop only after committing the last step."""
+
+
+def _positive_steps(value):
+    if isinstance(value, bool) or not isinstance(value, int) or not 0 < value <= 9007199254740991:
+        raise ValueError("end_after_steps must use positive safe integer step counts")
+    return value
+
+
+def encode_end_after_steps(value):
+    if not isinstance(value, dict):
+        return {"steps": _positive_steps(value)}
+    if set(value) != {"parameter", "steps"} or not isinstance(value["parameter"], str) or not value["parameter"].strip():
+        raise ValueError("end_after_steps requires parameter and steps")
+    choices = value["steps"]
+    if not isinstance(choices, dict) or not choices:
+        raise ValueError("end_after_steps steps must map option values to positive step counts")
+    entries = []
+    for option, steps in choices.items():
+        if isinstance(option, bool) or not isinstance(option, int) or abs(option) > 9007199254740991:
+            raise ValueError("end_after_steps option values must be safe integers")
+        entries.append({"value": option, "steps": _positive_steps(steps)})
+    return {"parameter": value["parameter"], "steps": entries}
+
+
+def parse_end_after_steps(raw, parameters):
+    """Validate compiled metadata once and return a resolver for a new run."""
+    def fail():
+        raise ValueError("Invalid end_after_steps: map every option to a positive safe integer step count")
+
+    if not isinstance(raw, dict):
+        fail()
+    if set(raw) == {"steps"}:
+        steps = _positive_steps(raw["steps"])
+        return lambda read: steps
+    if set(raw) != {"parameter", "steps"} or not isinstance(raw["parameter"], str):
+        fail()
+    options = parameters.get(raw["parameter"], {}).get("options")
+    entries = raw["steps"]
+    if not options or not isinstance(entries, list) or len(entries) != len(options):
+        fail()
+    counts = {}
+    for entry in entries:
+        if not isinstance(entry, dict) or set(entry) != {"value", "steps"}:
+            fail()
+        value = entry["value"]
+        if isinstance(value, bool) or not isinstance(value, int) or abs(value) > 9007199254740991:
+            fail()
+        if value in counts or not any(option["value"] == value for option in options):
+            fail()
+        counts[value] = _positive_steps(entry["steps"])
+
+    def resolve(read):
+        selected = read(raw["parameter"])
+        if selected not in counts:
+            fail()
+        return counts[selected]
+    return resolve

tests/fixtures/processors.json

@@ -1474,6 +1474,19 @@
     {
       "type": "proc-action-hook",
       "config": {
+        "endAfterSteps": {
+          "parameter": "motion",
+          "steps": [
+            {
+              "value": 0,
+              "steps": 3
+            },
+            {
+              "value": 2,
+              "steps": 5
+            }
+          ]
+        },
         "resetOnChange": {
           "parameter": "motion",
           "robotId": "robot",

tests/test_headless.py

@@ -250,7 +250,8 @@ def modify(action, frame):
         assert run.step().frame.index == 0
 
 
-def test_compiled_motion_reset_switch_replays_fresh_run(pendulum_project):
+@pytest.mark.parametrize("limited", [False, True])
+def test_compiled_motion_reset_switch_replays_fresh_run(pendulum_project, limited):
     import numpy as np
 
     from botlab._project import load_project
@@ -262,6 +263,8 @@ def test_compiled_motion_reset_switch_replays_fresh_run(pendulum_project):
     source = source.replace('def create_controller(sim):', 'motions = {0: Motion("a.npz"), 2: Motion("b.npz")}\n\ndef create_controller(sim):')
     source = source.replace('concat([robot.commands("gain")])', 'history(clock(unit="steps") + robot.commands("gain"), 2)[[0]]')
     source = source.replace('PolicyController(robot, control=control)', 'PolicyController(robot, control=control, reset_on_change={"parameter": "gain", "motions": motions})')
+    if limited:
+        source = source.replace('reset_on_change={', 'end_after_steps={"parameter": "gain", "steps": {0: 7, 2: 12}}, reset_on_change={')
     path.write_text(source)
     write_motion_npz(pendulum_project / 'a.npz', dofs=1)
     write_motion_npz(pendulum_project / 'b.npz', dofs=1)
@@ -277,3 +280,72 @@ def test_compiled_motion_reset_switch_replays_fresh_run(pendulum_project):
         assert record.frame.index == 0
         np.testing.assert_array_equal(record.action, expected[0][0])
         np.testing.assert_array_equal(run.physics.data.qpos, expected[0][1])
+        if limited:
+            assert len(run.steps(100)) == 11
+            assert run.completed
+            run.reset()
+            assert len(run.steps(100)) == 12
+        # Even a reset after only frame zero must reapply the selected motion pose.
+        for _ in range(2):
+            run.reset()
+            np.testing.assert_array_equal(run.step().action, expected[0][0])
+            np.testing.assert_array_equal(run.physics.data.qpos, expected[0][1])
+
+
+@pytest.mark.parametrize("limit", [1, 3, 8])
+def test_end_after_steps_stops_physics_and_reset_replays(pendulum_project, limit):
+    import numpy as np
+
+    path = pendulum_project / "controller.py"
+    path.write_text(path.read_text().replace(
+        "PolicyController(robot, control=control)",
+        f"PolicyController(robot, control=control, end_after_steps={limit})",
+    ))
+    with open_headless(pendulum_project) as run:
+        states = []
+        records = run.steps(limit + 10, on_step=lambda _: states.append(run.physics.data.qpos.copy()))
+        assert len(records) == limit
+        assert [r.frame.index for r in records] == list(range(limit))
+        assert run.physics.data.time == pytest.approx(limit / run.control_hz)
+        assert run.completed
+        assert run.steps(2) == []
+        with pytest.raises(RuntimeError, match="completed"):
+            run.step()
+        assert run.physics.data.time == pytest.approx(limit / run.control_hz)
+        run.reset()
+        assert not run.completed
+        replay = []
+        assert len(run.steps(limit + 1, on_step=lambda _: replay.append(run.physics.data.qpos.copy()))) == limit
+        np.testing.assert_array_equal(replay, states)
+
+
+def test_parameter_end_limit_is_sampled_on_reset_and_native_run_stops(pendulum_project):
+    path = pendulum_project / "controller.py"
+    source = path.read_text().replace(
+        '{"default": 1.0, "range": [0.0, 2.0]}',
+        '{"default": 0, "options": [{"label": "A", "value": 0}, {"label": "B", "value": 2}]}',
+    ).replace("PolicyController(robot, control=control)",
+              'PolicyController(robot, control=control, end_after_steps={"parameter": "gain", "steps": {0: 3, 2: 5}})')
+    path.write_text(source)
+    with open_headless(pendulum_project) as run:
+        run.step()
+        run.set_parameter("gain", 2)
+        assert len(run.steps(20)) == 2  # this run retains its original budget
+        run.reset()
+        assert len(run.steps(20)) == 5
+        run.set_parameter("gain", 0)
+        assert run.session.run() == 3  # the CLI uses this same loop
+        assert run.completed
+
+
+@pytest.mark.parametrize("limit", [0, -1, True, 1.2, 2**53, {"steps": {0: 2}},
+    {"parameter": "gain", "steps": {0: 0}}])
+def test_end_after_steps_rejects_invalid_declarations(pendulum_project, limit):
+    path = pendulum_project / "controller.py"
+    path.write_text(path.read_text().replace(
+        "PolicyController(robot, control=control)",
+        f"PolicyController(robot, control=control, end_after_steps={limit!r})",
+    ))
+    with pytest.raises(ValueError, match="end_after_steps"):
+        with open_headless(pendulum_project):
+            pass