"""Stringman world-model data: task-free trajectories in the twin, the SkyJEPA recipe (scripts/arm_datagen.py's, on the cable robot). No target, no cost, no task label in the files: only how the machine moves when commanded. A NEW RANDOM ROOM PER TRAJECTORY (sim/scene.random_room; --fixed-room for a test set of one real room): the model is a model of the machine, not of a place (2026-10-08), and a room enters its dynamics only through the line geometry, which is stored with each trajectory (mounts, runs) and given to the model as an input. python scripts/cdpr_datagen.py --scenario configs/scenarios/clear_playroom.yaml --n 20 --out data/cdpr_v0_sample python scripts/cdpr_datagen.py --scenario configs/scenarios/clear_playroom.yaml --n 4000 --workers 7 --out data/cdpr_v0_playroom Files (shards of 100): states (N, T+1, 14) p, v, swing s, s_dot, wrist, wrist_dot, closure, closure_dot (sim.kin), AS ESTIMATED (noise added) actions (N, T, 5) the command SENT at each record step: gantry velocity (3), wrist rate, closure target tensions (N, T+1, 4) line tensions (N): an observation, NOT a model input (uncalibrated on the real robot) payload (N, T+1) kg held at the grasp point (bookkeeping: the model must infer it) eta (N, P) the trajectory's randomised physics (eta_names) mounts (N, 4, 3) the room's mount points (line order of the firmware), runs (N, 4) the indirect lines' fixed runs Recorded at 10 Hz (dt 0.1 s = half a planner step: the planner's 0.2 s is 2 model steps), 30 s per trajectory. The machine is driven exactly as the real host drives it (sim.host.SpoolState: velocity -> line speeds by finite-difference IK at the ESTIMATED position through MISCALIBRATED mounts, 50 Hz spool loop with the tension floor and soft mute, acceleration-limited motors, 60 Hz gripper loop), with a command latency. Planning model = the machine-only twin (no furniture, no contacts): the grasp point is kept >= 4 cm above the floor (floor contact is a task's business, not the machine's). Motion = a chain of segments chosen at random until the trajectory is full: goto to a random reachable point (inside the mount polygon by 0.35 m, 3-line worst tension <= 26 N), 0.1-0.35 m/s dwell hold still 1-4 s (the robot mostly hovers and waits) stop a fast goto cut to zero half way: swing excitation and its decay pump horizontal sinusoid at 0.7-1.3 x the swing frequency, small amplitude: the swing's resonance descend down to a grasp height (G 4-15 cm), hold, back up (the pick's vertical regime) plan what MPPI asks: a new command every planner step (0.2 s), AR(0.9) random walk, 1-3 x the planner's sigma Wrist: rate pulses (half the segments); fingers: a new closure target (40 % of segments). Payload: attached (20-500 g) or dropped at 35 % of segment boundaries: the "weight changes" regime. Randomised per trajectory (eta): gantry and gripper mass x0.85-1.15, line EA 2-15 kN (log), line damping x0.5-2, spool acceleration 0.10-0.30 m/s^2, command latency 0-0.15 s, mount calibration error sigma 0-4 cm (the host's IK only), swing damping x0.3-3. Estimator noise as the sim estimator's (perception/oracle_estimator.StateEstimate). """ from __future__ import annotations import argparse import copy import json import sys import time from collections import deque from pathlib import Path import mujoco import numpy as np import torch ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) from cablepilot.perception.oracle_estimator import StateEstimate # noqa: E402 from cablepilot.sim.host import HostModel, SpoolState # noqa: E402 from cablepilot.sim.kin import ccw_polygon, worst_hang_tensions # noqa: E402 from cablepilot.sim.scene import (Indices, build_model, grasp_point_local, load_config, mounts, random_room, room_centre, # noqa: E402 set_room_geometry, settle_spools, total_mass) from cablepilot.sim.state import to_state # noqa: E402 SEGMENTS = {"goto": 0.25, "dwell": 0.10, "stop": 0.10, "pump": 0.10, "descend": 0.15, "plan": 0.30} ETA_NAMES = ["gantry_mass_scale", "gripper_mass_scale", "EA_N", "line_damping_scale", "line_accel_mps2", "latency_s", "calib_sigma_m", "swing_damping_scale", "payload_max_kg"] DT = 0.1 D_G = -0.5147 # grasp point on the pole, Stringman hand (tasks/pick_place.PickPlaceTask.dG); per hand: Generator.dG class Generator: def __init__(self, scenario: str, seconds: float = 30.0, random_rooms: bool = True, hand: str | None = None): self.cfg = load_config(scenario=scenario, hand=hand) # grasp point's distance down the pole (SO-101: hand origin + gripperframe, -0.472; its 7.9 mm lateral offset # does not matter for keeping it off the floor) self.dG = float(self.cfg["body"]["gripper_origin_z_m"] + grasp_point_local(self.cfg)[2]) self.base_cfg = self.cfg self.random_rooms = random_rooms self.m = build_model(self.cfg, world=False, payload_kg=1e-4) self.d = mujoco.MjData(self.m) self.idx = Indices(self.m) self.h = self.m.opt.timestep self.sub = int(round(DT / self.h)) self.T = int(round(seconds / DT)) self.M = mounts(self.cfg) self.poly = ccw_polygon(self.M[:, :2]) self.pay = self.m.body("payload").id m = self.m self.orig = {k: getattr(m, k).copy() for k in ("body_mass", "body_inertia", "actuator_gainprm", "actuator_biasprm", "dof_damping")} self.gripper, self.gantry = m.body("gripper").id, m.body("gantry").id self.swing_dofs = slice(self.idx.gv + 3, self.idx.gv + 6) # ------------------------------------------------------------------------------------- geometry def inside_by(self, xy): a, b = self.poly, np.roll(self.poly, -1, 0) e = b - a n = np.c_[-e[:, 1], e[:, 0]] / np.linalg.norm(e, axis=1)[:, None] return float(((np.asarray(xy) - a) * n).sum(1).min()) def tension_ok(self, p, lim=26.0): t = worst_hang_tensions(torch.as_tensor(self.M, dtype=torch.float32), torch.as_tensor(p, dtype=torch.float32), total_mass(self.m) + 0.5) return float(t.max()) <= lim def random_point(self, rng, z_lo=None): lo, hi = self.poly.min(0), self.poly.max(0) for _ in range(1000): xy = rng.uniform(lo, hi) if self.inside_by(xy) < 0.35: continue gz = rng.uniform(0.04, 0.15) if z_lo else rng.uniform(0.15, 1.2) p = np.r_[xy, gz - self.dG] if self.tension_ok(p): return p return np.r_[self.poly.mean(0), 1.0] # ------------------------------------------------------------------------------------- randomisation def randomise(self, rng) -> tuple[dict, HostModel]: m, o = self.m, self.orig for k, v in o.items(): getattr(m, k)[:] = v if self.random_rooms: # a new room every trajectory: the model learns the machine, not a place self.cfg = set_room_geometry(self.base_cfg, m, self.idx, random_room(rng)) self.M = mounts(self.cfg) self.poly = ccw_polygon(self.M[:, :2]) e = {"gantry_mass_scale": rng.uniform(0.85, 1.15), "gripper_mass_scale": rng.uniform(0.85, 1.15), "EA_N": float(np.exp(rng.uniform(np.log(2000), np.log(15000)))), "line_damping_scale": rng.uniform(0.5, 2.0), "line_accel_mps2": rng.uniform(0.10, 0.30), "latency_s": rng.uniform(0.0, 0.15), "calib_sigma_m": rng.uniform(0.0, 0.04), "swing_damping_scale": float(np.exp(rng.uniform(np.log(0.3), np.log(3.0)))), "payload_max_kg": rng.uniform(0.0, 0.5)} m.body_mass[self.gantry] *= e["gantry_mass_scale"] m.body_inertia[self.gantry] *= e["gantry_mass_scale"] m.body_mass[self.gripper] *= e["gripper_mass_scale"] m.body_inertia[self.gripper] *= e["gripper_mass_scale"] m.dof_damping[self.swing_dofs] *= e["swing_damping_scale"] run = np.asarray(self.cfg["room"]["indirect_run_m"], float) span = np.linalg.norm(self.M - np.r_[room_centre(self.cfg), 1.2], axis=1) for i, a in enumerate(self.idx.spool): k = e["EA_N"] / (span[i] + run[i]) m.actuator_gainprm[a, 0] = k m.actuator_biasprm[a, 1] = -k m.actuator_biasprm[a, 2] *= e["line_damping_scale"] cfg_h = copy.deepcopy(self.cfg) # the host's belief of where the mounts are cfg_h["room"]["mounts_xyz"] = (self.M + rng.normal(0, e["calib_sigma_m"], self.M.shape)).tolist() host = HostModel(cfg_h, self.h, self.idx.travel) host.acc = e["line_accel_mps2"] return e, host # ------------------------------------------------------------------------------------- commands def segments(self, rng): """Endless generator of (kind, function(state, t_in_segment) -> (action, done)).""" sig = np.array([0.06, 0.06, 0.03, 0.4]) omega = np.sqrt(9.81 / 0.3089) names, probs = list(SEGMENTS), np.array(list(SEGMENTS.values())) while True: kind = str(rng.choice(names, p=probs / probs.sum())) wrist_rate = rng.choice([0.0, 0.0, rng.uniform(-1.5, 1.5)]) close = None if rng.random() > 0.4 else float(rng.choice([0.0, 1.0, rng.uniform(0, 1)])) if kind in ("goto", "stop", "descend"): tgt = self.random_point(rng, z_lo=(kind == "descend")) vmax = rng.uniform(0.1, 0.35) cut = rng.uniform(1.0, 4.0) if kind == "stop" else None hold = rng.uniform(0.5, 2.5) state = {"arrived": None} def f(x, t, tgt=tgt, vmax=vmax, cut=cut, hold=hold, state=state): if cut is not None and t > cut: return np.zeros(3), t > cut + rng.uniform(2.0, 5.0) d = tgt - x[:3] v = d * 0.8 n = np.linalg.norm(v) v = v * min(1.0, vmax / max(n, 1e-9)) if np.linalg.norm(d) < 0.03: state["arrived"] = state["arrived"] or t done = state["arrived"] is not None and t - state["arrived"] > hold or t > 25 return v, done elif kind == "dwell": dur = rng.uniform(1.0, 4.0) def f(x, t, dur=dur): return np.zeros(3), t > dur elif kind == "pump": w = omega * rng.uniform(0.7, 1.3) amp = rng.uniform(0.02, 0.08) ang = rng.uniform(0, 2 * np.pi) dur = rng.uniform(2.0, 5.0) def f(x, t, w=w, amp=amp, ang=ang, dur=dur): s = amp * np.sin(w * t) if t < dur else 0.0 return np.array([s * np.cos(ang), s * np.sin(ang), 0.0]), t > dur + 3.0 else: # plan: AR(0.9) random walk, a new command every 0.2 s scale = rng.uniform(1, 3) st = {"eps": np.zeros(4), "base": None, "k": -1} steps = int(rng.integers(8, 30)) def f(x, t, scale=scale, st=st, steps=steps): k = int(round(t / 0.2)) if st["base"] is None: st["base"] = np.r_[x[3:6], 0.0] if k != st["k"]: st["k"] = k st["eps"] = 0.9 * st["eps"] + np.sqrt(1 - 0.81) * rng.normal(0, 1, 4) * sig * scale v = st["base"][:3] * 0.9 ** k + st["eps"][:3] return v, k >= steps yield kind, f, wrist_rate, close # ------------------------------------------------------------------------------------- one trajectory def trajectory(self, seed: int) -> dict | None: rng = np.random.default_rng(seed) m, d, idx = self.m, self.d, self.idx eta, host = self.randomise(rng) mujoco.mj_resetData(m, d) p0 = self.random_point(rng) d.qpos[idx.gq:idx.gq + 3] = p0 d.qpos[idx.fq] = d.qpos[idx.frq] = -idx.travel m.body_mass[self.pay] = 1e-4 settle_spools(m, d, idx, self.cfg) d.ctrl[idx.finger_ctrl] = -idx.travel for _ in range(int(1.0 / self.h)): mujoco.mj_step(m, d) x = to_state(m, d, idx) spools = SpoolState(host, self.cfg, x) noise = StateEstimate(np.random.default_rng(seed + 1)) # command latency, in PHYSICS steps: a command sent at tick k reaches the host lat_sub physics steps later. # (v0 counted it in physics steps but delayed it by that many 0.1 s ticks: up to 3.7 s; the data's velocity # followed its command at ~0.4 /s instead of ~2-3 /s and the model learned that commands barely move the robot.) lat_sub = int(round(eta["latency_s"] / self.h)) pending = deque() # (physics step at which it applies, action) applied = np.r_[np.zeros(3), 0.0, x[12]] n_phys = 0 T = self.T S = np.empty((T + 1, 14)) A = np.empty((T, 5)) TEN = np.empty((T + 1, 4)) PAY = np.empty(T + 1) kinds = {} segs = self.segments(rng) kind, f, wr, cl = next(segs) t_seg, close = 0.0, x[12] v_prev = np.zeros(3) def rec(k): S[k] = noise(to_state(m, d, idx)) TEN[k] = -d.actuator_force[idx.spool] PAY[k] = m.body_mass[self.pay] rec(0) for k in range(T): xe = S[k] v, done = f(xe, t_seg) if done: kinds[kind] = kinds.get(kind, 0) + 1 kind, f, wr, cl = next(segs) t_seg = 0.0 if cl is not None: close = cl if rng.random() < 0.35: # the payload changes at a segment boundary on = m.body_mass[self.pay] > 1e-3 m.body_mass[self.pay] = 1e-4 if on else max(1e-4, rng.uniform(0.02, 0.5) * eta["payload_max_kg"] / 0.5) m.body_inertia[self.pay] = 1e-6 v, _ = f(xe, t_seg) # near the floor, descend no faster than the spools can stop: stopping distance v^2 / 2a with a = half the # line acceleration (a 0.35 m/s descent needs ~0.3 m; the first sample put G 23 cm below the floor). The fw # visual servo does the same ("descent speed as a function of how far there is left to fall"). G = xe[:3] + self.dG * np.r_[xe[6:8], 1.0] v = v.copy() v[2] = max(v[2], -np.sqrt(2 * 0.5 * eta["line_accel_mps2"] * max(G[2] - 0.04, 0.0))) # keep inside: a command that would leave the reachable box is turned back toward the centre if self.inside_by(xe[:2]) < 0.25 or G[2] < 0.08 or not self.tension_ok(xe[:3], 30.0): c = np.r_[self.poly.mean(0), 1.0] v = 0.15 * (c - xe[:3]) / max(np.linalg.norm(c - xe[:3]), 1e-6) sp = np.linalg.norm(v) v = v * min(1.0, 0.35 / max(sp, 1e-9)) a = np.r_[v, wr if t_seg < 3.0 else 0.0, close] A[k] = a pending.append((n_phys + lat_sub, a, xe[:3].copy())) for _ in range(self.sub): while pending and pending[0][0] <= n_phys: _, applied, p_at_send = pending.popleft() spools.command(p_at_send, applied) n_phys += 1 u = spools.substep(-d.actuator_force[idx.spool], self.h) d.ctrl[idx.spool] = u[:4] d.ctrl[idx.wrist_ctrl], d.ctrl[idx.finger_ctrl] = u[4], u[5] mujoco.mj_step(m, d) if not np.isfinite(d.qpos).all() or np.abs(d.qvel[:6]).max() > 20: return None rec(k + 1) t_seg += DT v_prev = v return {"states": S.astype(np.float32), "actions": A.astype(np.float32), "tensions": TEN.astype(np.float32), "payload": PAY.astype(np.float32), "eta": np.array([eta[k] for k in ETA_NAMES], np.float32), "kinds": kinds, "mounts": self.M.astype(np.float32), "runs": np.asarray(self.cfg["room"]["indirect_run_m"], np.float32)} def _shard_job(args): out, scenario, shard_i, seed0, n, seconds, random_rooms, hand = args gen = Generator(scenario, seconds, random_rooms, hand) rows, rejected, kinds = [], 0, {} t0 = time.perf_counter() for i in range(n): tr = gen.trajectory(seed0 + i) if tr is None: rejected += 1 continue for k, v in tr["kinds"].items(): kinds[k] = kinds.get(k, 0) + v rows.append(tr) shard = {k: np.stack([r[k] for r in rows]) for k in ("states", "actions", "tensions", "payload", "eta", "mounts", "runs")} np.savez_compressed(Path(out) / f"shard_{shard_i:04d}.npz", **shard, eta_names=np.array(ETA_NAMES), dt=DT) return len(rows), rejected, kinds, time.perf_counter() - t0 def main(): ap = argparse.ArgumentParser() ap.add_argument("--scenario", default="configs/scenarios/clear_playroom.yaml") ap.add_argument("--n", type=int, default=20) ap.add_argument("--seed", type=int, default=0) ap.add_argument("--seconds", type=float, default=30.0) ap.add_argument("--shard-size", type=int, default=100) ap.add_argument("--workers", type=int, default=1) ap.add_argument("--hand", default=None, help="stringman | so101 (default: configs/stringman.yaml `hand`)") ap.add_argument("--fixed-room", action="store_true", help="the scenario's room for every trajectory (test sets of a real room)") ap.add_argument("--out", default=str(ROOT / "data/cdpr_v0_sample")) args = ap.parse_args() out = Path(args.out) out.mkdir(parents=True, exist_ok=True) size = min(args.shard_size, args.n) jobs = [(str(out), args.scenario, i, args.seed + i * size, size, args.seconds, not args.fixed_room, args.hand) for i in range(args.n // size)] t0 = time.perf_counter() done, rej, kinds, cpu = 0, 0, {}, 0.0 if args.workers > 1: from multiprocessing import Pool with Pool(args.workers) as pool: it = pool.imap_unordered(_shard_job, jobs) for k, (n, r, kd, s) in enumerate(it): done, rej, cpu = done + n, rej + r, cpu + s for a, v in kd.items(): kinds[a] = kinds.get(a, 0) + v print(f" {k + 1}/{len(jobs)} shards, {done} trajectories, {time.perf_counter() - t0:.0f} s", flush=True) else: for job in jobs: n, r, kd, s = _shard_job(job) done, rej, cpu = done + n, rej + r, cpu + s for a, v in kd.items(): kinds[a] = kinds.get(a, 0) + v meta = {"trajectories": done, "rejected": rej, "seconds_each": args.seconds, "rate_hz": 1 / DT, "scenario": args.scenario, "hand": Generator(args.scenario, 1.0, False, args.hand).cfg["hand"], "dG": Generator(args.scenario, 1.0, False, args.hand).dG, "random_rooms": not args.fixed_room, "wall_s": time.perf_counter() - t0, "cpu_s": cpu, "cpu_s_per_trajectory": cpu / max(done, 1), "segments": kinds, "eta_names": ETA_NAMES, "model_state": "states (14): p, v, s, s_dot, wrist, wrist_dot, closure, closure_dot. tensions NOT a model input."} (out / "manifest.json").write_text(json.dumps(meta, indent=2)) print(json.dumps(meta, indent=1)) if __name__ == "__main__": main()