lucen-cli 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (100) hide show
  1. lucen/__init__.py +72 -0
  2. lucen/errors.py +141 -0
  3. lucen/hub.py +889 -0
  4. lucen/progress.py +96 -0
  5. lucen/py.typed +0 -0
  6. lucen/results.py +885 -0
  7. lucen_cli/__init__.py +3 -0
  8. lucen_cli/checkpoint_dir.py +226 -0
  9. lucen_cli/checkpoint_publisher.py +269 -0
  10. lucen_cli/client.py +1730 -0
  11. lucen_cli/cmd_auth.py +256 -0
  12. lucen_cli/cmd_datasets.py +190 -0
  13. lucen_cli/cmd_holds.py +76 -0
  14. lucen_cli/cmd_hub.py +245 -0
  15. lucen_cli/cmd_model.py +175 -0
  16. lucen_cli/cmd_org.py +154 -0
  17. lucen_cli/cmd_quickstart.py +320 -0
  18. lucen_cli/cmd_repo.py +87 -0
  19. lucen_cli/cmd_robot.py +784 -0
  20. lucen_cli/cmd_rollout.py +576 -0
  21. lucen_cli/cmd_runs.py +615 -0
  22. lucen_cli/cmd_serve.py +438 -0
  23. lucen_cli/cmd_worker.py +135 -0
  24. lucen_cli/config.py +268 -0
  25. lucen_cli/context.py +43 -0
  26. lucen_cli/contract_derive.py +491 -0
  27. lucen_cli/embodiment.py +522 -0
  28. lucen_cli/errors.py +210 -0
  29. lucen_cli/gate_scorer.py +168 -0
  30. lucen_cli/group.py +104 -0
  31. lucen_cli/hf_dataset.py +696 -0
  32. lucen_cli/live.py +189 -0
  33. lucen_cli/local.py +160 -0
  34. lucen_cli/main.py +277 -0
  35. lucen_cli/model_card.py +523 -0
  36. lucen_cli/offscreen.py +52 -0
  37. lucen_cli/onnx_manifest.py +121 -0
  38. lucen_cli/onnx_writer.py +313 -0
  39. lucen_cli/ort_quiet.py +75 -0
  40. lucen_cli/output.py +868 -0
  41. lucen_cli/progress.py +133 -0
  42. lucen_cli/pull.py +158 -0
  43. lucen_cli/push.py +384 -0
  44. lucen_cli/rl/__init__.py +8 -0
  45. lucen_cli/rl/ppo.py +316 -0
  46. lucen_cli/rl/tasks.py +298 -0
  47. lucen_cli/rl/vecenv.py +246 -0
  48. lucen_cli/robot/__init__.py +66 -0
  49. lucen_cli/robot/checks.py +1416 -0
  50. lucen_cli/robot/confine.py +176 -0
  51. lucen_cli/robot/draft.py +523 -0
  52. lucen_cli/robot/hf.py +218 -0
  53. lucen_cli/robot/hub.py +227 -0
  54. lucen_cli/robot/intake.py +590 -0
  55. lucen_cli/robot/model.py +334 -0
  56. lucen_cli/robot/pose.py +132 -0
  57. lucen_cli/robot/publish.py +197 -0
  58. lucen_cli/robot/refs.py +212 -0
  59. lucen_cli/robot/thumbnail.py +586 -0
  60. lucen_cli/robot/urdf.py +435 -0
  61. lucen_cli/robot/worker.py +599 -0
  62. lucen_cli/robot_assets.py +127 -0
  63. lucen_cli/rollout/__init__.py +12 -0
  64. lucen_cli/rollout/assets.py +135 -0
  65. lucen_cli/rollout/conditions.py +486 -0
  66. lucen_cli/rollout/contract.py +314 -0
  67. lucen_cli/rollout/engine.py +1808 -0
  68. lucen_cli/rollout/harness.py +201 -0
  69. lucen_cli/rollout/obs.py +262 -0
  70. lucen_cli/rollout/plant.py +456 -0
  71. lucen_cli/rollout/posture.py +96 -0
  72. lucen_cli/rollout/video.py +548 -0
  73. lucen_cli/rollout/worker.py +505 -0
  74. lucen_cli/serve/__init__.py +22 -0
  75. lucen_cli/serve/lerobot_backend.py +337 -0
  76. lucen_cli/serve/msgpack_numpy.py +59 -0
  77. lucen_cli/serve/policy.py +276 -0
  78. lucen_cli/serve/server.py +474 -0
  79. lucen_cli/serve/session.py +142 -0
  80. lucen_cli/serve/stub_skill.py +264 -0
  81. lucen_cli/serve/vla.py +788 -0
  82. lucen_cli/serve/worker.py +578 -0
  83. lucen_cli/sim_words.py +55 -0
  84. lucen_cli/state.py +237 -0
  85. lucen_cli/templates/__init__.py +94 -0
  86. lucen_cli/templates/ppo_isaac.py +914 -0
  87. lucen_cli/templates/ppo_mujoco.py +664 -0
  88. lucen_cli/templates/smolvla_lora.py +986 -0
  89. lucen_cli/templates/stub_ppo.py +597 -0
  90. lucen_cli/templates/stub_skill.py +180 -0
  91. lucen_cli/tfevents.py +323 -0
  92. lucen_cli/training_curve.py +70 -0
  93. lucen_cli/transfer.py +209 -0
  94. lucen_cli/vla_manifest.py +322 -0
  95. lucen_cli/worker.py +1433 -0
  96. lucen_cli-0.1.0.dist-info/METADATA +113 -0
  97. lucen_cli-0.1.0.dist-info/RECORD +100 -0
  98. lucen_cli-0.1.0.dist-info/WHEEL +4 -0
  99. lucen_cli-0.1.0.dist-info/entry_points.txt +3 -0
  100. lucen_cli-0.1.0.dist-info/licenses/LICENSE +202 -0
@@ -0,0 +1,1808 @@
1
+ """MuJoCo + onnxruntime: run a policy under its `io_contract`, and measure.
2
+
3
+ What runs here is the deployment bridge's loop, in simulation: at the
4
+ contract's control rate the observation is assembled in the contract's order
5
+ and scales, the ONNX graph is asked for an action, the action is clipped and
6
+ scaled onto the default pose to give joint targets, and a PD loop drives the
7
+ joints at the physics rate, torque-clipped at the robot card's deployed
8
+ limits ("simulation clips at the same value" — the card's own words). Laika's
9
+ MJCF has no actuators, so the PD torque is applied through `qfrc_applied`;
10
+ a model that does carry actuators is driven through `ctrl` instead.
11
+
12
+ Measurement, not verdict: an episode reports whether the base stayed upright
13
+ for the whole horizon (and when it stopped), and the mean body-frame
14
+ velocity per commanded axis as a ratio of the command. The hub applies the
15
+ harness rule to those numbers (`lucen_api.rollout.gate`).
16
+ """
17
+
18
+ from __future__ import annotations
19
+
20
+ import hashlib
21
+ import math
22
+ import time
23
+ from collections.abc import Callable
24
+ from dataclasses import dataclass, field
25
+ from pathlib import Path
26
+ from typing import Any
27
+
28
+ # onnxruntime's wheels embed Microsoft's 1DS telemetry client. It is created
29
+ # with the runtime environment, starts an HTTP uploader thread, and on macOS
30
+ # aborts the whole process at exit when a late response handler locks a mutex
31
+ # libc++ has already destroyed (`recursive_mutex lock failed`, one pytest run
32
+ # in three; the crash report's faulting thread is
33
+ # `Microsoft::Applications::Events::HttpClientManager::onHttpResponse`). The
34
+ # binary honours `ORT_DISABLE_TELEMETRY` at environment creation — which is
35
+ # before `disable_telemetry_events()` can run, so only the variable prevents
36
+ # the thread from existing. It has to be set before the import, which is why
37
+ # the runtime comes from `lucen_cli.ort_quiet` (which also keeps the runtime's
38
+ # device-discovery warnings off stderr). A policy gate on a user's machine has
39
+ # no business reporting anywhere in any case.
40
+ import mujoco
41
+ import numpy as np
42
+
43
+ from lucen_cli.onnx_manifest import OnnxReadError, read_lucen_manifest
44
+ from lucen_cli.ort_quiet import ort
45
+
46
+ if hasattr(ort, "disable_telemetry_events"): # belt to the variable's braces
47
+ ort.disable_telemetry_events()
48
+ from lucen_cli.embodiment import Embodiment, InterfaceError
49
+ from lucen_cli.embodiment import card_limits as card_limits_of
50
+ from lucen_cli.rollout import conditions as conditions_mod
51
+ from lucen_cli.rollout import harness
52
+ from lucen_cli.rollout import obs as obs_registry
53
+ from lucen_cli.rollout import posture as posture_mod
54
+ from lucen_cli.rollout.contract import ContractMismatch, IoContract, validate
55
+ from lucen_cli.rollout.plant import Plant, PlantError, build_plant, default_qpos
56
+
57
+ #: Upright: the base keeps at least this fraction of its initial height…
58
+ UPRIGHT_HEIGHT_FRACTION = 0.4
59
+ #: …and tilts less than 60 degrees (projected gravity's z component).
60
+ UPRIGHT_GRAVITY_Z = -0.5
61
+ #: Tracking is averaged after the policy has had this long to settle.
62
+ SETTLE_S = 1.0
63
+ #: Per-seed initial-state noise: joint offsets (rad) and base yaw (rad).
64
+ JOINT_NOISE_RAD = 0.02
65
+ YAW_NOISE_RAD = 0.3
66
+ DEFAULT_GAINS = (40.0, 1.0)
67
+ #: M17b. `joint_target` / `reach` cells: the error is averaged over the last
68
+ #: this-fraction of the horizon ("settle"), and over everything after SETTLE_S
69
+ #: ("tracking"). The PASS tolerances themselves are the hub's (gate.py).
70
+ SETTLE_WINDOW_FRACTION = 0.2
71
+ CELL_KINDS = ("velocity", "joint_target", "reach", "recovery")
72
+ #: C2. How a battery measures: `m8` (M8's definitions, the default) or
73
+ #: `sim2sim` (the founder's harness: its start pose, its survival reading, its
74
+ #: yaw rate from the net heading change, and eval_c_matrix's median tracking).
75
+ PROTOCOL_M8 = "m8"
76
+ PROTOCOL_SIM2SIM = "sim2sim"
77
+ PROTOCOLS = (PROTOCOL_M8, PROTOCOL_SIM2SIM)
78
+ #: sim2sim's absolute survival numbers on Laika (0.384 m standing), as
79
+ #: fractions of the robot's standing height: an episode ends below 0.25 m, and
80
+ #: is alive when it ends above 0.30 m and tilted less than 25°.
81
+ HARNESS_FALL_FRACTION = 0.25 / 0.384
82
+ HARNESS_ALIVE_FRACTION = 0.30 / 0.384
83
+ HARNESS_ALIVE_TILT_DEG = 25.0
84
+ #: The recovery criterion (RECOVERY_V0_SPEC §5, sim2sim `--init-fallen`,
85
+ #: `accept_recovery.py`): tilt, height, both feet, angular rate, held 0.5 s.
86
+ RECOVERY_TILT_DEG = 15.0
87
+ RECOVERY_HEIGHT_FRACTION = 0.85
88
+ RECOVERY_FOOT_N = 5.0
89
+ RECOVERY_OMEGA = 1.0
90
+ RECOVERY_HOLD_S = 0.5
91
+ RECOVERY_TORQUE_SETTLE_S = 0.5
92
+ RECOVERY_ACTION_SAT = 0.99
93
+ #: C3b. After a push on a posture cell, the base has RECOVERED at the first
94
+ #: control step from which, to the end of the episode, it is tilted no more
95
+ #: than the stand battery's tilt bar and moving no faster than 5 cm/s across
96
+ #: the floor. Fixed with the calibration (docs/DECISIONS.md "C3b") and
97
+ #: recorded beside the verdict; the hub gates on the tilt, not on this.
98
+ PUSH_RECOVER_TILT_DEG = 3.0
99
+ PUSH_RECOVER_SPEED_M_S = 0.05
100
+
101
+ FrameSink = Callable[[mujoco.MjData, float], None]
102
+ Log = Callable[[str], None]
103
+
104
+
105
+ def versions() -> dict[str, str]:
106
+ return {"mujoco": str(mujoco.__version__), "onnxruntime": str(ort.__version__)}
107
+
108
+
109
+ # ---------------------------------------------------------------------------
110
+ # The policy
111
+ # ---------------------------------------------------------------------------
112
+
113
+
114
+ @dataclass
115
+ class Policy:
116
+ session: ort.InferenceSession
117
+ input_name: str
118
+ output_name: str
119
+ in_dim: int | None
120
+ out_dim: int | None
121
+ sha256: str
122
+ path: Path
123
+ manifest: dict[str, Any]
124
+
125
+ def __call__(self, obs: np.ndarray) -> np.ndarray:
126
+ out = self.session.run([self.output_name], {self.input_name: obs.reshape(1, -1)})[0]
127
+ return np.asarray(out, dtype=np.float64).reshape(-1)
128
+
129
+ def close(self) -> None:
130
+ """Drop the onnxruntime session now rather than at interpreter exit.
131
+
132
+ Native resources are released when the work is done, not whenever the
133
+ garbage collector gets to them; the worker closes its policy in a
134
+ `finally` and tests close their fixtures.
135
+ """
136
+ self.session = None
137
+
138
+
139
+ def _session_options() -> ort.SessionOptions:
140
+ """One thread, no pool: a locomotion policy is a few thousand MACs a step.
141
+
142
+ With one thread onnxruntime creates no intra-op workers at all, so a
143
+ battery's per-step latency is lower and nothing of the session outlives
144
+ its last use. (The exit abort this was first suspected of turned out to
145
+ be the telemetry thread above.)
146
+ """
147
+ options = ort.SessionOptions()
148
+ options.intra_op_num_threads = 1
149
+ options.inter_op_num_threads = 1
150
+ options.log_severity_level = 3
151
+ return options
152
+
153
+
154
+ def _last_dim(shape: list[Any]) -> int | None:
155
+ if not shape:
156
+ return None
157
+ last = shape[-1]
158
+ return int(last) if isinstance(last, int) else None
159
+
160
+
161
+ def load_policy(path: Path, manifest: dict[str, Any] | None = None) -> Policy:
162
+ """An onnxruntime session plus the contract stamped in the file.
163
+
164
+ `manifest` overrides the stamp only when the file carries none (the
165
+ `io_contract.json` sidecar case); an unstamped ONNX with no sidecar is a
166
+ ContractMismatch, exactly as the two real consumers treat it.
167
+ """
168
+ data = path.read_bytes()
169
+ try:
170
+ stamped = read_lucen_manifest(path)
171
+ except (OnnxReadError, ValueError) as exc:
172
+ raise ContractMismatch(f"{path.name} is not a readable ONNX file: {exc}") from exc
173
+ chosen = stamped or manifest
174
+ if not chosen:
175
+ raise ContractMismatch(
176
+ f"{path.name} carries no `lucen_manifest` and no io_contract.json sidecar was "
177
+ "found; a policy with no contract is refused, as both real consumers refuse it"
178
+ )
179
+ try:
180
+ session = ort.InferenceSession(
181
+ str(path), sess_options=_session_options(), providers=["CPUExecutionProvider"]
182
+ )
183
+ except Exception as exc:
184
+ raise ContractMismatch(f"{path.name}: onnxruntime cannot load it ({exc})") from exc
185
+ inputs, outputs = session.get_inputs(), session.get_outputs()
186
+ if len(inputs) != 1 or len(outputs) < 1:
187
+ raise ContractMismatch(
188
+ f"{path.name}: expected one input and one output, found {len(inputs)} and "
189
+ f"{len(outputs)}"
190
+ )
191
+ return Policy(
192
+ session=session,
193
+ input_name=inputs[0].name,
194
+ output_name=outputs[0].name,
195
+ in_dim=_last_dim(list(inputs[0].shape)),
196
+ out_dim=_last_dim(list(outputs[0].shape)),
197
+ sha256=hashlib.sha256(data).hexdigest(),
198
+ path=path,
199
+ manifest=dict(chosen),
200
+ )
201
+
202
+
203
+ # ---------------------------------------------------------------------------
204
+ # The simulator
205
+ # ---------------------------------------------------------------------------
206
+
207
+
208
+ def model_joints(model: mujoco.MjModel) -> dict[str, int]:
209
+ """`{joint name: DoF count}` for every named joint in the model."""
210
+ out: dict[str, int] = {}
211
+ for j in range(model.njnt):
212
+ name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, j)
213
+ if not name:
214
+ continue
215
+ kind = model.jnt_type[j]
216
+ dofs = {mujoco.mjtJoint.mjJNT_FREE: 6, mujoco.mjtJoint.mjJNT_BALL: 3}.get(kind, 1)
217
+ out[name] = dofs
218
+ return out
219
+
220
+
221
+ def load_model(mjcf_path: Path) -> mujoco.MjModel:
222
+ try:
223
+ return mujoco.MjModel.from_xml_path(str(mjcf_path))
224
+ except Exception as exc:
225
+ raise ContractMismatch(f"MuJoCo refused to compile {mjcf_path.name}: {exc}") from exc
226
+
227
+
228
+ def _yaw_of(quat: np.ndarray) -> float:
229
+ """World yaw of a wxyz quaternion — sim2sim's `yaw_now`."""
230
+ w, x, y, z = (float(v) for v in quat)
231
+ return float(math.atan2(2 * (w * z + x * y), 1 - 2 * (y * y + z * z)))
232
+
233
+
234
+ class Simulator:
235
+ """One robot, one contract: reset, observe, step, and say whether it is upright."""
236
+
237
+ def __init__(
238
+ self,
239
+ model: mujoco.MjModel,
240
+ contract: IoContract,
241
+ *,
242
+ plant: dict[str, Any] | None = None,
243
+ card: dict[str, Any] | None = None,
244
+ emb: Embodiment | None = None,
245
+ built: Plant | None = None,
246
+ ) -> None:
247
+ self.model = model
248
+ self.data = mujoco.MjData(model)
249
+ self.contract = contract
250
+ #: M17b. The robot's own declaration and the plant compiled from it.
251
+ #: Both None on the legacy path (a card with no `interface`).
252
+ self.emb = emb
253
+ self.built = built
254
+ plant = plant or {}
255
+ n = len(contract.joint_order)
256
+ joint_ids = [
257
+ mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, j) for j in contract.joint_order
258
+ ]
259
+ self.joint_ids = joint_ids
260
+ self.qpos_adr = np.asarray([model.jnt_qposadr[j] for j in joint_ids], dtype=np.intp)
261
+ self.dof_adr = np.asarray([model.jnt_dofadr[j] for j in joint_ids], dtype=np.intp)
262
+
263
+ # The floating base, if any.
264
+ self.base_qpos: int | None = None
265
+ self.base_dof: int | None = None
266
+ self.base_body = 1 if model.nbody > 1 else 0
267
+ for j in range(model.njnt):
268
+ if model.jnt_type[j] == mujoco.mjtJoint.mjJNT_FREE:
269
+ self.base_qpos = int(model.jnt_qposadr[j])
270
+ self.base_dof = int(model.jnt_dofadr[j])
271
+ self.base_body = int(model.jnt_bodyid[j])
272
+ break
273
+ self.fixed_base = self.base_qpos is None
274
+
275
+ # Torque limits: the card's deployed `effort` first, then the plant's
276
+ # scalar, then the MJCF's own actuatorfrcrange, then unlimited.
277
+ limits = np.full(n, np.inf)
278
+ card_limits = (card or {}).get("limits") if isinstance(card, dict) else None
279
+ tau_plant = plant.get("tau_limit")
280
+ for k, (name, j) in enumerate(zip(contract.joint_order, joint_ids, strict=True)):
281
+ effort = None
282
+ if isinstance(card_limits, dict) and isinstance(card_limits.get(name), dict):
283
+ effort = card_limits[name].get("effort")
284
+ if effort is None and isinstance(tau_plant, int | float):
285
+ effort = float(tau_plant)
286
+ if effort is None and model.jnt_actfrclimited[j]:
287
+ effort = float(model.jnt_actfrcrange[j][1])
288
+ if effort is not None:
289
+ limits[k] = abs(float(effort))
290
+ if emb is not None:
291
+ # The robot's declared effort limit is what its plant clips at.
292
+ for k in range(n):
293
+ if emb.effort_limit[k] is not None:
294
+ limits[k] = abs(float(emb.effort_limit[k])) # type: ignore[arg-type]
295
+ self.tau_limit = limits
296
+
297
+ gains = plant.get("gains_rl_default")
298
+ if isinstance(gains, list | tuple) and len(gains) == 2:
299
+ self.kp, self.kd = float(gains[0]), float(gains[1])
300
+ else:
301
+ self.kp, self.kd = DEFAULT_GAINS
302
+ self._base_kp, self._base_kd = self.kp, self.kd
303
+
304
+ # Actuators, when the model has them: one per contract joint at most.
305
+ self.actuator_of: dict[int, tuple[int, bool]] = {}
306
+ for a in range(model.nu):
307
+ if model.actuator_trntype[a] != mujoco.mjtTrn.mjTRN_JOINT:
308
+ continue
309
+ joint = int(model.actuator_trnid[a][0])
310
+ position_servo = (
311
+ model.actuator_biastype[a] == mujoco.mjtBias.mjBIAS_AFFINE
312
+ and model.actuator_biasprm[a][1] < 0
313
+ )
314
+ for k, j in enumerate(joint_ids):
315
+ if j == joint:
316
+ self.actuator_of[k] = (a, bool(position_servo))
317
+
318
+ self.dt = float(model.opt.timestep)
319
+ # The robot declares its control rate; a contract that disagrees was
320
+ # already refused by `validate` via the fingerprint, so when there is
321
+ # an interface its rate is the one that holds.
322
+ rate = emb.control_rate_hz if emb is not None else contract.control_rate_hz
323
+ self.control_rate_hz = float(rate)
324
+ self.decimation = max(1, round(1.0 / (self.dt * self.control_rate_hz)))
325
+ self.control_dt = self.dt * self.decimation
326
+
327
+ # Soft joint limits for limit-violation counting: the card's
328
+ # `soft_lower`/`soft_upper` when it has them, else its `lower`/`upper`
329
+ # (else the MJCF range) shrunk by the interface's soft_limit_factor.
330
+ self.soft_lo, self.soft_hi = self._soft_limits(card, joint_ids)
331
+ # Hard ranges (the MJCF's, else the card's): where the founder's fallen
332
+ # starts draw their joints from.
333
+ hard = card_limits_of(card, contract.joint_order)
334
+ self.range_lo = np.full(n, -math.pi)
335
+ self.range_hi = np.full(n, math.pi)
336
+ for k, j in enumerate(joint_ids):
337
+ if model.jnt_limited[j]:
338
+ self.range_lo[k], self.range_hi[k] = model.jnt_range[j]
339
+ else:
340
+ if hard["lower"][k] is not None:
341
+ self.range_lo[k] = float(hard["lower"][k])
342
+ if hard["upper"][k] is not None:
343
+ self.range_hi[k] = float(hard["upper"][k])
344
+ self.ee_name: str | None = None
345
+ if built is not None and built.ee_sensor:
346
+ self.ee_name = next(iter(built.ee_sensor))
347
+
348
+ # C2: contract features past the base 45 dims (policy_io's own).
349
+ try:
350
+ self.lat_ff = harness.LateralFeedForward.from_contract(
351
+ contract.raw, contract.joint_order, contract.cycle_time_s
352
+ )
353
+ self.beta = harness.beta_vector(contract.raw, contract.joint_order)
354
+ except harness.HarnessContractError as exc:
355
+ raise ContractMismatch(str(exc)) from exc
356
+ self.history_dim = sum(b.dim for b in contract.obs_spec if b.name == harness.HISTORY_BLOCK)
357
+ self.history: harness.FrameHistory | None = (
358
+ harness.FrameHistory(self.history_dim, n) if self.history_dim else None
359
+ )
360
+ # C2: per-joint velocity limits for the bridge slew, from the interface.
361
+ self.velocity_limit = np.full(n, np.inf)
362
+ if emb is not None:
363
+ for k in range(n):
364
+ if emb.velocity_limit[k] is not None:
365
+ self.velocity_limit[k] = float(emb.velocity_limit[k]) # type: ignore[arg-type]
366
+ # C2: what conditions scale is always scaled from these, never compounded.
367
+ self._base_friction = np.asarray(model.geom_friction, dtype=np.float64).copy()
368
+ self._base_gainprm = np.asarray(model.actuator_gainprm, dtype=np.float64).copy()
369
+ self._base_biasprm = np.asarray(model.actuator_biasprm, dtype=np.float64).copy()
370
+ self.cond = conditions_mod.Conditions()
371
+ self._total_mass = float(model.body_subtreemass[0])
372
+ self._foot_geoms: list[set[int]] = []
373
+ if emb is not None:
374
+ for foot in emb.feet:
375
+ body = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, foot)
376
+ if body >= 0:
377
+ self._foot_geoms.append(
378
+ {g for g in range(model.ngeom) if int(model.geom_bodyid[g]) == body}
379
+ )
380
+
381
+ # Initial state: the `standing` keyframe, else keyframe 0, else qpos0.
382
+ self.keyframe = -1
383
+ self.standing_keyframe = -1
384
+ for k in range(model.nkey):
385
+ if mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_KEY, k) == "standing":
386
+ self.keyframe = k
387
+ self.standing_keyframe = k
388
+ if self.keyframe < 0 and model.nkey > 0:
389
+ self.keyframe = 0
390
+ self._home_qpos: np.ndarray | None = None
391
+ if emb is not None and built is not None:
392
+ try:
393
+ self._home_qpos = default_qpos(built, emb)
394
+ except PlantError as exc:
395
+ raise ContractMismatch(str(exc)) from exc
396
+ # The policy's own nominal pose is what its actions are offsets from
397
+ # and what `joint_pos_rel` is measured against — the contract's, as
398
+ # both of the founder's consumers read it (policy_io's `default`). M17b
399
+ # took the robot's `default_pose` instead; for every policy trained
400
+ # against the interface they are the same numbers, and for a published
401
+ # Lucen policy whose profile has its own nominal (stand: all zeros;
402
+ # recovery: mid-range) the robot's pose was simply the wrong reference.
403
+ self._reset_state()
404
+ init = self.data.qpos[self.qpos_adr].copy()
405
+ self.default_pose = (
406
+ contract.default_pose.astype(np.float64)
407
+ if contract.default_pose is not None
408
+ else init.astype(np.float64)
409
+ )
410
+ self._reset_state()
411
+ self.init_height = float(self.base_position()[2]) if not self.fixed_base else 0.0
412
+ # C2: the founder's start (`sim2sim.py`): the `standing` keyframe (else
413
+ # the home pose), the policy joints at the policy's nominal. Its base
414
+ # height is the robot's standing height the harness's absolute numbers
415
+ # were written against (0.384 m on Laika).
416
+ self._harness_qpos = self._harness_home()
417
+ #: C2, parity only: read the base ROTATION the way `sim2sim.py` does —
418
+ #: from `xmat`, one physics substep old (see `_fresh_base`). Never set
419
+ #: by a battery; the parity harness sets it to show that, with this one
420
+ #: difference removed, the gate and the founder's harness agree step for step.
421
+ self.stale_rotation = False
422
+ self.stand_height = (
423
+ float(self._harness_qpos[self.base_qpos + 2]) if self.base_qpos is not None else 0.0
424
+ )
425
+ self.time_s = 0.0
426
+
427
+ # -- state -------------------------------------------------------------
428
+
429
+ def _soft_limits(
430
+ self, card: dict[str, Any] | None, joint_ids: list[int]
431
+ ) -> tuple[np.ndarray, np.ndarray]:
432
+ n = len(joint_ids)
433
+ lo, hi = np.full(n, -np.inf), np.full(n, np.inf)
434
+ for k, j in enumerate(joint_ids):
435
+ if self.model.jnt_limited[j]:
436
+ lo[k], hi[k] = self.model.jnt_range[j]
437
+ hard = card_limits_of(card, self.contract.joint_order)
438
+ for k in range(n):
439
+ if hard["lower"][k] is not None:
440
+ lo[k] = float(hard["lower"][k])
441
+ if hard["upper"][k] is not None:
442
+ hi[k] = float(hard["upper"][k])
443
+ factor = self.emb.soft_limit_factor if self.emb is not None else 1.0
444
+ centre, half = (lo + hi) / 2.0, (hi - lo) / 2.0
445
+ with np.errstate(invalid="ignore"):
446
+ soft_lo = np.where(np.isfinite(lo), centre - factor * half, -np.inf)
447
+ soft_hi = np.where(np.isfinite(hi), centre + factor * half, np.inf)
448
+ limits = (card or {}).get("limits") if isinstance(card, dict) else None
449
+ for k, name in enumerate(self.contract.joint_order):
450
+ entry = limits.get(name) if isinstance(limits, dict) else None
451
+ if isinstance(entry, dict):
452
+ if isinstance(entry.get("soft_lower"), int | float):
453
+ soft_lo[k] = float(entry["soft_lower"])
454
+ if isinstance(entry.get("soft_upper"), int | float):
455
+ soft_hi[k] = float(entry["soft_upper"])
456
+ return soft_lo, soft_hi
457
+
458
+ def limit_violation(self) -> bool:
459
+ """Is any policy joint outside its soft limit right now?"""
460
+ q = self.data.qpos[self.qpos_adr]
461
+ return bool(np.any(q < self.soft_lo - 1e-9) or np.any(q > self.soft_hi + 1e-9))
462
+
463
+ def ee_position(self) -> np.ndarray | None:
464
+ """The first declared end effector's position, or None when there is none."""
465
+ if self.built is None or self.ee_name is None:
466
+ return None
467
+ start, stop = self.built.ee_sensor[self.ee_name]
468
+ return np.asarray(self.data.sensordata[start:stop], dtype=np.float64).copy()
469
+
470
+ def _reset_state(self) -> None:
471
+ if self._home_qpos is not None:
472
+ mujoco.mj_resetData(self.model, self.data)
473
+ self.data.qpos[:] = self._home_qpos
474
+ elif self.keyframe >= 0:
475
+ mujoco.mj_resetDataKeyframe(self.model, self.data, self.keyframe)
476
+ else:
477
+ mujoco.mj_resetData(self.model, self.data)
478
+ pose = getattr(self, "default_pose", None)
479
+ if pose is not None:
480
+ self.data.qpos[self.qpos_adr] = pose
481
+ mujoco.mj_forward(self.model, self.data)
482
+
483
+ def _harness_home(self) -> np.ndarray:
484
+ data = mujoco.MjData(self.model)
485
+ if self.standing_keyframe >= 0:
486
+ mujoco.mj_resetDataKeyframe(self.model, data, self.standing_keyframe)
487
+ elif self._home_qpos is not None:
488
+ data.qpos[:] = self._home_qpos
489
+ elif self.keyframe >= 0:
490
+ mujoco.mj_resetDataKeyframe(self.model, data, self.keyframe)
491
+ data.qpos[self.qpos_adr] = self.default_pose
492
+ return np.asarray(data.qpos, dtype=np.float64).copy()
493
+
494
+ def configure(self, cond: conditions_mod.Conditions) -> list[str]:
495
+ """Apply a cell's plant conditions (friction, gains); returns what could not be applied.
496
+
497
+ Always from the pristine copies taken at construction, so one cell's
498
+ `kd_scale` never leaks into the next.
499
+ """
500
+ notes: list[str] = []
501
+ self.cond = cond
502
+ self.model.geom_friction[:] = self._base_friction
503
+ self.model.actuator_gainprm[:] = self._base_gainprm
504
+ self.model.actuator_biasprm[:] = self._base_biasprm
505
+ self.kp, self.kd = self._base_kp, self._base_kd
506
+ if cond.friction is not None:
507
+ # sim2sim `--friction`: the sliding coefficient of every geom.
508
+ self.model.geom_friction[:, 0] = cond.friction
509
+ if cond.kp_scale != 1.0 or cond.kd_scale != 1.0:
510
+ if self.built is not None:
511
+ for a in self.built.ctrl_adr:
512
+ if self.model.actuator_biasprm[a, 1] < 0: # position servo
513
+ self.model.actuator_gainprm[a, 0] *= cond.kp_scale
514
+ self.model.actuator_biasprm[a, 1] *= cond.kp_scale
515
+ self.model.actuator_biasprm[a, 2] *= cond.kd_scale
516
+ elif self.model.actuator_biasprm[a, 2] < 0: # velocity servo
517
+ self.model.actuator_gainprm[a, 0] *= cond.kd_scale
518
+ self.model.actuator_biasprm[a, 2] *= cond.kd_scale
519
+ else:
520
+ self.kp = self._base_kp * cond.kp_scale
521
+ self.kd = self._base_kd * cond.kd_scale
522
+ if cond.slew and not np.all(np.isfinite(self.velocity_limit)):
523
+ unlimited = [
524
+ self.contract.joint_order[k]
525
+ for k in range(len(self.velocity_limit))
526
+ if not np.isfinite(self.velocity_limit[k])
527
+ ]
528
+ notes.append(
529
+ "bridge slew: no `velocity_limit` declared for "
530
+ + ", ".join(unlimited)
531
+ + " — those joints are held to 1 rad/s until homed and not limited after"
532
+ )
533
+ if cond.fallen is not None and not self._foot_geoms:
534
+ raise ContractMismatch(
535
+ "a fallen start is scored on both feet touching the ground, and the robot's "
536
+ "interface declares no `feet` in the MJCF"
537
+ )
538
+ return notes
539
+
540
+ def reset(
541
+ self,
542
+ seed: int,
543
+ *,
544
+ protocol: str = PROTOCOL_M8,
545
+ category: str | None = None,
546
+ ) -> conditions_mod.FallenState | None:
547
+ """Put the robot in seed `seed`'s start state.
548
+
549
+ `m8` — M8's start, unchanged apart from the policy's own nominal pose:
550
+ joint noise ±0.02 rad and base yaw ±0.3 rad (a condition's `jitter` /
551
+ `yaw_jitter` replaces either). `sim2sim` — the founder's: the
552
+ `standing` keyframe with the policy joints at their nominal, jittered by
553
+ `default_rng(seed).uniform(±jitter)` and clipped to the joint ranges,
554
+ yaw untouched. A fallen category (a recovery cell) replaces both with
555
+ `sample_fallen`, draw for draw the founder's `--init-fallen`.
556
+ """
557
+ cond = self.cond
558
+ fallen_state: conditions_mod.FallenState | None = None
559
+ if category is not None:
560
+ fallen = cond.fallen
561
+ jitter = fallen.jitter_rad if fallen is not None else 0.26
562
+ mujoco.mj_resetData(self.model, self.data)
563
+ self.data.qpos[:] = self._harness_qpos
564
+ (
565
+ fallen_state,
566
+ base_pos,
567
+ base_quat,
568
+ joints,
569
+ base_vel,
570
+ joint_vel,
571
+ ) = conditions_mod.sample_fallen(
572
+ category=category,
573
+ seed=seed,
574
+ jitter_rad=jitter,
575
+ lower=self.range_lo,
576
+ upper=self.range_hi,
577
+ stand_height=self.stand_height or conditions_mod.LAIKA_STAND_HEIGHT_M,
578
+ )
579
+ assert self.base_qpos is not None and self.base_dof is not None
580
+ self.data.qpos[self.base_qpos : self.base_qpos + 3] = base_pos
581
+ self.data.qpos[self.base_qpos + 3 : self.base_qpos + 7] = base_quat
582
+ self.data.qpos[self.qpos_adr] = joints
583
+ self.data.qvel[:] = 0.0
584
+ self.data.qvel[self.base_dof : self.base_dof + 6] = base_vel
585
+ self.data.qvel[self.dof_adr] = joint_vel
586
+ elif protocol == PROTOCOL_SIM2SIM:
587
+ mujoco.mj_resetData(self.model, self.data)
588
+ self.data.qpos[:] = self._harness_qpos
589
+ if cond.jitter:
590
+ rng = np.random.default_rng(seed)
591
+ q = self.data.qpos[self.qpos_adr] + rng.uniform(
592
+ -cond.jitter, cond.jitter, len(self.qpos_adr)
593
+ )
594
+ self.data.qpos[self.qpos_adr] = np.clip(q, self.range_lo, self.range_hi)
595
+ if cond.yaw_jitter and self.base_qpos is not None:
596
+ self._set_yaw(float(rng.uniform(-cond.yaw_jitter, cond.yaw_jitter)))
597
+ elif cond.yaw_jitter and self.base_qpos is not None:
598
+ rng = np.random.default_rng(seed)
599
+ self._set_yaw(float(rng.uniform(-cond.yaw_jitter, cond.yaw_jitter)))
600
+ self.data.qvel[:] = 0.0
601
+ else:
602
+ self._reset_state()
603
+ rng = np.random.default_rng(seed)
604
+ joint_noise = JOINT_NOISE_RAD if cond.jitter is None else cond.jitter
605
+ yaw_noise = YAW_NOISE_RAD if cond.yaw_jitter is None else cond.yaw_jitter
606
+ self.data.qpos[self.qpos_adr] += rng.uniform(
607
+ -joint_noise, joint_noise, len(self.qpos_adr)
608
+ )
609
+ if self.base_qpos is not None:
610
+ self._set_yaw(float(rng.uniform(-yaw_noise, yaw_noise)))
611
+ self.data.qvel[:] = 0.0
612
+ self.data.qfrc_applied[:] = 0.0
613
+ mujoco.mj_forward(self.model, self.data)
614
+ self.time_s = 0.0
615
+ if self.history is not None:
616
+ self.history.reset()
617
+ return fallen_state
618
+
619
+ def _set_yaw(self, yaw: float) -> None:
620
+ assert self.base_qpos is not None
621
+ quat = np.zeros(4)
622
+ mujoco.mju_axisAngle2Quat(quat, np.array([0.0, 0.0, 1.0]), yaw)
623
+ self.data.qpos[self.base_qpos + 3 : self.base_qpos + 7] = quat
624
+
625
+ def _fresh_base(self) -> bool:
626
+ """Read the base from `qpos` (current) rather than `xpos`/`xmat` (one substep old)?
627
+
628
+ `mj_step` runs its forward pass BEFORE it integrates, so after a step
629
+ `data.xpos` / `data.xmat` describe the state the step STARTED from
630
+ while `data.qpos` is already the new one. M8 read the derived arrays,
631
+ so its `projected_gravity` (and the rotation it applies to the base
632
+ velocity) lagged the joints by one physics substep. Nothing computed
633
+ the same quantity a second way, so nobody saw it — until M17b's
634
+ trainer, which derives the base from the rollout's output `qpos`, and
635
+ a parity test found the two disagreeing by 1.3% on a moving Laika.
636
+
637
+ On the interface path the gate now reads `qpos` too: the policy is
638
+ shown the current state, and trainer and gate are identical. The
639
+ legacy path (a card with no `interface`) is deliberately left exactly
640
+ as M8 shipped it — "every existing battery is unchanged" includes the
641
+ numbers the published policies were gated under; the staleness there
642
+ is reported in docs/DECISIONS.md rather than silently changed.
643
+ """
644
+ return self.built is not None and self.base_qpos is not None
645
+
646
+ def base_position(self) -> np.ndarray:
647
+ if self._fresh_base():
648
+ assert self.base_qpos is not None
649
+ return np.asarray(
650
+ self.data.qpos[self.base_qpos : self.base_qpos + 3], dtype=np.float64
651
+ ).copy()
652
+ return np.asarray(self.data.xpos[self.base_body], dtype=np.float64)
653
+
654
+ def base_rotation(self) -> np.ndarray:
655
+ if self._fresh_base() and not self.stale_rotation:
656
+ assert self.base_qpos is not None
657
+ quat = np.asarray(
658
+ self.data.qpos[self.base_qpos + 3 : self.base_qpos + 7], dtype=np.float64
659
+ )
660
+ matrix = np.zeros(9)
661
+ mujoco.mju_quat2Mat(matrix, quat / np.linalg.norm(quat))
662
+ return matrix.reshape(3, 3)
663
+ return np.asarray(self.data.xmat[self.base_body], dtype=np.float64).reshape(3, 3)
664
+
665
+ def base_velocity_body(self) -> tuple[np.ndarray, np.ndarray]:
666
+ """`(linear, angular)` in the base frame. Zero on a fixed base."""
667
+ if self.base_dof is None:
668
+ return np.zeros(3), np.zeros(3)
669
+ rot = self.base_rotation()
670
+ lin_world = np.asarray(self.data.qvel[self.base_dof : self.base_dof + 3])
671
+ ang_body = np.asarray(self.data.qvel[self.base_dof + 3 : self.base_dof + 6])
672
+ return np.asarray(rot.T @ lin_world, dtype=np.float64), ang_body.astype(np.float64)
673
+
674
+ def projected_gravity(self) -> np.ndarray:
675
+ return np.asarray(self.base_rotation().T @ np.array([0.0, 0.0, -1.0]), dtype=np.float64)
676
+
677
+ def mirror_pairs(self) -> list[tuple[int, int, int]]:
678
+ """`(left, right, sign)` per mirrored joint pair of the contract (C3, `posture`).
679
+
680
+ Indices into `contract.joint_order`; the sign is read from the joint
681
+ axes at the model's reference pose, in the base frame. Computed once.
682
+ """
683
+ cached = getattr(self, "_mirror_pairs", None)
684
+ if cached is not None:
685
+ return list(cached)
686
+ scratch = mujoco.MjData(self.model)
687
+ scratch.qpos[:] = self.model.qpos0
688
+ mujoco.mj_kinematics(self.model, scratch)
689
+ base_rot = np.asarray(scratch.xmat[self.base_body], dtype=np.float64).reshape(3, 3)
690
+ out: list[tuple[int, int, int]] = []
691
+ for left, right in posture_mod.mirror_name_pairs(list(self.contract.joint_order)):
692
+ axis_l = base_rot.T @ np.asarray(scratch.xaxis[self.joint_ids[left]], dtype=np.float64)
693
+ axis_r = base_rot.T @ np.asarray(scratch.xaxis[self.joint_ids[right]], dtype=np.float64)
694
+ sign = posture_mod.mirror_sign(axis_l, axis_r)
695
+ if sign is not None:
696
+ out.append((left, right, sign))
697
+ self._mirror_pairs = out
698
+ return list(out)
699
+
700
+ def tilt_deg(self) -> float:
701
+ """Angle between the base's z axis and the world's — sim2sim's `tilts`."""
702
+ return float(math.degrees(math.acos(max(-1.0, min(1.0, self.base_rotation()[2, 2])))))
703
+
704
+ def yaw(self) -> float:
705
+ if self.base_qpos is None:
706
+ return 0.0
707
+ return _yaw_of(self.data.qpos[self.base_qpos + 3 : self.base_qpos + 7])
708
+
709
+ def upright(self) -> bool:
710
+ if self.fixed_base:
711
+ return True
712
+ height = float(self.base_position()[2])
713
+ return (
714
+ height >= UPRIGHT_HEIGHT_FRACTION * self.init_height
715
+ and float(self.projected_gravity()[2]) <= UPRIGHT_GRAVITY_Z
716
+ )
717
+
718
+ def foot_forces(self) -> list[float]:
719
+ """Normal contact force on each declared foot (N), summed over its contacts."""
720
+ out = [0.0] * len(self._foot_geoms)
721
+ force = np.zeros(6)
722
+ for i in range(self.data.ncon):
723
+ contact = self.data.contact[i]
724
+ for k, geoms in enumerate(self._foot_geoms):
725
+ if contact.geom1 in geoms or contact.geom2 in geoms:
726
+ mujoco.mj_contactForce(self.model, self.data, i, force)
727
+ out[k] += abs(float(force[0]))
728
+ return out
729
+
730
+ # -- the contract's observation ---------------------------------------
731
+
732
+ def frame(
733
+ self,
734
+ command: np.ndarray,
735
+ last_action: np.ndarray,
736
+ *,
737
+ joint_target: np.ndarray | None = None,
738
+ ee_target: np.ndarray | None = None,
739
+ ) -> obs_registry.Frame:
740
+ """This instant as a one-env :class:`Frame` — what every block builder reads."""
741
+ lin, ang = self.base_velocity_body()
742
+ ee = self.ee_position()
743
+ return obs_registry.Frame(
744
+ joint_pos=np.asarray(self.data.qpos[self.qpos_adr], dtype=np.float64)[None],
745
+ joint_vel=np.asarray(self.data.qvel[self.dof_adr], dtype=np.float64)[None],
746
+ default_pose=self.default_pose,
747
+ last_action=np.asarray(last_action, dtype=np.float64)[None],
748
+ base_lin_vel=lin[None],
749
+ base_ang_vel=ang[None],
750
+ projected_gravity=self.projected_gravity()[None],
751
+ base_height=np.array([float(self.base_position()[2])]),
752
+ time_s=np.array([self.time_s]),
753
+ command=np.asarray(command, dtype=np.float64)[None],
754
+ cycle_time_s=self.contract.cycle_time_s,
755
+ joint_target=None if joint_target is None else joint_target[None],
756
+ ee_pos=None if ee is None else ee[None],
757
+ ee_target=None if ee_target is None else ee_target[None],
758
+ )
759
+
760
+ def observe(
761
+ self,
762
+ command: np.ndarray,
763
+ last_action: np.ndarray,
764
+ *,
765
+ joint_target: np.ndarray | None = None,
766
+ ee_target: np.ndarray | None = None,
767
+ ) -> np.ndarray:
768
+ """The contract's observation, in its order and scales.
769
+
770
+ `frame_hist` (C2) is built from the same scaled terms as the head of
771
+ the observation and advances exactly once per call — `observe` is
772
+ called once per control step, as `PolicyIO.build_obs` requires.
773
+ """
774
+ frame = self.frame(command, last_action, joint_target=joint_target, ee_target=ee_target)
775
+ try:
776
+ parts: list[np.ndarray] = []
777
+ scaled: dict[str, np.ndarray] = {}
778
+ for block in self.contract.obs_spec:
779
+ if block.name == harness.HISTORY_BLOCK:
780
+ parts.append(np.zeros(0, dtype=np.float32)) # placeholder, filled below
781
+ continue
782
+ value = obs_registry.build_observation(frame, [(block.name, block.scale)])[0]
783
+ scaled[obs_registry.resolve(block.name).name] = value
784
+ parts.append(value)
785
+ if self.history is not None:
786
+ current = np.concatenate(
787
+ [scaled[obs_registry.resolve(t).name] for t in harness.HIST_FRAME_TERMS]
788
+ )
789
+ block_value = self.history.block(current)
790
+ for index, block in enumerate(self.contract.obs_spec):
791
+ if block.name == harness.HISTORY_BLOCK:
792
+ parts[index] = block_value
793
+ return np.asarray(np.concatenate(parts), dtype=np.float32)
794
+ except (obs_registry.UnknownObsTerm, ValueError, KeyError) as exc:
795
+ raise ContractMismatch(f"cannot build the observation: {exc}") from exc
796
+
797
+ # -- one control step --------------------------------------------------
798
+
799
+ def phase(self) -> float | None:
800
+ cycle = self.contract.cycle_time_s
801
+ return None if cycle is None else (self.time_s / cycle) % 1.0
802
+
803
+ def command_for(self, action: np.ndarray, *, command: np.ndarray | None = None) -> np.ndarray:
804
+ """The per-joint command an action means, under the contract's semantics.
805
+
806
+ `position_offset` (every contract before M17b): `default_pose + scale *
807
+ action`. `position_absolute`, `velocity` and `torque` drop the pose.
808
+ Position targets are clamped to the contract's `target_lo/hi` when it
809
+ carries them and to the joint range when the plant came from an
810
+ interface — a target past a hard stop is a fight with the limit, not a
811
+ command. C2 adds policy_io's two extras: the `omni_ff` lateral
812
+ feed-forward (on the commanded `vy` and the phase clock) and the
813
+ β-anchored recovery contract, `target = q + β · action`.
814
+ """
815
+ clip = self.contract.action_clip
816
+ action = np.clip(np.asarray(action, dtype=np.float64), -clip, clip)
817
+ semantics = self.contract.action_semantics
818
+ if self.beta is not None:
819
+ out = np.asarray(self.data.qpos[self.qpos_adr], dtype=np.float64) + self.beta * action
820
+ else:
821
+ scaled = self.contract.action_scale * action
822
+ out = self.default_pose + scaled if semantics == "position_offset" else scaled
823
+ if self.lat_ff is not None:
824
+ vy = float(command[1]) if command is not None else 0.0
825
+ out = out + self.lat_ff.offset(len(out), vy, self.phase() or 0.0)
826
+ if semantics.startswith("position") or self.beta is not None:
827
+ if self.contract.target_lo is not None and self.contract.target_hi is not None:
828
+ out = np.clip(out, self.contract.target_lo, self.contract.target_hi)
829
+ if self.built is not None:
830
+ out = np.clip(out, self.built.lower, self.built.upper)
831
+ return np.asarray(out, dtype=np.float64)
832
+
833
+ def gains(self) -> tuple[np.ndarray, np.ndarray]:
834
+ """Per-joint `(kp, kd)` the plant drives with right now (after `kp_scale`/`kd_scale`)."""
835
+ n = len(self.qpos_adr)
836
+ if self.built is not None:
837
+ kp = np.asarray([self.model.actuator_gainprm[a, 0] for a in self.built.ctrl_adr])
838
+ kd = np.asarray([-self.model.actuator_biasprm[a, 2] for a in self.built.ctrl_adr])
839
+ return kp, kd
840
+ return np.full(n, self.kp), np.full(n, self.kd)
841
+
842
+ def step(
843
+ self,
844
+ action: np.ndarray,
845
+ *,
846
+ target: np.ndarray | None = None,
847
+ torque_probe: Callable[[np.ndarray, bool], None] | None = None,
848
+ ) -> np.ndarray:
849
+ """Apply one policy action for `decimation` physics steps; returns the joint command.
850
+
851
+ `target` (C2) overrides the command the action would give — the bridge
852
+ slew and the delay line sit between the policy and the PD loop.
853
+ `torque_probe` is called before every physics substep with the PD
854
+ demand BEFORE the effort clip and whether it is the control step's
855
+ last substep (the recovery gate's torque bookkeeping: sim2sim's
856
+ `tau_demand_peak` per substep and `tau_demand_peak_ctrlrate` at the
857
+ control-step boundary).
858
+ """
859
+ if target is None:
860
+ target = self.command_for(action)
861
+ kp_v, kd_v = self.gains() if torque_probe is not None else (None, None)
862
+ if self.built is not None:
863
+ # Interface path: the servo (or motor) the plant builder compiled
864
+ # runs the loop inside mj_step, at the physics rate, with the
865
+ # robot's own per-group gains and effort limits.
866
+ self.data.ctrl[self.built.ctrl_adr] = target
867
+ for sub in range(self.decimation):
868
+ if torque_probe is not None:
869
+ assert kp_v is not None and kd_v is not None
870
+ torque_probe(
871
+ kp_v * (target - self.data.qpos[self.qpos_adr])
872
+ - kd_v * self.data.qvel[self.dof_adr],
873
+ sub == self.decimation - 1,
874
+ )
875
+ mujoco.mj_step(self.model, self.data)
876
+ self.time_s += self.control_dt
877
+ return target
878
+ if self.contract.action_semantics != "position_offset" and self.beta is None:
879
+ raise ContractMismatch(
880
+ f"action semantics `{self.contract.action_semantics}` needs a robot card with "
881
+ "an `interface` (it declares the actuators); this card has none"
882
+ )
883
+ for sub in range(self.decimation):
884
+ q = self.data.qpos[self.qpos_adr]
885
+ qd = self.data.qvel[self.dof_adr]
886
+ demand = self.kp * (target - q) - self.kd * qd
887
+ if torque_probe is not None:
888
+ torque_probe(np.asarray(demand, dtype=np.float64), sub == self.decimation - 1)
889
+ tau = np.clip(demand, -self.tau_limit, self.tau_limit)
890
+ self.data.qfrc_applied[:] = 0.0
891
+ for k, dof in enumerate(self.dof_adr):
892
+ actuator = self.actuator_of.get(k)
893
+ if actuator is None:
894
+ self.data.qfrc_applied[dof] = tau[k]
895
+ else:
896
+ index, position_servo = actuator
897
+ self.data.ctrl[index] = target[k] if position_servo else tau[k]
898
+ mujoco.mj_step(self.model, self.data)
899
+ self.time_s += self.control_dt
900
+ return target
901
+
902
+ def push(self, impulse: float, axis: int) -> None:
903
+ """sim2sim `--push`: `impulse / total mass` added to the base's world velocity."""
904
+ if self.base_dof is None:
905
+ return
906
+ self.data.qvel[self.base_dof + axis] += impulse / self._total_mass
907
+
908
+
909
+ # ---------------------------------------------------------------------------
910
+ # Episodes and cells
911
+ # ---------------------------------------------------------------------------
912
+
913
+
914
+ def _wire_metrics(metrics: dict[str, float]) -> dict[str, float | None]:
915
+ """`inf` (a window with no sample) is not JSON; on the wire it is null.
916
+
917
+ Null means "could not be measured" and the hub scores it as such — it is
918
+ deliberately NOT zero, which would read as a perfect result.
919
+ """
920
+ return {k: (float(v) if math.isfinite(v) else None) for k, v in metrics.items()}
921
+
922
+
923
+ @dataclass
924
+ class EpisodeResult:
925
+ seed: int
926
+ alive: bool
927
+ fell_at_s: float | None
928
+ tracking: dict[str, float]
929
+ steps: int
930
+ trajectory: dict[str, Any] | None = None
931
+ #: M17b. `settle_error` / `tracking_error` / `limit_violations`; empty for
932
+ #: a velocity cell on a floating base, which is scored on `alive` + `tracking`.
933
+ metrics: dict[str, float] = field(default_factory=dict)
934
+ #: C2. The fallen category a recovery episode started in.
935
+ category: str | None = None
936
+
937
+ def as_measurement(self) -> dict[str, Any]:
938
+ out: dict[str, Any] = {
939
+ "seed": self.seed,
940
+ "alive": self.alive,
941
+ "fell_at_s": self.fell_at_s,
942
+ "tracking": dict(self.tracking),
943
+ }
944
+ if self.metrics:
945
+ out["metrics"] = _wire_metrics(self.metrics)
946
+ if self.category is not None:
947
+ out["category"] = self.category
948
+ return out
949
+
950
+
951
+ @dataclass
952
+ class CellResult:
953
+ name: str
954
+ command: dict[str, float]
955
+ episodes: list[EpisodeResult] = field(default_factory=list)
956
+ video_path: Path | None = None
957
+ traj_path: Path | None = None
958
+ kind: str = "velocity"
959
+ fixed_base: bool = False
960
+ #: C2. `sim2sim` aggregates tracking the harness's way.
961
+ protocol: str = PROTOCOL_M8
962
+ #: C2. `{view: mp4}` — side / front / feet — when the three-view recorder ran.
963
+ videos: dict[str, Path] = field(default_factory=dict)
964
+ #: C3. The cell's rule (`posture` logs its posture medians).
965
+ score: str | None = None
966
+ #: L1. The poster (`video: poster`) and when in seed 0's episode it was taken.
967
+ poster_path: Path | None = None
968
+ poster_t_s: float | None = None
969
+
970
+ @property
971
+ def metrics(self) -> dict[str, float]:
972
+ """Worst case over the seeds: a gate is about the seed that went wrong.
973
+
974
+ Errors take the max; `limit_violations` sums (it is a count of steps).
975
+ Only M17b's three target metrics aggregate this way; the recovery and
976
+ harness metrics live per episode and the hub aggregates those.
977
+ """
978
+ out: dict[str, float] = {}
979
+ for episode in self.episodes:
980
+ for key, value in episode.metrics.items():
981
+ if key not in TARGET_METRICS:
982
+ continue
983
+ if key == "limit_violations":
984
+ out[key] = out.get(key, 0.0) + value
985
+ else:
986
+ out[key] = max(out.get(key, value), value)
987
+ return out
988
+
989
+ @property
990
+ def seeds(self) -> int:
991
+ return len(self.episodes)
992
+
993
+ @property
994
+ def alive(self) -> int:
995
+ return sum(1 for e in self.episodes if e.alive)
996
+
997
+ @property
998
+ def tracking(self) -> dict[str, float]:
999
+ """Per commanded axis over the seeds: M8's mean over every seed, or —
1000
+ under the `sim2sim` protocol — `eval_c_matrix.cell`'s median over the
1001
+ seeds that stayed alive (the aggregate the imported scorecards were
1002
+ converted from)."""
1003
+ axes = [a for a in ("vx", "vy", "yaw") if abs(self.command.get(a, 0.0)) > 1e-9]
1004
+ out: dict[str, float] = {}
1005
+ for axis in axes:
1006
+ if self.protocol == PROTOCOL_SIM2SIM:
1007
+ values = sorted(
1008
+ e.tracking[axis] for e in self.episodes if e.alive and axis in e.tracking
1009
+ )
1010
+ if values:
1011
+ out[axis] = float(values[len(values) // 2])
1012
+ continue
1013
+ values = [e.tracking[axis] for e in self.episodes if axis in e.tracking]
1014
+ if values:
1015
+ out[axis] = float(np.mean(values))
1016
+ return out
1017
+
1018
+ def as_measurement(self) -> dict[str, Any]:
1019
+ out: dict[str, Any] = {
1020
+ "name": self.name,
1021
+ "command": dict(self.command),
1022
+ "seeds": self.seeds,
1023
+ "alive": self.alive,
1024
+ "tracking": self.tracking,
1025
+ "episodes": [e.as_measurement() for e in self.episodes],
1026
+ "kind": self.kind,
1027
+ "fixed_base": self.fixed_base,
1028
+ }
1029
+ if self.metrics:
1030
+ out["metrics"] = _wire_metrics(self.metrics)
1031
+ return out
1032
+
1033
+
1034
+ #: M17b's three per-cell target metrics.
1035
+ TARGET_METRICS = frozenset({"settle_error", "tracking_error", "limit_violations"})
1036
+
1037
+
1038
+ def resolve_targets(sim: Simulator, cell: dict[str, Any]) -> tuple[str, Any, Any, list[str]]:
1039
+ """`(kind, joint_target, ee_target, notes)` for one battery cell.
1040
+
1041
+ Targets are **offsets**, which is what makes a battery portable across
1042
+ robots: `joint_offsets` is radians from the default pose, `ee_offset` is
1043
+ metres from where the end effector sits AT the default pose. A joint
1044
+ target outside the soft limits is clamped to them and the clamp is
1045
+ reported — scoring a policy against a pose the robot may not take would
1046
+ measure the battery, not the policy.
1047
+ """
1048
+ kind = str(cell.get("kind") or "velocity")
1049
+ if kind not in CELL_KINDS:
1050
+ raise ContractMismatch(
1051
+ f"battery cell `{cell.get('name')}` has kind `{kind}`; this engine runs "
1052
+ f"{', '.join(CELL_KINDS)}"
1053
+ )
1054
+ notes: list[str] = []
1055
+ joint_target = ee_target = None
1056
+ if kind == "joint_target":
1057
+ offsets = cell.get("joint_offsets") or {}
1058
+ strangers = sorted(set(offsets) - set(sim.contract.joint_order))
1059
+ if strangers:
1060
+ raise ContractMismatch(
1061
+ f"cell `{cell.get('name')}`: joint_offsets names joints outside the contract: "
1062
+ + ", ".join(strangers)
1063
+ )
1064
+ wanted = sim.default_pose + np.asarray(
1065
+ [float(offsets.get(j, 0.0)) for j in sim.contract.joint_order]
1066
+ )
1067
+ joint_target = np.clip(wanted, sim.soft_lo, sim.soft_hi)
1068
+ moved = [
1069
+ sim.contract.joint_order[k]
1070
+ for k in range(len(wanted))
1071
+ if abs(wanted[k] - joint_target[k]) > 1e-9
1072
+ ]
1073
+ if moved:
1074
+ notes.append(
1075
+ f"{cell.get('name')}: target clamped to the soft limits on {', '.join(moved)}"
1076
+ )
1077
+ elif kind == "reach":
1078
+ if sim.ee_name is None:
1079
+ raise ContractMismatch(
1080
+ f"cell `{cell.get('name')}` is a `reach` cell, but the robot's interface "
1081
+ "declares no end effector with a site or body in the MJCF"
1082
+ )
1083
+ sim._reset_state()
1084
+ home = sim.ee_position()
1085
+ assert home is not None
1086
+ offset = cell.get("ee_offset") or [0.0, 0.0, 0.0]
1087
+ if len(offset) != 3:
1088
+ raise ContractMismatch(f"cell `{cell.get('name')}`: ee_offset must be [dx, dy, dz]")
1089
+ ee_target = home + np.asarray([float(v) for v in offset])
1090
+ return kind, joint_target, ee_target, notes
1091
+
1092
+
1093
+ class _RecoveryTally:
1094
+ """sim2sim's `--init-fallen` bookkeeping, step for step (C2).
1095
+
1096
+ Success: tilt < 15°, base height > 0.85 x standing height, both feet
1097
+ pressing the ground (> 5 N) and |ω| < 1 rad/s, held for 0.5 s — the first
1098
+ step of the first such window is the stand-up time. `end_good` is the last
1099
+ step's verdict (sim2sim); `end_hold` is whether the whole last 0.5 s held
1100
+ (`accept_recovery.py`'s re-fall reading, which does not trust one frame).
1101
+ Torque demand is the PD's BEFORE the effort clip, skipping the first
1102
+ 0.5 s (the drop onto the floor is not the stand-up's demand), read twice
1103
+ as sim2sim reads it: the peak over every physics substep (`tau_frac`)
1104
+ and the peak of the one sample at each control-step boundary
1105
+ (`tau_frac_ctrl`) — which is what `accept_recovery.py`'s ⑤ reads
1106
+ (Isaac's `computed_torque` at 50 Hz), and so what the hub gates on.
1107
+ """
1108
+
1109
+ def __init__(self, sim: Simulator) -> None:
1110
+ self.sim = sim
1111
+ self.hold = max(1, round(RECOVERY_HOLD_S * sim.control_rate_hz))
1112
+ self.settle = max(0, round(RECOVERY_TORQUE_SETTLE_S * sim.control_rate_hz))
1113
+ self.h_gate = RECOVERY_HEIGHT_FRACTION * (sim.stand_height or sim.init_height)
1114
+ self.run = 0
1115
+ self.first = -1
1116
+ self.good: list[bool] = []
1117
+ n = len(sim.qpos_adr)
1118
+ self.tau_peak = np.zeros(n)
1119
+ self.tau_peak_ctrl = np.zeros(n)
1120
+ self.tau_sat = np.zeros(n)
1121
+ self.tau_samples = 0
1122
+ self.act_sat = np.zeros(n)
1123
+ self.act_samples = 0
1124
+ self.step_index = 0
1125
+
1126
+ def probe(self, demand: np.ndarray, last_substep: bool = False) -> None:
1127
+ if self.step_index < self.settle:
1128
+ return
1129
+ magnitude = np.abs(demand)
1130
+ np.maximum(self.tau_peak, magnitude, out=self.tau_peak)
1131
+ if last_substep:
1132
+ np.maximum(self.tau_peak_ctrl, magnitude, out=self.tau_peak_ctrl)
1133
+ self.tau_sat += magnitude > self.sim.tau_limit
1134
+ self.tau_samples += 1
1135
+
1136
+ def after_step(self, index: int, raw_action: np.ndarray) -> None:
1137
+ sim = self.sim
1138
+ feet = sim.foot_forces()
1139
+ _lin, ang = sim.base_velocity_body()
1140
+ ok = (
1141
+ sim.tilt_deg() < RECOVERY_TILT_DEG
1142
+ and float(sim.base_position()[2]) > self.h_gate
1143
+ and len(feet) >= 2
1144
+ and all(f > RECOVERY_FOOT_N for f in feet)
1145
+ and float(np.linalg.norm(ang)) < RECOVERY_OMEGA
1146
+ )
1147
+ self.good.append(ok)
1148
+ self.run = self.run + 1 if ok else 0
1149
+ if self.run >= self.hold and self.first < 0:
1150
+ self.first = index - self.run + 1
1151
+ self.act_sat += np.abs(raw_action) >= RECOVERY_ACTION_SAT
1152
+ self.act_samples += 1
1153
+ self.step_index = index + 1
1154
+
1155
+ def metrics(self) -> dict[str, float]:
1156
+ rate = self.sim.control_rate_hz
1157
+ success = self.first >= 0
1158
+ end_hold = len(self.good) >= self.hold and all(self.good[-self.hold :])
1159
+ out: dict[str, float] = {
1160
+ "recover_success": 1.0 if success else 0.0,
1161
+ "recover_time_s": (self.first / rate) if success else float("inf"),
1162
+ "end_good": 1.0 if (self.good and self.good[-1]) else 0.0,
1163
+ "end_hold": 1.0 if end_hold else 0.0,
1164
+ "refall": 1.0 if (success and not end_hold) else 0.0,
1165
+ }
1166
+ limit = np.where(np.isfinite(self.sim.tau_limit), self.sim.tau_limit, np.nan)
1167
+ with np.errstate(invalid="ignore", divide="ignore"):
1168
+ frac = self.tau_peak / limit
1169
+ frac_ctrl = self.tau_peak_ctrl / limit
1170
+ for name, value, sampled in zip(
1171
+ self.sim.contract.joint_order, frac, frac_ctrl, strict=True
1172
+ ):
1173
+ out[f"tau_frac.{name}"] = float(value) if np.isfinite(value) else float("inf")
1174
+ out[f"tau_frac_ctrl.{name}"] = float(sampled) if np.isfinite(sampled) else float("inf")
1175
+ if self.act_samples:
1176
+ for name, value in zip(
1177
+ self.sim.contract.joint_order, self.act_sat / self.act_samples, strict=True
1178
+ ):
1179
+ out[f"sat.{name}"] = float(value)
1180
+ return out
1181
+
1182
+
1183
+ def run_episode(
1184
+ sim: Simulator,
1185
+ policy: Policy,
1186
+ *,
1187
+ command: dict[str, float],
1188
+ seed: int,
1189
+ horizon_s: float,
1190
+ record: bool = False,
1191
+ frame_sink: FrameSink | None = None,
1192
+ cell: dict[str, Any] | None = None,
1193
+ protocol: str = PROTOCOL_M8,
1194
+ category: str | None = None,
1195
+ ) -> EpisodeResult:
1196
+ """One episode: reset, run the contract's loop, measure.
1197
+
1198
+ A `velocity` cell on a floating base measures survival and per-axis
1199
+ tracking, as M8 did. Everything else — `joint_target`, `reach`, and ANY
1200
+ cell on a fixed base — also measures an error against a target (for a
1201
+ velocity cell on a fixed base that target is the default pose: "hold")
1202
+ plus soft-limit violations, because on a fixed base `upright()` is true by
1203
+ construction and proves nothing.
1204
+
1205
+ C2: the cell's conditions (`sim.configure` was called by the battery)
1206
+ shape the loop the way `sim2sim.py` shapes it — heading → observe → act →
1207
+ power scale → target (feed-forward / β anchor) → bridge slew → delay →
1208
+ push → PD substeps — and the `sim2sim` protocol measures survival and
1209
+ tracking with the harness's own definitions. A `recovery` cell starts
1210
+ fallen (`category`), never ends early, and measures the stand-up.
1211
+ """
1212
+ kind, joint_target, ee_target, _notes = resolve_targets(sim, cell or {})
1213
+ cond = sim.cond
1214
+ recovery = kind == "recovery"
1215
+ fallen = sim.reset(seed, protocol=protocol, category=category if recovery else None)
1216
+ cmd = np.array([command.get("vx", 0.0), command.get("vy", 0.0), command.get("yaw", 0.0)])
1217
+ last_action = np.zeros(sim.contract.action_dim)
1218
+ steps = max(1, math.ceil(horizon_s * sim.control_rate_hz))
1219
+ settle = min(SETTLE_S, horizon_s / 4.0)
1220
+ settle_from = horizon_s * (1.0 - SETTLE_WINDOW_FRACTION)
1221
+ measure_error = kind in ("joint_target", "reach") or (kind == "velocity" and sim.fixed_base)
1222
+ if kind == "velocity" and sim.fixed_base:
1223
+ # "Stand" on a fixed base means "hold the default pose". It is scored
1224
+ # that way, so a policy that observes `joint_target` must be shown
1225
+ # that same target — not nothing.
1226
+ joint_target = sim.default_pose.copy()
1227
+ hold_target = joint_target
1228
+ harness_mode = protocol == PROTOCOL_SIM2SIM
1229
+ # C3: posture — tilt, mirrored-pair asymmetry, drift — on a `posture` cell,
1230
+ # and on every sim2sim episode on a floating base (whose tilt it already read).
1231
+ track_posture = (
1232
+ not sim.fixed_base
1233
+ and not recovery
1234
+ and kind == "velocity"
1235
+ and (harness_mode or (cell or {}).get("score") == "posture")
1236
+ )
1237
+ pairs = sim.mirror_pairs() if track_posture else []
1238
+ pair_left = np.asarray([p[0] for p in pairs], dtype=np.intp)
1239
+ pair_right = np.asarray([p[1] for p in pairs], dtype=np.intp)
1240
+ pair_sign = np.asarray([p[2] for p in pairs], dtype=np.float64)
1241
+ asym_sum = np.zeros(len(pairs))
1242
+ asym_samples = 0
1243
+ start_q = np.asarray(sim.data.qpos[sim.qpos_adr], dtype=np.float64).copy()
1244
+ delay = conditions_mod.DelayLine(start_q, delay=cond.delay, jitter=cond.delay_jitter, seed=seed)
1245
+ slew = conditions_mod.Slew(start_q, sim.velocity_limit, sim.control_dt) if cond.slew else None
1246
+ yaw0 = sim.yaw()
1247
+ tally = _RecoveryTally(sim) if recovery else None
1248
+ fall_height = HARNESS_FALL_FRACTION * sim.stand_height
1249
+ sums = np.zeros(3)
1250
+ samples = 0
1251
+ err_track: list[float] = []
1252
+ err_settle: list[float] = []
1253
+ violations = 0
1254
+ fell_at: float | None = None
1255
+ # sim2sim's arrays (only what the gate and the parity table read).
1256
+ h_vels: list[tuple[float, float]] = []
1257
+ h_yaws: list[float] = [yaw0]
1258
+ h_heights: list[float] = []
1259
+ h_tilts: list[float] = []
1260
+ h_raw: list[np.ndarray] = []
1261
+ # C3b: the base's horizontal speed per control step, and the step a push
1262
+ # landed on — for the recovery time after it.
1263
+ h_speed: list[float] = []
1264
+ push_index: int | None = None
1265
+ xy0 = sim.base_position()[:2].copy() if not sim.fixed_base else np.zeros(2)
1266
+ traj: dict[str, list[Any]] | None = (
1267
+ {"t": [], "joint_pos": [], "action": [], "target": [], "base_pos": [], "base_vel": []}
1268
+ if record
1269
+ else None
1270
+ )
1271
+ if traj is not None and measure_error:
1272
+ traj["error"] = []
1273
+ if traj is not None and recovery:
1274
+ traj["tilt_deg"] = []
1275
+ if frame_sink is not None:
1276
+ frame_sink(sim.data, 0.0)
1277
+ for index in range(steps):
1278
+ if cond.heading and not sim.fixed_base:
1279
+ error = conditions_mod.wrap(yaw0 - sim.yaw())
1280
+ cmd[2] = float(
1281
+ np.clip(
1282
+ cond.heading_k * error,
1283
+ -conditions_mod.HEADING_MAX_RATE,
1284
+ conditions_mod.HEADING_MAX_RATE,
1285
+ )
1286
+ )
1287
+ obs = sim.observe(cmd, last_action, joint_target=joint_target, ee_target=ee_target)
1288
+ raw = policy(obs)
1289
+ clip = sim.contract.action_clip
1290
+ acted = np.clip(raw, -clip, clip) * cond.power_scale
1291
+ target = sim.command_for(acted, command=cmd)
1292
+ if slew is not None:
1293
+ target = slew(target)
1294
+ target = delay(target)
1295
+ if cond.push is not None and abs(index / sim.control_rate_hz - cond.push.at_s) < (
1296
+ 0.5 / sim.control_rate_hz
1297
+ ):
1298
+ sim.push(cond.push.impulse, cond.push.axis)
1299
+ push_index = index
1300
+ sim.step(acted, target=target, torque_probe=tally.probe if tally else None)
1301
+ last_action = acted
1302
+ t = sim.time_s
1303
+ lin, ang = sim.base_velocity_body()
1304
+ if t >= settle:
1305
+ sums += np.array([lin[0], lin[1], ang[2]])
1306
+ samples += 1
1307
+ if harness_mode or tally is not None or track_posture:
1308
+ h_vels.append((float(lin[0]), float(lin[1])))
1309
+ h_yaws.append(sim.yaw())
1310
+ h_heights.append(float(sim.base_position()[2]))
1311
+ h_tilts.append(sim.tilt_deg())
1312
+ h_raw.append(np.asarray(raw, dtype=np.float64))
1313
+ if track_posture and sim.base_dof is not None:
1314
+ v = sim.data.qvel[sim.base_dof : sim.base_dof + 2]
1315
+ h_speed.append(float(math.hypot(float(v[0]), float(v[1]))))
1316
+ if pairs and t >= settle:
1317
+ q_now = np.asarray(sim.data.qpos[sim.qpos_adr], dtype=np.float64)
1318
+ asym_sum += q_now[pair_left] - pair_sign * q_now[pair_right]
1319
+ asym_samples += 1
1320
+ if tally is not None:
1321
+ tally.after_step(index, np.asarray(raw, dtype=np.float64))
1322
+ error_value: float | None = None
1323
+ if measure_error:
1324
+ if kind == "reach":
1325
+ ee = sim.ee_position()
1326
+ assert ee is not None and ee_target is not None
1327
+ error_value = float(np.linalg.norm(ee - ee_target))
1328
+ else:
1329
+ assert hold_target is not None
1330
+ # The WORST joint, not an RMS: one joint far off must not hide
1331
+ # behind five that are close.
1332
+ error_value = float(np.max(np.abs(sim.data.qpos[sim.qpos_adr] - hold_target)))
1333
+ if t >= settle:
1334
+ err_track.append(error_value)
1335
+ if t >= settle_from:
1336
+ err_settle.append(error_value)
1337
+ if sim.limit_violation():
1338
+ violations += 1
1339
+ if traj is not None:
1340
+ traj["t"].append(round(t, 4))
1341
+ traj["joint_pos"].append([round(float(v), 4) for v in sim.data.qpos[sim.qpos_adr]])
1342
+ traj["action"].append([round(float(v), 4) for v in last_action])
1343
+ traj["target"].append([round(float(v), 4) for v in target])
1344
+ traj["base_pos"].append([round(float(v), 4) for v in sim.base_position()])
1345
+ traj["base_vel"].append(
1346
+ [round(float(lin[0]), 4), round(float(lin[1]), 4), round(float(ang[2]), 4)]
1347
+ )
1348
+ if error_value is not None:
1349
+ traj["error"].append(round(error_value, 5))
1350
+ if recovery:
1351
+ traj["tilt_deg"].append(round(h_tilts[-1], 2))
1352
+ if frame_sink is not None:
1353
+ frame_sink(sim.data, t)
1354
+ if recovery:
1355
+ continue # lying down is a legal state; getting back up is the task
1356
+ if harness_mode and not sim.fixed_base:
1357
+ if h_heights[-1] < fall_height:
1358
+ fell_at = round(t, 3)
1359
+ break
1360
+ elif not sim.upright():
1361
+ fell_at = round(t, 3)
1362
+ break
1363
+ tracking: dict[str, float] = {}
1364
+ metrics: dict[str, float] = {}
1365
+ alive = fell_at is None
1366
+ if harness_mode and not sim.fixed_base and not recovery:
1367
+ # sim2sim's own readings: alive = the last height and tilt, body-frame
1368
+ # vx/vy averaged after the first second, and the yaw rate from the
1369
+ # net world-yaw change (the body-frame rate is biased by the sway).
1370
+ rate = sim.control_rate_hz
1371
+ duration = len(h_heights) / rate if h_heights else 0.0
1372
+ alive = bool(
1373
+ h_heights
1374
+ and h_heights[-1] > HARNESS_ALIVE_FRACTION * sim.stand_height
1375
+ and h_tilts[-1] < HARNESS_ALIVE_TILT_DEG
1376
+ )
1377
+ if not alive and fell_at is None:
1378
+ fell_at = round(duration, 3)
1379
+ window = h_vels[int(rate) :] or h_vels
1380
+ vx_mean = float(np.mean([v[0] for v in window])) if window else 0.0
1381
+ vy_mean = float(np.mean([v[1] for v in window])) if window else 0.0
1382
+ dyaw = float(np.unwrap(np.asarray(h_yaws))[-1] - h_yaws[0]) if len(h_yaws) > 1 else 0.0
1383
+ yaw_rate = dyaw / duration if duration else 0.0
1384
+ measured = {"vx": vx_mean, "vy": vy_mean, "yaw": yaw_rate}
1385
+ for axis in ("vx", "vy", "yaw"):
1386
+ if abs(cmd_of(command, axis)) > 1e-9:
1387
+ tracking[axis] = measured[axis] / cmd_of(command, axis)
1388
+ dxy = sim.base_position()[:2] - xy0
1389
+ metrics = {
1390
+ "vx_mean": vx_mean,
1391
+ "vy_mean": vy_mean,
1392
+ "dyaw_deg": math.degrees(dyaw),
1393
+ "dx": float(dxy[0]),
1394
+ "dy": float(dxy[1]),
1395
+ "tilt_max_deg": float(max(h_tilts)) if h_tilts else float("inf"),
1396
+ "t": duration,
1397
+ }
1398
+ if h_raw:
1399
+ tail = np.abs(np.asarray(h_raw[len(h_raw) // 3 :]))
1400
+ if len(tail):
1401
+ metrics["sat_max_pct"] = float(
1402
+ ((tail >= sim.contract.action_clip - 1e-6).mean(axis=0) * 100).max()
1403
+ )
1404
+ elif not recovery:
1405
+ # M8: the mean body-frame velocity after the settle, per commanded axis
1406
+ # of the CELL's command (a heading loop may rewrite the yaw it sends).
1407
+ measured_v = sums / samples if samples else np.zeros(3)
1408
+ for index_axis, axis in enumerate(("vx", "vy", "yaw")):
1409
+ commanded = cmd_of(command, axis)
1410
+ if abs(commanded) > 1e-9:
1411
+ tracking[axis] = float(measured_v[index_axis] / commanded) if samples else 0.0
1412
+ if track_posture:
1413
+ # C3 (lucen_cli.rollout.posture): the worst pair's settled asymmetry,
1414
+ # the episode's maximum tilt, and how far the base wandered. A robot
1415
+ # with no mirrored pair reports `asym_pairs: 0` and no asymmetry; an
1416
+ # episode that ended before the settle reports inf, never a zero.
1417
+ dxy_end = sim.base_position()[:2] - xy0
1418
+ metrics["tilt_max_deg"] = float(max(h_tilts)) if h_tilts else float("inf")
1419
+ metrics["drift_m"] = float(math.hypot(float(dxy_end[0]), float(dxy_end[1])))
1420
+ metrics["asym_pairs"] = float(len(pairs))
1421
+ if pairs:
1422
+ metrics["asym_max_deg"] = (
1423
+ float(posture_mod.pair_asymmetry_deg(asym_sum / asym_samples).max())
1424
+ if asym_samples
1425
+ else float("inf")
1426
+ )
1427
+ if cond.push is not None:
1428
+ # C3b: how long the base took to settle after the push. A seed
1429
+ # that fell, or never settled, reports the whole window that was
1430
+ # left after the push — a number at least as bad as any bar, never
1431
+ # a null that a median would quietly leave out.
1432
+ window_s = max(0.0, horizon_s - cond.push.at_s)
1433
+ metrics["recover_s"] = push_recovery_s(
1434
+ h_tilts,
1435
+ h_speed,
1436
+ push_index=push_index,
1437
+ rate_hz=sim.control_rate_hz,
1438
+ fell=fell_at is not None,
1439
+ window_s=window_s,
1440
+ )
1441
+ metrics["recovered"] = 1.0 if metrics["recover_s"] < window_s else 0.0
1442
+ if measure_error:
1443
+ # No sample in a window means the episode ended before it: report inf,
1444
+ # never a flattering zero.
1445
+ metrics.update(
1446
+ {
1447
+ "settle_error": float(np.mean(err_settle)) if err_settle else float("inf"),
1448
+ "tracking_error": float(np.mean(err_track)) if err_track else float("inf"),
1449
+ "limit_violations": float(violations),
1450
+ }
1451
+ )
1452
+ if tally is not None:
1453
+ metrics = tally.metrics()
1454
+ alive = bool(metrics["recover_success"]) and not bool(metrics["refall"])
1455
+ fell_at = None
1456
+ trajectory: dict[str, Any] | None = None
1457
+ if traj is not None:
1458
+ trajectory = {
1459
+ "command": {k: float(command.get(k, 0.0)) for k in ("vx", "vy", "yaw")},
1460
+ "kind": kind,
1461
+ "seed": seed,
1462
+ "rate_hz": sim.control_rate_hz,
1463
+ "joints": list(sim.contract.joint_order),
1464
+ "alive": alive,
1465
+ "fell_at_s": fell_at,
1466
+ "tracking": tracking,
1467
+ "metrics": {k: (v if math.isfinite(v) else None) for k, v in metrics.items()},
1468
+ "joint_target": None if joint_target is None else [float(v) for v in joint_target],
1469
+ "ee_target": None if ee_target is None else [float(v) for v in ee_target],
1470
+ "conditions": cond.label() or None,
1471
+ "category": fallen.category if fallen is not None else None,
1472
+ "start": fallen.describe() if fallen is not None else None,
1473
+ **traj,
1474
+ }
1475
+ return EpisodeResult(
1476
+ seed=seed,
1477
+ alive=alive,
1478
+ fell_at_s=fell_at,
1479
+ tracking=tracking,
1480
+ steps=len(traj["t"]) if traj is not None else steps,
1481
+ trajectory=trajectory,
1482
+ metrics=metrics,
1483
+ category=fallen.category if fallen is not None else None,
1484
+ )
1485
+
1486
+
1487
+ def push_recovery_s(
1488
+ tilts: list[float],
1489
+ speeds: list[float],
1490
+ *,
1491
+ push_index: int | None,
1492
+ rate_hz: float,
1493
+ fell: bool,
1494
+ window_s: float,
1495
+ tilt_tol: float = PUSH_RECOVER_TILT_DEG,
1496
+ speed_tol: float = PUSH_RECOVER_SPEED_M_S,
1497
+ ) -> float:
1498
+ """Seconds from a push until the base stays settled to the end of the episode (C3b).
1499
+
1500
+ Settled = tilted at most `tilt_tol` and moving at most `speed_tol` across
1501
+ the floor, from that control step to the last one. `window_s` (the time
1502
+ the episode had left after the push) when the robot fell, the push never
1503
+ landed, or the base was still unsettled at the end.
1504
+ """
1505
+ if fell or push_index is None or push_index >= len(tilts) or len(speeds) != len(tilts):
1506
+ return float(window_s)
1507
+ after = range(push_index, len(tilts))
1508
+ last_bad: int | None = None
1509
+ for k in after:
1510
+ if tilts[k] > tilt_tol or speeds[k] > speed_tol:
1511
+ last_bad = k
1512
+ if last_bad is None:
1513
+ return 0.0
1514
+ if last_bad == len(tilts) - 1:
1515
+ return float(window_s)
1516
+ return min(float(window_s), (last_bad + 1 - push_index) / rate_hz)
1517
+
1518
+
1519
+ def cmd_of(command: dict[str, float], axis: str) -> float:
1520
+ return float(command.get(axis, 0.0) or 0.0)
1521
+
1522
+
1523
+ def build(
1524
+ mjcf_path: Path,
1525
+ policy: Policy,
1526
+ *,
1527
+ card: dict[str, Any] | None,
1528
+ ) -> tuple[Simulator, IoContract]:
1529
+ """Compile the robot, read the contract, and refuse a mismatch before stepping.
1530
+
1531
+ When the card carries an `interface` (M17b) the model is compiled by the
1532
+ shared plant builder — the robot's own joint order, per-group gains,
1533
+ default pose and rates — and `validate` checks the contract against that
1534
+ declaration, fingerprint first. When it does not, this is M8's path,
1535
+ unchanged: the MJCF as written, one scalar kp/kd, the joint set inferred
1536
+ from the keys of `card.limits`.
1537
+ """
1538
+ contract = IoContract.from_manifest(policy.manifest)
1539
+ try:
1540
+ emb = Embodiment.from_card(card)
1541
+ except InterfaceError as exc:
1542
+ raise ContractMismatch(f"the robot card's interface is unreadable: {exc}") from exc
1543
+ built: Plant | None = None
1544
+ if emb is not None:
1545
+ try:
1546
+ built = build_plant(mjcf_path, emb)
1547
+ except PlantError as exc:
1548
+ raise ContractMismatch(str(exc)) from exc
1549
+ model = built.model
1550
+ else:
1551
+ model = load_model(mjcf_path)
1552
+ validate(
1553
+ contract,
1554
+ model_joints=model_joints(model),
1555
+ policy_in_dim=policy.in_dim,
1556
+ policy_out_dim=policy.out_dim,
1557
+ card=card,
1558
+ )
1559
+ plant = policy.manifest.get("plant") if isinstance(policy.manifest.get("plant"), dict) else None
1560
+ return Simulator(model, contract, plant=plant, card=card, emb=emb, built=built), contract
1561
+
1562
+
1563
+ #: The target a cell kind supplies, as the observation block that consumes it.
1564
+ _TARGET_BLOCK = {"joint_target": "joint_target", "reach": "ee_target"}
1565
+
1566
+
1567
+ def check_battery(sim: Simulator, spec: dict[str, Any]) -> None:
1568
+ """Refuse a battery this policy cannot sit — before a single physics step.
1569
+
1570
+ A policy is trained to follow one kind of target. One whose contract has
1571
+ no `ee_target` block has no input through which a `reach` cell's target
1572
+ could reach it; running it anyway would score the robot's default
1573
+ behaviour and call the number a reach error. Every cell is checked first,
1574
+ so a bad cell five deep does not waste four. C2: every cell's conditions
1575
+ are parsed here too, and a recovery cell needs a floating base and feet.
1576
+ """
1577
+ blocks = {
1578
+ obs_registry.resolve(b.name).name
1579
+ for b in sim.contract.obs_spec
1580
+ if obs_registry.is_known(b.name)
1581
+ }
1582
+ problems: list[str] = []
1583
+ protocol = str(spec.get("protocol") or PROTOCOL_M8)
1584
+ if protocol not in PROTOCOLS:
1585
+ problems.append(f"protocol `{protocol}` (this engine runs {', '.join(PROTOCOLS)})")
1586
+ for cell in spec.get("cells") or []:
1587
+ name = str(cell.get("name"))
1588
+ kind = str(cell.get("kind") or "velocity")
1589
+ if kind not in CELL_KINDS:
1590
+ problems.append(
1591
+ f"`{name}` has kind `{kind}` (this engine runs {', '.join(CELL_KINDS)})"
1592
+ )
1593
+ continue
1594
+ try:
1595
+ cond = conditions_mod.parse(
1596
+ conditions_mod.merge(spec.get("conditions"), cell.get("conditions"))
1597
+ )
1598
+ except conditions_mod.ConditionError as exc:
1599
+ problems.append(f"`{name}`: {exc}")
1600
+ continue
1601
+ if kind == "recovery":
1602
+ if sim.fixed_base:
1603
+ problems.append(f"`{name}` is a recovery cell, but the robot's base is fixed")
1604
+ elif not sim._foot_geoms:
1605
+ problems.append(
1606
+ f"`{name}` is a recovery cell, scored on both feet on the ground, and the "
1607
+ "robot's interface declares no `feet`"
1608
+ )
1609
+ if cond.fallen is None:
1610
+ problems.append(f"`{name}` is a recovery cell with no `init_pose` condition")
1611
+ elif cond.fallen is not None:
1612
+ problems.append(f"`{name}` has an `init_pose` but is not a `recovery` cell")
1613
+ needed = _TARGET_BLOCK.get(kind)
1614
+ if needed is not None and needed not in blocks:
1615
+ problems.append(
1616
+ f"`{name}` is a `{kind}` cell, but the policy observes no `{needed}` block — "
1617
+ "it has no input for that target"
1618
+ )
1619
+ if kind == "velocity" and not sim.fixed_base and "commands" not in blocks:
1620
+ moving = any(abs(float(cell.get(a) or 0.0)) > 1e-9 for a in ("vx", "vy", "yaw"))
1621
+ if moving:
1622
+ problems.append(
1623
+ f"`{name}` commands a velocity, but the policy observes no `commands` block"
1624
+ )
1625
+ if kind == "reach" and sim.ee_name is None:
1626
+ problems.append(
1627
+ f"`{name}` is a `reach` cell, but the robot's interface declares no end "
1628
+ "effector with a site or body in the MJCF"
1629
+ )
1630
+ if problems:
1631
+ raise ContractMismatch(
1632
+ "battery mismatch — refusing to run cells this policy cannot sit: "
1633
+ + "; ".join(problems)
1634
+ + "."
1635
+ )
1636
+
1637
+
1638
+ def run_battery(
1639
+ sim: Simulator,
1640
+ policy: Policy,
1641
+ spec: dict[str, Any],
1642
+ *,
1643
+ seed_base: int,
1644
+ out_dir: Path,
1645
+ recorder: Any | None = None,
1646
+ log: Log = lambda _line: None,
1647
+ ) -> tuple[list[CellResult], list[str]]:
1648
+ """Every cell x seed of a battery. Seed 0 of each cell is recorded and rendered."""
1649
+ import json
1650
+
1651
+ warnings: list[str] = []
1652
+ check_battery(sim, spec)
1653
+ horizon = float(spec.get("horizon_s") or 5.0)
1654
+ seeds = int(spec.get("seeds") or 1)
1655
+ protocol = str(spec.get("protocol") or PROTOCOL_M8)
1656
+ results: list[CellResult] = []
1657
+ for cell in spec["cells"]:
1658
+ name = str(cell["name"])
1659
+ command = {k: float(cell.get(k) or 0.0) for k in ("vx", "vy", "yaw")}
1660
+ cond = conditions_mod.parse(
1661
+ conditions_mod.merge(spec.get("conditions"), cell.get("conditions"))
1662
+ )
1663
+ for note in sim.configure(cond):
1664
+ if note not in warnings:
1665
+ warnings.append(note)
1666
+ kind, _jt, _ee, notes = resolve_targets(sim, cell)
1667
+ warnings.extend(notes)
1668
+ result = CellResult(
1669
+ name=name,
1670
+ command=command,
1671
+ kind=kind,
1672
+ fixed_base=sim.fixed_base,
1673
+ protocol=protocol,
1674
+ score=cell.get("score"),
1675
+ )
1676
+ cell_dir = out_dir / "cells" / _slug(name)
1677
+ cell_dir.mkdir(parents=True, exist_ok=True)
1678
+ started = time.perf_counter()
1679
+ for k in range(seeds):
1680
+ seed = seed_base + k
1681
+ category = (
1682
+ cond.fallen.category_for(k, seeds)
1683
+ if cond.fallen is not None and kind == "recovery"
1684
+ else None
1685
+ )
1686
+ sink: FrameSink | None = None
1687
+ recording: Any | None = None
1688
+ if k == 0 and recorder is not None and recorder.available:
1689
+ recording = recorder
1690
+ recording.start(cell_dir / "video.mp4")
1691
+ sink = recording.frame
1692
+ episode = run_episode(
1693
+ sim,
1694
+ policy,
1695
+ command=command,
1696
+ seed=seed,
1697
+ horizon_s=horizon,
1698
+ record=(k == 0),
1699
+ frame_sink=sink,
1700
+ cell=cell,
1701
+ protocol=protocol,
1702
+ category=category,
1703
+ )
1704
+ if recording is not None:
1705
+ finished = recording.finish()
1706
+ if getattr(recording, "poster", False):
1707
+ # L1: one still, not a video.
1708
+ if finished is not None:
1709
+ result.poster_path, result.poster_t_s = finished
1710
+ elif isinstance(finished, dict):
1711
+ result.videos = dict(finished)
1712
+ result.video_path = finished.get("side") or next(iter(finished.values()), None)
1713
+ elif finished is not None:
1714
+ result.video_path = finished
1715
+ if not finished and recording.reason:
1716
+ warnings.append(f"{name}: {recording.reason}")
1717
+ if episode.trajectory is not None:
1718
+ traj_path = cell_dir / "trajectory.json"
1719
+ traj_path.write_text(
1720
+ json.dumps({"cell": name, **episode.trajectory}, separators=(",", ":")),
1721
+ encoding="utf-8",
1722
+ )
1723
+ result.traj_path = traj_path
1724
+ result.episodes.append(episode)
1725
+ log(_cell_line(result, sim, cond, time.perf_counter() - started))
1726
+ results.append(result)
1727
+ sim.configure(conditions_mod.Conditions())
1728
+ if recorder is not None and not recorder.available and recorder.reason:
1729
+ warnings.insert(0, f"video unavailable: {recorder.reason}")
1730
+ return results, warnings
1731
+
1732
+
1733
+ def _cell_line(
1734
+ result: CellResult, sim: Simulator, cond: conditions_mod.Conditions, seconds: float
1735
+ ) -> str:
1736
+ """The one log line per cell. A velocity cell on a floating base with no
1737
+ conditions logs exactly the line M8 shipped — it is quoted in the
1738
+ Quickstart, and "every existing battery is unchanged" includes what it
1739
+ prints. Only the new shapes say more."""
1740
+ kind = result.kind
1741
+ tracking = ", ".join(f"{a} {v * 100:.0f}%" for a, v in sorted(result.tracking.items()))
1742
+ label = cond.label()
1743
+ if kind == "recovery":
1744
+ eps = result.episodes
1745
+ ok = [e for e in eps if e.metrics.get("recover_success")]
1746
+ times = sorted(e.metrics["recover_time_s"] for e in ok)
1747
+ refall = sum(1 for e in ok if e.metrics.get("refall"))
1748
+ median = f"{times[len(times) // 2]:.2f} s" if times else "—"
1749
+ return (
1750
+ f"test {result.name} [recovery]: {len(ok)}/{len(eps)} stood up, median {median}, "
1751
+ f"{refall} re-fell" + (f" · {label}" if label else "") + f" ({seconds:.1f}s)"
1752
+ )
1753
+ if result.score == "posture":
1754
+ medians = _posture_medians(result.episodes)
1755
+ shown = ", ".join(
1756
+ f"{label} {medians[key]:.2f}{unit}"
1757
+ for key, label, unit in (
1758
+ ("tilt_max_deg", "tilt max", "°"),
1759
+ ("asym_max_deg", "asym", "°"),
1760
+ ("drift_m", "drift", " m"),
1761
+ ("recover_s", "recovered in", " s"),
1762
+ )
1763
+ if key in medians
1764
+ )
1765
+ return (
1766
+ f"test {result.name} [posture]: {result.alive}/{result.seeds} upright"
1767
+ + (f", {shown} (median over seeds)" if shown else "")
1768
+ + (f" · {label}" if label else "")
1769
+ + f" ({seconds:.1f}s)"
1770
+ )
1771
+ measured = result.metrics
1772
+ unit = "m" if kind == "reach" else "rad"
1773
+ legacy = kind == "velocity" and not sim.fixed_base
1774
+ head = f"test {result.name}: " if legacy else f"test {result.name} [{kind}]: "
1775
+ return (
1776
+ head
1777
+ + ("fixed base" if sim.fixed_base else f"{result.alive}/{result.seeds} upright")
1778
+ + (f", tracking {tracking}" if tracking else "")
1779
+ + (
1780
+ f", settle error {measured['settle_error']:.4f} {unit}, "
1781
+ f"{int(measured['limit_violations'])} limit violation step(s)"
1782
+ if "settle_error" in measured
1783
+ else ""
1784
+ )
1785
+ + (f" · {label}" if label else "")
1786
+ + f" ({seconds:.1f}s)"
1787
+ )
1788
+
1789
+
1790
+ def _posture_medians(episodes: list[EpisodeResult]) -> dict[str, float]:
1791
+ """The founder's median (`sorted(v)[n // 2]`) of each posture metric over the seeds."""
1792
+ out: dict[str, float] = {}
1793
+ for key in ("tilt_max_deg", "asym_max_deg", "drift_m", "recover_s"):
1794
+ values = sorted(
1795
+ e.metrics[key] for e in episodes if key in e.metrics and math.isfinite(e.metrics[key])
1796
+ )
1797
+ if values:
1798
+ out[key] = float(values[len(values) // 2])
1799
+ return out
1800
+
1801
+
1802
+ def _slug(name: str) -> str:
1803
+ """A cell name as a path segment (`vy-0.10` → `vy-minus-0p10`), lower-case, no punctuation."""
1804
+ swap = {"&": "-and-", "+": "-plus-", "-": "-minus-", ".": "p", "@": "-at-"}
1805
+ text = "".join(swap.get(ch, ch) for ch in name.lower())
1806
+ while "--" in text:
1807
+ text = text.replace("--", "-")
1808
+ return "".join(ch if ch.isalnum() or ch in "-_" else "-" for ch in text).strip("-") or "cell"