| """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 |
| from cablepilot.sim.host import HostModel, SpoolState |
| from cablepilot.sim.kin import ccw_polygon, worst_hang_tensions |
| from cablepilot.sim.scene import (Indices, build_model, grasp_point_local, load_config, mounts, random_room, room_centre, |
| set_room_geometry, settle_spools, total_mass) |
| from cablepilot.sim.state import to_state |
|
|
| 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 |
|
|
|
|
| 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) |
| |
| |
| 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) |
|
|
| |
| 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] |
|
|
| |
| 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: |
| 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) |
| 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 |
|
|
| |
| 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: |
| 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 |
|
|
| |
| 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)) |
| |
| |
| |
| lat_sub = int(round(eta["latency_s"] / self.h)) |
| pending = deque() |
| 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: |
| 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) |
| |
| |
| |
| 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))) |
| |
| 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() |
|
|
