Blame for scripts/cdpr_datagen.py

"""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()