Add files using upload-large-folder tool

applications/multi_robot/README.md

@@ -0,0 +1,32 @@
+# Multi-robot motion planning
+
+Reproduces the multi-robot results of *Continuous-Time Constrained Motion Planning with Global
+Convergence Guarantees*, in two experiments:
+
+| folder | experiment | paper |
+|---|---|---|
+| [`robotarium/`](robotarium/README.md) | the Robotarium arena: 5 maps x N in {3,6,9,12}, single-plan success, the hardware runs; GlobMP behind a diffusion planner's projector | Sec. 4.2, App. `sec:steering` |
+| [`multi_robot_extension/`](multi_robot_extension/README.md) | the simulated suite: 8 maps x N = 3..21 x 3 velocity caps, the eight suite tables and figures | App. `app:exp-multirobot` |
+
+Each folder has the same layout and runs on its own:
+
+```
+<experiment>/
+  *.py          GlobMP's side of the experiment (instances, evaluate and figures are shared, identical copies)
+  run/          the entry scripts
+  data/         instances and the result files the paper's tables and figures read
+  *_baselines/  the baselines and the benchmark they run in: config/, data_prep/, train/, run/, eval/, tools/,
+                checkpoints/ and one folder per original repository (same layout as steering/xhand_baselines/)
+```
+
+GlobMP itself is the `globmp` package at the repository root; nothing here changes it.
+
+## Environments
+
+* GlobMP side: Python >= 3.10 with numpy, scipy, torch (CPU is enough), matplotlib, and the `globmp` package
+  on the path: `export PYTHONPATH=<repo root>` (or `pip install -e <repo root>`).  The runs behind the paper
+  used one thread per job (`OMP_NUM_THREADS=1`, `run_grid.py --threads 1`) on pinned cores.
+* Benchmark side (only for the benchmark's predicate, the instance export and the Robotarium sampler): the
+  co-authors' reproduction of the DiRecT/MMD multi-robot benchmark (`<experiment>_baselines/` in each folder) with its `direct` conda env.  Two quirks:
+  `PYTHONPATH` must put `multi_robot/` and `multi_robot/deps/torch_robotics` first, and `PYTHONNOUSERSITE`
+  must not be set.  `EnvEmpty2D` has no checkpoint; `EnvEmptyNoWait2D` is used.

applications/multi_robot/robotarium/README.md

@@ -0,0 +1,114 @@
+# Multi-robot: Robotarium arena (Sec. 4.2)
+
+Single-plan success of GlobMP on the Robotarium arena (5 maps x N in {3,6,9,12}) and the baselines it is compared
+with, plus GlobMP behind a diffusion planner's projector interface (App. `sec:steering`).
+
+| paper item | what | produced by |
+|---|---|---|
+| Sec. 4.2, Fig. `robotarium_sr_one_compact.pdf` | single-plan success on the Robotarium arena, 5 maps x N in {3,6,9,12} | `run/run_cell.sh` (GlobMP as the sampler's projector), `run/score.py`, `figures.py robotarium` |
+| Sec. 4.2, Fig. `robotarium_gallery.pdf` | frames of the real Robotarium runs | not reproducible from this package (see [Hardware gallery](#hardware-gallery)) |
+| App. `sec:steering` (Steering a generative planner) | GlobMP behind a diffusion planner's projector interface | `projector.py`, `batch.py` (this is the configuration of the Robotarium runs) |
+
+## Files
+
+```
+projector.py          GlobMP as DiRecT's projector (BaseProjector.project): steering calls, final certified solve,
+                      robot-by-robot certification with wall, timing and detour repairs
+batch.py              all candidate teams of one projector call as one batched problem (CPU or GPU)
+instances.py          map names, obstacle scales; load_robotarium; disc covers of boxes        (shared, identical copy)
+evaluate.py           exact continuous-time robot-robot check and obstacle checks              (shared, identical copy)
+figures.py            the multi-robot figures from result files (`figures.py robotarium`)      (shared, identical copy)
+run/
+  direct_hook.py      the three-line wiring of projector.py into the benchmark's DiRecT fork
+  run_cell.sh         one Robotarium cell: diffusion sampler + GlobMP projector, the paper's configuration
+  score.py            select (first certified candidate per instance), bench (the comparison's predicate),
+                      exact (offline re-check without the benchmark)
+data/robotarium/      instances.npz (10 per cell), maps.json, geometry.json (boxes in metres)
+data/results/         robotarium_cells.json (baselines), robotarium_gpa_sr_one.json, robotarium_gpa_sr_one_detail.json
+robotarium_baselines/ the baselines (DiRecT, CD, Final Proj., PCD-SHD, PCD-LB, MMD-CBS) on this arena: the forks, data
+                      preparation, training, runners, scorer, models, the N = 12 upload scripts (robotarium_baselines/README.md)
+```
+
+`instances.py`, `evaluate.py` and `figures.py` are shared by both experiments; each folder keeps an identical copy,
+so either folder runs on its own.
+
+## Environments
+
+* GlobMP side: Python >= 3.10 with numpy, scipy, torch (CPU is enough), matplotlib, and the `globmp` package
+  on the path: `export PYTHONPATH=<repo root>` (or `pip install -e <repo root>`).  The runs behind the paper
+  used one thread per job on pinned cores.
+* Benchmark side (only for the benchmark's predicate, the instance export and the Robotarium sampler): the
+  co-authors' reproduction of the DiRecT/MMD multi-robot benchmark (in `robotarium_baselines/`) with its `direct` conda env.  Two quirks:
+  `PYTHONPATH` must put `multi_robot/` and `multi_robot/deps/torch_robotics` first, and `PYTHONNOUSERSITE`
+  must not be set.  `EnvEmpty2D` has no checkpoint; `EnvEmptyNoWait2D` is used.
+
+The baselines' code is in `robotarium_baselines/` (`robotarium_baselines/README.md`).  Their records
+(`data/results/robotarium_cells.json`, a copy of `robotarium_baselines/eval/robotarium_cells.json`) were written by its
+Robotarium scorer (`robotarium_baselines/eval/original/`; `robotarium_baselines/eval/score_cells.py` is the scorer for
+new runs and reproduces them). The MMD-CBS cells are from the 2026-10-15 re-run with the isotropic normalizer
+(`robotarium_baselines/eval/mmd_cbs_rerun.md`).
+
+The trained single-robot models of the paper's setting (v2) are in `robotarium_baselines/checkpoints/<EnvClass>-RobotPlanarDisk/`
+(`args.yaml` and the final EMA weights `checkpoints/ema_model_current_state_dict.pth`), byte-identical to the models
+GlobMP's runs loaded.  Conveyor's is the wide-channel model retrained for `RBT_WIDE=1`; the other four are shared with
+the earlier v1 setting.  Point the baselines' `MODELS_ROOT` / `--models-root`, and the `data_trained_models/` of the
+fork behind `DIRECT_MR`, at this folder.  The planners rebuild their normalizer from the training demonstrations at
+start-up; those (about 4 GB) are not included and are regenerated by `robotarium_baselines/data_prep/generate_data.sh`.
+
+GlobMP is scored with the same predicates: robot pairs at least 2r apart at the 64 waypoints and the
+benchmark's `task.get_trajs_collision_and_free` at its default interpolation (k = 0).
+
+## Robotarium (Sec. 4.2)
+
+The Robotarium numbers come from GlobMP used as the projector of the comparison's own diffusion sampler
+(DiRecT with the single-robot models trained for the arena, r = 0.13 m, isotropic normalizer), on the
+comparison's 10 instances per cell.  At each of the 25 projection calls of an instance the sampler's batch of
+64 candidate teams is the reference of a tracking cost; intermediate calls steer (short restoration), the final
+call solves every candidate team, and the returned plan is the first candidate GlobMP certifies (its
+piecewise-linear motion is exact-checked against walls, arena, speed and 2r).  If no candidate certifies with
+seed 18, fresh diffusion batches are drawn (seeds 19, 20, ...) and the first certified plan is returned.
+
+```bash
+export DIRECT_MR=$PWD/robotarium_baselines/DiRecT          # the baselines' DiRecT fork, with run/direct_hook.py wired in
+export DIRECT_PY=<conda>/envs/direct/bin/python OUT=$PWD/rbt_out
+run/run_cell.sh crossroom 12 0                       # map N gpu: 10 instances x 64 candidates
+SEED=19 EXP_TAG=_s19i5 NINITS=1 INST_NPZ=<single-instance npz> run/run_cell.sh crossroom 12 0   # a retry
+python run/score.py select --results $DIRECT_MR/results --certs $OUT --out ret
+(cd $DIRECT_MR && <Robotarium env vars of run_cell.sh> $DIRECT_PY <here>/run/score.py bench crossroom 12 ret/returned_crossroom_N12.npz)
+python run/score.py exact --returned ret --sr-json data/results/robotarium_gpa_sr_one.json
+python figures.py robotarium --cells data/results/robotarium_cells.json --gpa data/results/robotarium_gpa_sr_one.json --out robotarium_sr_one_compact.pdf
+```
+
+The cells were run on one A100 each with 40 CPU workers for the robot-by-robot certification; Conveyor N = 9
+and N = 12 were split into five jobs of two instances (`score.py` SPLITS), and the retried instances are listed
+in `score.py` RETRIES (Conveyor N = 12: instances 1, 3, 4, 5, 8, 9; Cross-Room N = 12: instance 5).  Conveyor
+N = 12 instances 3 and 4 never certified (seeds 18 to 24), which is the 0.8 in the figure.  Wall time per cell
+ranged from minutes (Empty, small N) to several hours (Conveyor N = 12).  The comparison's predicate does not
+require the goal to be reached exactly: the sampler pins only the start, and the returned plans end on average
+4 to 5 cm (at most 22 cm) from their goals; the same holds for every method in the comparison.
+
+### Hardware gallery
+
+`robotarium_gallery.pdf` shows frames of the overhead videos of the real runs (12-robot ring instance,
+Cross-Room and Conveyor) with tracked paths; it is drawn by the co-authors' `fig_robotarium_gallery.py` from the
+videos and tracks, which are not part of this package.  The GlobMP plans executed there ("best picks") were
+produced on 2026-09-24 by the projector above, on the antipodal ring instances of an earlier benchmark
+version, with a selection loop over steering settings (keep-right, circulation `GPA_CIRC_MODE=reflect`,
+`GPA_SHIFT_SMOOTH`, intermediate budget) that kept the first certified plan passing a review (channel use, path
+length at most 1.3 times the straight lines, a clean replay in the Robotarium simulator).  The submission
+scripts were generated by the co-authors' uploader.  `projector.py` supports those settings; the loop and the
+uploader are not included.
+
+## Validation of this port (2026-10-12)
+
+Robotarium: `projector.py` + `batch.py` against the original projector, on identical inputs (Cross-Room N = 3
+and N = 6, a steering call and a final call, including the robot-by-robot certification): outputs agree to
+3e-15 in iteration-limited runs (each code differs from its own repeat by 1e-14) and to 6e-4 (normalized units)
+in full runs, with identical certified teams.  The final-call certificate screen of `projector.py` on the saved
+64 x 10 candidates of Empty N = 3, Cross-Room N = 6 and Drop-Region N = 12 reproduces the original certificate
+logs on 30 of 30 instances (625, 267 and 314 certified candidates; same first certified candidate).
+`score.py exact` on the 20 saved returned-plan files gives the published single-plan success on all 20 cells
+(Conveyor N = 12: 0.8, instances 3 and 4).  An end-to-end rerun of the diffusion sampler was not done here (it
+needs the benchmark's Robotarium fork and a GPU).
+
+Figures from the files in `data/results/`: `figures.py robotarium` matches the paper's PDF.

applications/multi_robot/robotarium/batch.py

@@ -0,0 +1,231 @@
+"""All candidate teams of one projector call solved as ONE batched problem with a leading team axis G.
+
+Same algorithm as globmp.Problem in 'orthant' mode: Algorithm 1 (restoration, Barzilai-Borwein steps with a
+nonmonotone backtracking reference) and a few blocks of Algorithm 2 (gradient projection), with per-team
+masks instead of a Python loop over teams. The chart (N, pseudoinverse), the cost Hessian and the constraint
+layout are shared by all teams; endpoints, tracking reference and step sizes are per team. Runs on CPU or GPU.
+
+Constraint layout (the same for every team): pairwise separation 2r, a per-robot speed bound, the workspace
+box (hard) and, per robot, every obstacle disc within `reach` of that robot's reference in ANY team.
+
+Used by projector.py (the GlobMP projector behind the diffusion sampler of the Robotarium runs)."""
+import math
+import time
+
+import numpy as np
+import torch
+
+from globmp import Problem, spline as bern
+from globmp import problem as _problem
+
+DT = torch.float64
+
+
+class TeamBatch:
+    def __init__(self, W, n_agents, d, M, delta, robot_r, obstacles, ws_bounds, dx_max, De=30, reach=0.25,
+                 device='cuda'):
+        """W: (G, L+1, n) references (sampler frame). obstacles: list of (center (2,), radius)."""
+        self.dev = torch.device(device)
+        G, L1, n = W.shape; L = L1 - 1
+        self.G, self.n, self.na, self.d, self.M, self.dl = G, n, n // 2, d, M, float(delta)
+        tk = dict(dtype=DT, device=self.dev)
+        # shared chart from one CPU problem (cached per (d, ell, ell_j, M) inside globmp)
+        q = Problem(n, d, 4, 0, M, W[0, 0], W[0, -1], ell_j=2, delta=delta, mode='orthant', De=De)
+        Api, N = _problem._CHART_CACHE[('chart', d, 0, 2, M)]
+        self.N = N.to(**tk); self.P = M * (d + 1)
+        self.NpT = torch.linalg.pinv(N).T.to(**tk)                                   # for the initial fit
+        self.u = bern.ones_u(d).to(**tk); self.uu = torch.outer(self.u, self.u)
+        self.D1 = q.D1.to(**tk); self.E1 = q.E1.to(**tk); self.Eo = q.Eo.to(**tk)
+        C_, idx_ = bern._S_coeffs(d); self.SC = C_.to(**tk); self.Sidx = idx_.to(self.dev)
+        C1_, idx1_ = bern._S_coeffs(d - 1); self.SC1 = C1_.to(**tk); self.Sidx1 = idx1_.to(self.dev)
+        # per-team endpoints -> Gamma0 = start (x) Api[:,0] + goal (x) Api[:,1]
+        Wt = torch.as_tensor(W, **tk)
+        A0 = Api[0].to(**tk); A1 = Api[1].to(**tk)
+        self.Gamma0 = Wt[:, 0, :, None] * A0 + Wt[:, -1, :, None] * A1              # (G, n, P)
+        # tracking cost against the piecewise-linear reference (linear in W; exact Gauss-Legendre)
+        Tm, Qm = self._tracking_maps(L, M, d)
+        Gk = torch.block_diag(*[bern.cost_gram(d, 0)] * M).to(**tk)
+        self.H = self.N.T @ Gk @ self.N; self.H = 0.5 * (self.H + self.H.T)
+        Rd = torch.einsum('lc,gln->gnc', Tm.to(**tk), Wt)                            # (G, n, P)
+        cref = torch.einsum('lk,gln,gkn->g', Qm.to(**tk), Wt, Wt)
+        A = self.Gamma0
+        self.G0 = torch.einsum('gnp,pq,qf->gnf', A, Gk, self.N) - torch.einsum('gnp,pf->gnf', Rd * M, self.N)
+        self.c0 = torch.einsum('pq,gip,giq->g', Gk, A, A) - 2.0 * (Rd * M * A).sum((1, 2)) + cref * M
+        ev = torch.linalg.eigvalsh(self.H); self.beta = 2 * float(ev.max())
+        # constraint layout
+        na = self.na; r2 = 2.0 * robot_r
+        pairs = [(a, b) for a in range(na) for b in range(a + 1, na)]
+        self.sepI = torch.tensor([[2 * a, 2 * a + 1] for a, _ in pairs], device=self.dev)
+        self.sepJ = torch.tensor([[2 * b, 2 * b + 1] for _, b in pairs], device=self.dev)
+        self.sepR = r2
+        self.spdI = torch.tensor([[2 * a, 2 * a + 1] for a in range(na)], device=self.dev)
+        self.vmax = float(dx_max) * L / M
+        rows, cen, rad = [], [], []
+        if obstacles:
+            Cc = np.array([c for c, _ in obstacles], float); Rc = np.array([r for _, r in obstacles], float)
+            for a in range(na):
+                Pa = W[:, :, 2 * a:2 * a + 2].reshape(-1, 2)
+                dmin = np.linalg.norm(Pa[:, None] - Cc[None], axis=-1).min(0)
+                for k in np.where(dmin <= Rc + robot_r + delta + reach)[0]:
+                    rows.append([2 * a, 2 * a + 1]); cen.append(Cc[k]); rad.append(Rc[k] + robot_r)
+        self.nsph = len(rows)
+        if rows:
+            self.sphI = torch.tensor(rows, device=self.dev); self.sphC = torch.as_tensor(np.array(cen), **tk)
+            self.sphR = torch.as_tensor(np.array(rad), **tk)
+        lo, hi = np.asarray(ws_bounds[0], float), np.asarray(ws_bounds[1], float)
+        self.lo = torch.as_tensor(np.tile(lo + robot_r - delta + 1e-6, na), **tk)
+        self.hi = torch.as_tensor(np.tile(hi - robot_r + delta - 1e-6, na), **tk)
+        # tolerances per block, in pi_all order
+        dl = self.dl; tv = []
+        if self.nsph: tv.append(((2 * self.sphR * dl + dl * dl) ** 2)[:, None].expand(-1, M).reshape(-1))
+        tv.append(torch.full((len(pairs) * M,), (2 * r2 * dl + dl * dl) ** 2, **tk))
+        tv.append(torch.full((na * M,), 1e-12, **tk))                              # speed
+        tv.append(torch.full((2 * M * n,), 1e-12, **tk))                           # box: hard
+        self.tv = torch.cat(tv)
+        # initial chart coordinates: least squares to the M+1 waypoints, linear inside each segment
+        wp = Wt[:, np.round(np.linspace(0, L, M + 1)).astype(int), :]                 # (G, M+1, n)
+        ss = torch.linspace(0, 1, d + 1, **tk)
+        tgt = wp[:, :-1, :, None] * (1 - ss) + wp[:, 1:, :, None] * ss               # (G, M, n, d+1)
+        Gm = tgt.permute(0, 2, 1, 3).reshape(G, n, -1)
+        self.Z0 = (Gm - self.Gamma0) @ self.NpT                                      # (G, n, d_f)
+
+    @staticmethod
+    def _tracking_maps(L, M, d):
+        """Rd(W) = sum_l Tm[l,c] W[l,:] and cref(W) = W^T Qm W, both exact (piecewise-linear reference)."""
+        nq = int(math.ceil((d + 2) / 2)); xg, wg = np.polynomial.legendre.leggauss(nq)
+        xg = 0.5 * (xg + 1.0); wg = 0.5 * wg
+        Tm = np.zeros((L + 1, M * (d + 1))); Qm = np.zeros((L + 1, L + 1))
+        tb = np.linspace(0, 1, L + 1)
+        for a in range(M):
+            t0, t1 = a / M, (a + 1) / M
+            cuts = np.unique(np.r_[t0, t1, tb[(tb > t0) & (tb < t1)]])
+            for c0, c1 in zip(cuts[:-1], cuts[1:]):
+                h = c1 - c0; tq = c0 + h * xg; uq = (tq - t0) / (t1 - t0)
+                bu = bern.bernstein_basis_eval(d, torch.as_tensor(uq, dtype=DT)).numpy()   # (d+1, nq)
+                i = np.clip(np.floor(tq * L).astype(int), 0, L - 1); f = tq * L - i
+                for qq in range(nq):
+                    w_ = wg[qq] * h
+                    for (li, cf) in ((i[qq], 1 - f[qq]), (i[qq] + 1, f[qq])):
+                        Tm[li, a * (d + 1):(a + 1) * (d + 1)] += w_ * cf * bu[:, qq]
+                    Qm[i[qq], i[qq]] += w_ * (1 - f[qq]) ** 2; Qm[i[qq] + 1, i[qq] + 1] += w_ * f[qq] ** 2
+                    Qm[i[qq], i[qq] + 1] += w_ * (1 - f[qq]) * f[qq]; Qm[i[qq] + 1, i[qq]] += w_ * (1 - f[qq]) * f[qq]
+        return torch.as_tensor(Tm, dtype=DT), torch.as_tensor(Qm, dtype=DT)
+
+    # ---- blocks
+    def _S(self, Q, d1=False):
+        C, idx = (self.SC1, self.Sidx1) if d1 else (self.SC, self.Sidx)
+        W = (Q * C).reshape(*Q.shape[:-2], -1); dd = (self.d - 1) if d1 else self.d
+        out = torch.zeros(*Q.shape[:-2], 2 * dd + 1, dtype=DT, device=self.dev)
+        out.index_add_(-1, idx, W)
+        return out
+
+    def _elev(self, c):
+        Mx = 0.5 * (c[..., :, None] * self.u + self.u[:, None] * c[..., None, :])
+        return self._S(Mx)
+
+    def segs(self, Z):
+        Gm = self.Gamma0 + Z @ self.N.T                                              # (G, n, P)
+        return Gm.view(self.G, self.n, self.M, self.d + 1).permute(0, 2, 1, 3)       # (G, M, n, d+1)
+
+    def pi_all(self, Z):
+        S = self.segs(Z); dl = self.dl; out = []
+        if self.nsph:
+            Gc = S[:, :, self.sphI, :] - self.sphC[None, None, :, :, None] * self.u  # (G,M,K,2,d+1)
+            Mm = torch.einsum('gmkip,gmkiq->gmkpq', Gc, Gc) - ((self.sphR + dl) ** 2)[None, None, :, None, None] * self.uu
+            out.append(self._S(Mm).permute(0, 2, 1, 3).reshape(self.G, -1, 2 * self.d + 1))
+        Gc = S[:, :, self.sepI, :] - S[:, :, self.sepJ, :]                           # (G,M,Pp,2,d+1)
+        Mm = torch.einsum('gmkip,gmkiq->gmkpq', Gc, Gc) - (self.sepR + dl) ** 2 * self.uu
+        out.append(self._S(Mm).permute(0, 2, 1, 3).reshape(self.G, -1, 2 * self.d + 1))
+        Gd = S[:, :, self.spdI, :] @ self.D1                                         # (G,M,na,2,d)
+        Mm = torch.einsum('gmkip,gmkiq->gmkpq', Gd, Gd)
+        pe = self._S(Mm, d1=True) @ self.E1.T                                        # (G,M,na,2d+1)
+        out.append((self.vmax ** 2 - pe).permute(0, 2, 1, 3).reshape(self.G, -1, 2 * self.d + 1))
+        lo = self.lo + dl; hi = self.hi - dl
+        pl = self._elev(S - lo[None, None, :, None] * self.u).reshape(self.G, -1, 2 * self.d + 1)
+        ph = self._elev(hi[None, None, :, None] * self.u - S).reshape(self.G, -1, 2 * self.d + 1)
+        out += [pl, ph]
+        return torch.cat(out, 1)                                                     # (G, B, 2d+1)
+
+    def per_block(self, Z):
+        pe = self.pi_all(Z) @ self.Eo.T
+        return (torch.clamp(-pe, min=0) ** 2).sum(-1)                                # (G, B)
+
+    def cost(self, Z):
+        return self.c0 + 2 * (Z * self.G0).sum((1, 2)) + torch.einsum('pq,gip,giq->g', self.H, Z, Z)
+
+    # ---- Algorithm 1, batched
+    def project(self, Z, cap=3000, deadline=None, active=None):
+        G = self.G; tv = self.tv
+        eta = torch.full((G,), 1e-3, dtype=DT, device=self.dev)
+        Zp = None; gp = None; hist = []
+        live = torch.ones(G, dtype=torch.bool, device=self.dev) if active is None else active.clone()
+        ok = torch.zeros(G, dtype=torch.bool, device=self.dev)
+        for it in range(cap):
+            with torch.no_grad():
+                d2 = self.per_block(Z); act = d2 > tv
+            done = ~act.any(1); ok |= done & live; live &= ~done
+            if not bool(live.any()) or (deadline and time.time() > deadline):
+                break
+            Zg = Z.detach().requires_grad_(True)
+            with torch.enable_grad():
+                per = self.per_block(Zg); f = (act.to(DT) * per).sum(1)
+                (g,) = torch.autograd.grad(f.sum(), [Zg])
+            f0 = f.detach()
+            if Zp is not None:
+                s_ = (Z - Zp).reshape(G, -1); y_ = (g - gp).reshape(G, -1); sy = (s_ * y_).sum(1)
+                eta = torch.where(sy > 1e-300, (s_ * s_).sum(1) / sy.clamp_min(1e-300),
+                                  torch.minimum(eta * 1.3, 2 * f0 / (g * g).sum((1, 2)).clamp_min(1e-300)))
+            else:
+                eta = torch.minimum(eta, 2 * f0 / (g * g).sum((1, 2)).clamp_min(1e-300))
+            hist.append(f0); fref = torch.stack(hist[-10:]).max(0).values if Zp is not None else f0 * (1 + 1e-12) + 1e-300
+            Zt = Z; acc = ~live
+            for _ in range(40):
+                cand = (Z - eta[:, None, None] * g).detach()
+                with torch.no_grad():
+                    ft = (act.to(DT) * self.per_block(cand)).sum(1)
+                good = (ft < fref) & ~acc
+                Zt = torch.where(good[:, None, None], cand, Zt); acc |= good
+                if bool(acc.all()): break
+                eta = torch.where(acc, eta, eta * 0.5)
+            Zp = torch.where(live[:, None, None], Z, Zp if Zp is not None else Z)
+            gp = torch.where(live[:, None, None], g, gp if gp is not None else g)
+            Z = torch.where(live[:, None, None], Zt, Z)
+        with torch.no_grad():
+            ok |= ~(self.per_block(Z) > tv).any(1)
+        return Z.detach(), ok
+
+    # ---- Algorithm 2, batched (a few blocks)
+    def solve(self, Z, ok, blocks=4, steps=3, deadline=None):
+        Zp = Z.clone(); tc = torch.zeros(self.G, dtype=DT, device=self.dev)
+        for _ in range(blocks):
+            if deadline and time.time() > deadline: break
+            tc = tc + 1
+            y = Z + ((tc - 1) / (tc + 2))[:, None, None] * (Z - Zp)
+            for _ in range(steps):
+                y = (y - (1 / self.beta) * 2 * (self.G0 + y @ self.H)).detach()
+            cand, okc = self.project(y, cap=600, deadline=deadline, active=ok)
+            better = ok & okc & (self.cost(cand) <= self.cost(Z))
+            Zp = torch.where(better[:, None, None], Z, Zp); Z = torch.where(better[:, None, None], cand, Z)
+            tc = torch.where(better, tc, torch.zeros_like(tc))
+        return Z
+
+    def curve(self, Z, L):
+        tt = torch.linspace(0, 1, L + 1, dtype=DT, device=self.dev)
+        seg = torch.clamp((tt * self.M).floor().long(), 0, self.M - 1)
+        bu = bern.bernstein_basis_eval(self.d, (tt * self.M - seg).cpu()).to(device=self.dev, dtype=DT)
+        S = self.segs(Z)                                                             # (G, M, n, d+1)
+        return torch.einsum('ginp,pi->gin', S[:, seg], bu)                           # (G, L+1, n)
+
+
+def solve_teams(W, n_agents, robot_r, obstacles, ws_bounds, dx_max, d=5, M=21, delta=0.005, budget=20.0,
+                cap=3000, alg2_blocks=4, partial=False, device='cuda'):
+    """W: (G, L+1, n) numpy. Returns ((G, L+1, n) numpy, ok (G,) bool). With partial=True, teams that did not
+    restore return their partially restored curve (intermediate steering calls)."""
+    L = W.shape[1] - 1; M = min(M, L)
+    tb = TeamBatch(W, n_agents, d, M, delta, robot_r, obstacles, ws_bounds, dx_max, device=device)
+    dl = time.time() + budget
+    Z, ok = tb.project(tb.Z0, cap=cap, deadline=dl)
+    if alg2_blocks > 0 and bool(ok.any()):
+        Z = tb.solve(Z, ok, blocks=alg2_blocks, deadline=dl)
+    out = tb.curve(Z, L).cpu().numpy()
+    return out, ok.cpu().numpy()

applications/multi_robot/robotarium/evaluate.py

@@ -0,0 +1,208 @@
+"""Scoring of waypoint plans.
+
+A plan is executed (and scored by the benchmark) as the straight lines between its waypoints.  Three checks:
+
+* robot-robot, exact in continuous time: for two robots moving linearly over one interval the minimum
+  distance has a closed form (pair_min_dist_exact), so the check is over every t, not a dense sample.
+* obstacles with the BENCHMARK'S OWN predicate, `task.get_trajs_collision_and_free(num_interpolation=k)`;
+  the tables use k = 0 (the 64 waypoints only, the benchmark's default).  Needs the benchmark code.
+* obstacles and arena, exact for discs and axis-aligned boxes (exact_obstacle_clearance), with no benchmark
+  code; used to re-check the Robotarium plans offline.
+
+The paper's success numbers use the benchmark's own definitions (suite_metrics, robotarium/run/score.py);
+nothing here re-defines success.
+
+    python evaluate.py suite <map> --results-dir DIR --trained-models DIR      # benchmark env: SR CS Coll DA Time
+    python evaluate.py robotarium <map> <N> <plan.npz> [--inst K]              # any env: exact offline check
+"""
+import argparse
+import glob
+import json
+import os
+import sys
+
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+
+# ---- robot-robot, exact
+def pair_min_dist_exact(A, B):
+    """A, B: (..., H, 2) waypoints.  Exact minimum distance over the piecewise-linear motion."""
+    P = A[..., :-1, :] - B[..., :-1, :]
+    D = (A[..., 1:, :] - A[..., :-1, :]) - (B[..., 1:, :] - B[..., :-1, :])
+    dd = (D * D).sum(-1)
+    s = np.clip(np.where(dd > 1e-18, -(P * D).sum(-1) / np.maximum(dd, 1e-18), 0.0), 0.0, 1.0)
+    Q = P + s[..., None] * D
+    return np.sqrt((Q * Q).sum(-1)).min(-1)
+
+
+def eval_trajs(pos, radius):
+    """pos: (S, I, A, H, 2).  Per-configuration robot-robot success at the waypoints and in continuous time."""
+    S, I, A, H, _ = pos.shape
+    need = 2.0 * radius
+    wp_ok = np.ones((S, I), bool); ct_ok = np.ones((S, I), bool)
+    wp_min = np.full((S, I), np.inf); ct_min = np.full((S, I), np.inf)
+    for a in range(A):
+        for b in range(a + 1, A):
+            Pa, Pb = pos[:, :, a], pos[:, :, b]
+            dw = np.linalg.norm(Pa - Pb, axis=-1).min(-1)
+            dc = pair_min_dist_exact(Pa, Pb)
+            wp_min = np.minimum(wp_min, dw); ct_min = np.minimum(ct_min, dc)
+            wp_ok &= dw >= need; ct_ok &= dc >= need
+    return dict(wp_success=float(wp_ok.any(0).mean()), ct_success=float(ct_ok.any(0).mean()),
+                wp_sample_rate=float(wp_ok.mean()), ct_sample_rate=float(ct_ok.mean()),
+                hidden_collision_rate=float((wp_ok & ~ct_ok).sum() / max(1, wp_ok.sum())),
+                wp_min=float(wp_min.min()), ct_min=float(ct_min.min()), need=float(need),
+                ct_ok=ct_ok, n_samples=S, n_inits=I, n_agents=A, horizon=H)
+
+
+# ---- obstacles, exact for discs and boxes
+def _seg_box_dist(P0, P1, c, h, iters=100):
+    """Distance from each segment P0[k]->P1[k] to the box |x - c| <= h (componentwise).  The distance to a
+    convex set along a segment is convex in the segment parameter, so ternary search converges to the exact
+    minimum ((2/3)^100 of the segment after 100 iterations)."""
+    def dist(s):
+        X = P0 + s[:, None] * (P1 - P0)
+        q = np.maximum(np.abs(X - c) - h, 0.0)
+        return np.sqrt((q * q).sum(-1))
+    lo = np.zeros(len(P0)); hi = np.ones(len(P0))
+    for _ in range(iters):
+        m1 = lo + (hi - lo) / 3; m2 = hi - (hi - lo) / 3
+        left = dist(m1) <= dist(m2)
+        hi = np.where(left, m2, hi); lo = np.where(left, lo, m1)
+    return np.minimum.reduce([dist(lo), dist(hi), dist(np.zeros(len(P0))), dist(np.ones(len(P0)))])
+
+
+def exact_obstacle_clearance(path, discs=(), boxes=(), ws=None):
+    """path: (H, 2) waypoints of one robot.  Minimum centre-to-obstacle distance of the piecewise-linear
+    motion (discs: distance minus radius; boxes: distance to the box), and the minimum distance to the arena
+    edges.  Returns (obstacle clearance, arena clearance)."""
+    P0, P1 = path[:-1], path[1:]
+    best = np.inf
+    for c, r in discs:
+        D = P1 - P0; w = np.asarray(c, float) - P0
+        t = np.clip((w * D).sum(-1) / np.maximum((D * D).sum(-1), 1e-18), 0, 1)
+        best = min(best, float((np.linalg.norm(w - t[:, None] * D, axis=-1) - r).min()))
+    for c, h in boxes:
+        best = min(best, float(_seg_box_dist(P0, P1, np.asarray(c, float), np.asarray(h, float)).min()))
+    edge = np.inf
+    if ws is not None:
+        edge = float(np.minimum(path - ws[0], ws[1] - path).min())     # a segment stays inside iff its ends do
+    return best, edge
+
+
+def robotarium_exact(plans, mp, N, inst=None, data_dir=None):
+    """Exact offline check of returned Robotarium plans (I, N, H, 2), metres: robot-robot separation 2r over
+    every segment, robot body (r = 0.13) clear of every box, inside the arena.  Returns per-instance dict."""
+    from instances import load_robotarium
+    G = load_robotarium(mp, N, data_dir)
+    r = G['robot_r']; boxes = list(zip(G['box_c'], G['box_h']))
+    T = np.asarray(plans, float)
+    out = []
+    for i in range(T.shape[0]):
+        rr = min([float(pair_min_dist_exact(T[i, a], T[i, b])) for a in range(N) for b in range(a + 1, N)] or [np.inf])
+        oc, ed = np.inf, np.inf
+        for a in range(N):
+            o, e = exact_obstacle_clearance(T[i, a], boxes=boxes, ws=G['ws'])
+            oc = min(oc, o); ed = min(ed, e)
+        k = i if inst is None else inst
+        ge = float(np.linalg.norm(T[i, :, -1] - G['goals'][k], axis=-1).max())
+        se = float(np.linalg.norm(T[i, :, 0] - G['starts'][k], axis=-1).max())
+        out.append(dict(inst=k, rr_min=rr, obs_clear=oc - r, arena_clear=ed - r, goal_err=ge, start_err=se,
+                        free=bool(rr >= 2 * r and oc >= r and ed >= r)))
+    return out
+
+
+# ---- the benchmark's own predicate (benchmark environment only)
+def build_task(env_model_id, trained_models_dir, device='cpu', extra_objects=False):
+    import torch
+    from mmd.utils.loading import load_params_from_yaml
+    from mmd.trainer import get_dataset
+    args = load_params_from_yaml(os.path.join(trained_models_dir, env_model_id, 'args.yaml'))
+    trainset, _, _, _ = get_dataset(dataset_class='TrajectoryDataset', use_extra_objects=extra_objects,
+                                    obstacle_cutoff_margin=0.05, **args,
+                                    tensor_args=dict(device=device, dtype=torch.float32))
+    return trainset.dataset.task
+
+
+def obstacle_rates(task, pos, k, device='cpu'):
+    """pos: (S, I, A, H, 2) -> per-(sample, configuration) mask of obstacle-free plans at interpolation k."""
+    import torch
+    S, I, A, H, _ = pos.shape
+    T = torch.from_numpy(np.ascontiguousarray(pos.reshape(S * I * A, H, 2))).to(device=device, dtype=torch.float32)
+    _, _, _, _, wpc = task.get_trajs_collision_and_free(T, return_indices=True, num_interpolation=int(k))
+    m = wpc.cpu().numpy().reshape(S, I, A, -1)
+    return ~m.any(-1).any(-1)
+
+
+def suite_metrics(mp, results_dir, trained_models):
+    """The five table metrics of every saved cell of one map, with the benchmark's definitions: SR (some
+    candidate collision-free; here the one plan), CS (fraction of candidates), Coll (mean colliding robot pairs
+    at the waypoints), DA (the environment's data adherence), Time (median wall seconds per configuration).
+    Collision-free = robot pairs >= 2r at every waypoint and the benchmark's obstacle predicate at k = 0."""
+    import torch
+    from instances import OBS_SCALE, ENV_CLASS, suite_robot_radius
+    os.environ['OBS_SCALE'] = str(OBS_SCALE[mp])
+    TA = dict(device='cpu', dtype=torch.float32)
+    task = build_task(ENV_CLASS[mp] + '-RobotPlanarDisk', trained_models)
+    import torch_robotics.environments as _env_module
+    envobj = getattr(_env_module, ENV_CLASS[mp])(tensor_args=TA)
+    rob = suite_robot_radius(mp)
+    rows = []
+    for f in sorted(glob.glob(os.path.join(results_dir, f'{mp}__N*__*.npz'))):
+        D = np.load(f, allow_pickle=True); P = np.asarray(D['paths'], float)
+        B, I, A, H, _ = P.shape; pos = P[..., :2]
+        if A > 1:
+            dd = np.linalg.norm(pos[:, :, :, None, :, :] - pos[:, :, None, :, :, :], axis=-1)
+            iu = np.triu_indices(A, 1)
+            pairhit = (dd[:, :, iu[0], iu[1], :] < 2 * rob).any(-1)
+            coll = pairhit.sum(-1).astype(float); pair_ok = ~pairhit.any(-1)
+        else:
+            coll = np.zeros((B, I)); pair_ok = np.ones((B, I), bool)
+        feas = pair_ok & obstacle_rates(task, pos, 0)
+        da = envobj.compute_traj_data_adherence_batch(
+            torch.from_numpy(pos.reshape(-1, H, 2)).to(**TA)).cpu().numpy().reshape(B, I, A)
+        n = os.path.basename(f)[:-4].split('__')
+        r = dict(map=mp, N=A, regime=n[2], B=int(B), I=int(I),
+                 SR=100 * float(feas.any(0).mean()), CS=100 * float(feas.mean()), CS_sd=100 * float(feas.mean(0).std()),
+                 Coll=float(coll.mean()), Coll_sd=float(coll.std()), DA=float(da.mean()), DA_sd=float(da.std()),
+                 Time=float(np.median(D['secs'])), Time_sd=float(np.std(D['secs'])))
+        rows.append(r); print(json.dumps(r), flush=True)
+    return rows
+
+
+def cells_json(metrics_jsonl, results_dir, out_json):
+    """figures/gpa_cells.json records (fractions, median and std of wall time) from suite_metrics output."""
+    recs = []
+    for l in open(metrics_jsonl):
+        if not l.startswith('{'): continue
+        r = json.loads(l)
+        z = np.load(os.path.join(results_dir, f"{r['map']}__N{r['N']}__{r['regime']}.npz")); sec = np.asarray(z['secs'], float)
+        recs.append(dict(method='gpa', map=r['map'], N=r['N'], regime=r['regime'], src=os.path.basename(os.path.normpath(results_dir)),
+                         SR=r['SR'] / 100.0, CS=r['CS'] / 100.0, CS_sd=r['CS_sd'] / 100.0, Coll=r['Coll'], Coll_sd=r['Coll_sd'],
+                         DA=r['DA'], DA_sd=r['DA_sd'], T=float(np.median(sec)), T_sd=float(sec.std())))
+    json.dump(recs, open(out_json, 'w'), indent=0)
+    print(f'wrote {len(recs)} records to {out_json}')
+
+
+if __name__ == '__main__':
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    sub = ap.add_subparsers(dest='cmd', required=True)
+    s = sub.add_parser('suite'); s.add_argument('map'); s.add_argument('--results-dir', required=True)
+    s.add_argument('--trained-models', required=True, help="the benchmark's data_trained_models directory")
+    c = sub.add_parser('cells'); c.add_argument('metrics_jsonl'); c.add_argument('--results-dir', required=True)
+    c.add_argument('--out', required=True)
+    r = sub.add_parser('robotarium'); r.add_argument('map'); r.add_argument('N', type=int); r.add_argument('plan')
+    r.add_argument('--inst', type=int, default=None, help='the plan file holds one instance, benchmark index K')
+    a = ap.parse_args()
+    if a.cmd == 'suite':
+        suite_metrics(a.map, a.results_dir, a.trained_models)
+    elif a.cmd == 'cells':
+        cells_json(a.metrics_jsonl, a.results_dir, a.out)
+    else:
+        T = np.load(a.plan)['trajs']
+        res = robotarium_exact(T, a.map, a.N, a.inst)
+        for x in res: print(json.dumps(x))
+        print(json.dumps(dict(map=a.map, N=a.N, SR_exact=float(np.mean([x['free'] for x in res])),
+                              failed=[x['inst'] for x in res if not x['free']])))

applications/multi_robot/robotarium/figures.py

@@ -0,0 +1,261 @@
+#!/usr/bin/env python
+"""Multi-robot figures of the paper, from result files only (no solving).
+
+    # Sec. 4.2, figures/robotarium_sr_one_compact.pdf: single-plan success on the Robotarium arena
+    python figures.py robotarium --cells data/results/robotarium_cells.json --gpa data/results/robotarium_gpa_sr_one.json --out robotarium_sr_one_compact.pdf
+    # appendix, figures/hardsuite_single_plan{,_compact}.pdf and hardsuite_best64.pdf
+    python figures.py suite single  --baselines data/results/baseline_cells.json --gpa data/results/gpa_cells.json --out hardsuite_single_plan.pdf
+    python figures.py suite compact ... --out hardsuite_single_plan_compact.pdf
+    python figures.py suite best64  ... --out hardsuite_best64.pdf
+    # appendix, figures/hardsuite_velocity.pdf, read back from the generated tables
+    python figures.py velocity --tables hardsuite_tables.tex --out hardsuite_velocity.pdf [--n 9]
+
+GlobMP returns one trajectory per instance; for a sampling baseline, single-plan success is that of its first
+candidate (SR_at['1'] in the baseline records), best-of-64 is the tables' SR.  Every figure is written as .pdf
+and .png.
+"""
+import argparse
+import glob as _glob
+import json
+import re
+
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib import font_manager as _fm
+
+# ---- print style of the paper (true ICLR text width, Times, thin axes; GlobMP in blue)
+for _f in _glob.glob('/usr/share/fonts/truetype/msttcorefonts/[Tt]imes*.ttf'):
+    _fm.fontManager.addfont(_f)
+TEXTWIDTH = 5.5
+BLUE = '#2a78d6'
+INK, INK2, GRID = '#222222', '#5a5a5a', '#e9e9e6'
+plt.rcParams.update({
+    'font.family': 'serif', 'font.serif': ['Times New Roman', 'Nimbus Roman', 'STIXGeneral'], 'mathtext.fontset': 'stix',
+    'font.size': 8, 'axes.labelsize': 8, 'axes.titlesize': 8, 'xtick.labelsize': 7, 'ytick.labelsize': 7, 'legend.fontsize': 7,
+    'axes.linewidth': 0.6, 'axes.edgecolor': INK2, 'axes.labelcolor': INK, 'xtick.color': INK2, 'ytick.color': INK2,
+    'xtick.major.width': 0.6, 'ytick.major.width': 0.6, 'xtick.major.size': 2.5, 'ytick.major.size': 2.5,
+    'xtick.minor.size': 1.5, 'ytick.minor.size': 1.5, 'xtick.minor.width': 0.4, 'ytick.minor.width': 0.4,
+    'axes.spines.top': False, 'axes.spines.right': False, 'axes.axisbelow': True,
+    'lines.linewidth': 1.1, 'lines.markersize': 3.8, 'errorbar.capsize': 1.6,
+    'legend.frameon': False, 'legend.handlelength': 2.2, 'legend.columnspacing': 1.2, 'legend.handletextpad': 0.5,
+    'pdf.fonttype': 42, 'ps.fonttype': 42, 'savefig.dpi': 300, 'savefig.bbox': 'tight', 'savefig.pad_inches': 0.02,
+})
+OURS = 'GlobMP (Ours)'
+STY = {'pdm': dict(color='#7f7f7f', marker='h', label='PDM'), 'cd': dict(color='#8c564b', marker='P', label='CD'),
+       'mmd_cbs': dict(color='#17becf', marker='<', label='MMD-CBS'), 'mdoc': dict(color='#bcbd22', marker='>', label='MDOC'),
+       'diffusion_policy': dict(color='#9467bd', marker='x', label='Diff. Policy'),
+       'pcd_lb': dict(color='#ff7f0e', marker='v', label='PCD-LB'), 'pcd_shd': dict(color='#2ca02c', marker='s', label='PCD-SHD'),
+       'final_projection': dict(color='#e377c2', marker='D', label='Final Proj.'), 'direct': dict(color='#d62728', marker='o', label='DiRecT')}
+
+
+def _save(fig, out):
+    for ext in ('.pdf', '.png'):
+        fig.savefig(out.rsplit('.', 1)[0] + ext, dpi=300)
+
+
+def _mfc(st):
+    return 'white' if st['marker'] not in ('x', 'P') else st['color']
+
+
+# ---- Sec. 4.2: Robotarium
+def robotarium(cells_json, gpa_json, out):
+    """One row of five map panels: single-plan success against team size."""
+    MAPS = ['empty', 'corridor', 'shelf', 'corridorbay', 'crossroom']
+    TITLES = {'empty': 'Empty', 'corridor': 'Conveyor', 'shelf': 'Drop-Region', 'corridorbay': 'Corridor-Bay', 'crossroom': 'Cross-Room'}
+    NS = [3, 6, 9, 12]
+    ORDER = ['cd', 'mmd_cbs', 'pcd_lb', 'pcd_shd', 'final_projection', 'direct']
+    C = {(r['method'], r['map'], r['N']): r.get('SR_one') for r in json.load(open(cells_json))}
+    G = {(r['map'], r['N']): r.get('SR_one', r.get('SR')) for r in json.load(open(gpa_json))}
+    order = [m for m in ORDER if any(C.get((m, mp, n)) is not None for mp in MAPS for n in NS)]
+    FH = 1.35
+    fig, axes = plt.subplots(1, len(MAPS), figsize=(TEXTWIDTH, FH), sharex=True, sharey=True)
+    fig.subplots_adjust(left=0.07, right=0.995, top=1 - 0.42 / FH, bottom=0.35 / FH, wspace=0.12)
+    H = []
+    for k, mp in enumerate(MAPS):
+        ax = axes[k]
+        for m in order:
+            y = np.array([np.nan if C.get((m, mp, n)) is None else C[m, mp, n] for n in NS], float)
+            if np.all(np.isnan(y)): continue
+            st = STY[m]
+            h, = ax.plot(NS, y, color=st['color'], marker=st['marker'], ms=3.1, lw=0.9, mew=0.6, mfc=_mfc(st), label=st['label'])
+            if k == 0: H.append(h)
+        yg = np.array([np.nan if G.get((mp, n)) is None else G[mp, n] for n in NS], float)
+        ax.plot(NS, yg, color='white', lw=3.2, zorder=5.5, solid_capstyle='round')     # casing where lines cross
+        h, = ax.plot(NS, yg, color=BLUE, marker='o', ms=3.8, lw=1.7, label=OURS, zorder=6)
+        if k == 0: H.insert(0, h)
+        ax.set_title(TITLES[mp], fontsize=8, pad=2.5, color=INK)
+        ax.set_ylim(-0.05, 1.05); ax.set_yticks([0, 0.5, 1]); ax.set_xticks(NS)
+        ax.grid(axis='y', color=GRID, lw=0.5); ax.set_axisbelow(True)
+    fig.supxlabel('number of robots', fontsize=8, y=0.015)
+    fig.supylabel('single-plan success', fontsize=7, x=0.012, y=(1 - 0.42 / FH + 0.35 / FH) / 2)
+    fig.legend(H, [h.get_label() for h in H], loc='upper center', bbox_to_anchor=(0.53, 1.0), ncol=len(H), frameon=False,
+               fontsize=7, handlelength=1.8, columnspacing=0.9)
+    _save(fig, out); print('wrote', out)
+
+
+# ---- appendix: simulated suite
+SUITE_MAPS = [('empty', 'Empty'), ('basic', 'Basic'), ('dense', 'Dense'), ('corridor', 'Conveyor'), ('room', 'Highways'),
+              ('shelf', 'Drop-Region'), ('corridorbay', 'Corridor-Bay'), ('crossroom', 'Cross-Room')]
+VMAX = {'empty': (0.483, 0.571, 0.675), 'basic': (0.477, 0.567, 0.675), 'dense': (0.565, 0.617, 0.675), 'room': (0.647, 0.783, 0.878),
+        'corridor': (1.21, 1.46, 1.76), 'corridorbay': (1.21, 1.46, 1.76), 'shelf': (0.928, 1.13, 1.34), 'crossroom': (0.928, 1.13, 1.34)}
+REGS = ['low', 'medium', 'high']; SNS = [3, 6, 9, 12, 15, 18, 21]
+ORDER_ALL = ['pdm', 'cd', 'mmd_cbs', 'mdoc', 'diffusion_policy', 'pcd_lb', 'pcd_shd', 'final_projection', 'direct']
+
+
+def _sr_at(r, b):
+    if not r or not r.get('SR_at'): return np.nan
+    s = r['SR_at']; ks = sorted(int(k) for k in s); bb = max([k for k in ks if k <= b] or [ks[0]]); return s[str(bb)]
+
+
+def _metric(r, kind):
+    if r is None: return np.nan
+    if kind == 'single': return _sr_at(r, 1)
+    return r.get('SR', np.nan) if r.get('SR') is not None else np.nan
+
+
+def suite(kind, baselines_json, gpa_json, out):
+    """single / best64: one row per map, one column per velocity cap (all 168 cells).  compact: one panel per
+    map, single-plan success averaged over the three caps."""
+    C = {(r['method'], r['map'], r['N'], r['regime']): r for r in json.load(open(baselines_json))}
+    G = {(r['map'], r['N'], r['regime']): r for r in json.load(open(gpa_json))}
+    mk = 'single' if kind in ('single', 'compact') else 'best64'
+    order = [m for m in ORDER_ALL if any(not np.isnan(_metric(C.get((m, mp, n, rg)), mk)) for mp, _ in SUITE_MAPS for n in SNS for rg in REGS)]
+    if kind == 'compact':
+        fig, axes = plt.subplots(2, 4, figsize=(TEXTWIDTH, 2.95), sharex=True, sharey=True)
+        fig.subplots_adjust(left=0.07, right=0.995, top=0.86, bottom=0.115, wspace=0.12, hspace=0.34)
+        H = []
+        for k, (mp, title) in enumerate(SUITE_MAPS):
+            ax = axes.ravel()[k]
+            for m in order:
+                y = np.nanmean([[_metric(C.get((m, mp, n, rg)), 'single') for n in SNS] for rg in REGS], axis=0)
+                if np.all(np.isnan(y)): continue
+                st = STY[m]; h, = ax.plot(SNS, y, color=st['color'], marker=st['marker'], ms=3.1, lw=0.9, mew=0.6, mfc=_mfc(st), label=st['label'])
+                if k == 0: H.append(h)
+            yg = np.nanmean([[(G.get((mp, n, rg)) or {}).get('SR', np.nan) for n in SNS] for rg in REGS], axis=0)
+            ax.plot(SNS, yg, color='white', lw=3.2, zorder=5.5, solid_capstyle='round')
+            h, = ax.plot(SNS, yg, color=BLUE, marker='o', ms=3.8, lw=1.7, label=OURS, zorder=6)
+            if k == 0: H.append(h)
+            ax.set_title(title, fontsize=8, pad=2.5, color=INK)
+            ax.set_ylim(-0.05, 1.05); ax.set_yticks([0, 0.5, 1]); ax.set_xticks([3, 9, 15, 21])
+            ax.grid(axis='y', color=GRID, lw=0.5); ax.set_axisbelow(True)
+        fig.supxlabel('number of agents', fontsize=8, y=0.015); fig.supylabel('single-plan success', fontsize=8, x=0.012)
+        H = [H[-1]] + H[:-1]
+        fig.legend(H, [h.get_label() for h in H], loc='upper center', bbox_to_anchor=(0.53, 1.0), ncol=len(H), frameon=False,
+                   fontsize=7, handlelength=1.8, columnspacing=0.9)
+        _save(fig, out); print('wrote', out); return
+    fig, axes = plt.subplots(len(SUITE_MAPS), 3, figsize=(TEXTWIDTH, 8.3), sharex=True, sharey=True)
+    for ri, (mp, title) in enumerate(SUITE_MAPS):
+        for ci, rg in enumerate(REGS):
+            ax = axes[ri, ci]
+            for m in order:
+                y = [_metric(C.get((m, mp, n, rg)), mk) for n in SNS]
+                if all(np.isnan(y)): continue
+                st = STY[m]; ax.plot(SNS, y, color=st['color'], marker=st['marker'], ms=2.4, lw=0.8, mew=0.5, mfc=_mfc(st), label=st['label'])
+            yg = [(G.get((mp, n, rg)) or {}).get('SR', np.nan) for n in SNS]
+            ax.plot(SNS, yg, color=BLUE, marker='o', ms=2.8, lw=1.5, label=OURS, zorder=6)
+            ax.set_xticks([3, 9, 15, 21]); ax.set_ylim(-0.05, 1.05); ax.set_yticks([0, 0.5, 1]); ax.grid(axis='y', color=GRID, lw=0.5); ax.set_axisbelow(True)
+            ax.set_title(f"{title}, $v_{{\\max}}={VMAX[mp][ci]:g}$ ({rg})", fontsize=7, pad=2, color=INK)
+    fig.supxlabel('number of agents', fontsize=8, y=0.005)
+    fig.supylabel({'single': 'single-plan success', 'best64': 'success, best of 64'}[kind], fontsize=8, x=0.005)
+    fig.subplots_adjust(left=0.07, right=0.995, top=0.945, bottom=0.045, wspace=0.08, hspace=0.42)
+    h, l = axes[0, 1].get_legend_handles_labels(); h = [h[-1]] + h[:-1]; l = [l[-1]] + l[:-1]
+    fig.legend(h, l, loc='upper center', ncol=min(len(l), 6), frameon=False, fontsize=7, bbox_to_anchor=(0.5, 1.0), handlelength=1.8, columnspacing=0.9)
+    _save(fig, out); print('wrote', out)
+
+
+# ---- appendix: velocity axis, read back from the generated tables
+def parse_tables(path):
+    """Rows of hardsuite_tables.tex: N & Method & (SR CS Coll DA Time) x 3 caps; bolding and deviations stripped,
+    SR and CS back to fractions."""
+    s = open(path).read()
+    out = []
+    for t in re.findall(r'\\begin\{table\}.*?\\end\{table\}', s, re.S):
+        lab = re.search(r'\\label\{tab:hard-([a-z]+)\}', t).group(1)
+        caps = re.findall(r'\\multicolumn\{5\}\{c\}\{(Low|Medium|High) \(\$v_\{\\max\}=([0-9.]+)\$\)\}', t)
+        regimes = [(r.lower(), float(v)) for r, v in caps]
+        cur_n = None
+        for line in t.split(r'\midrule', 1)[1].split('\\\\'):
+            line = line.replace(r'\midrule', '').replace(r'\bottomrule', '').replace(r'\end{tabular}', '').strip()
+            if not line or '&' not in line: continue
+            cells = [c.strip() for c in line.split('&')]
+            if len(cells) != 17: continue
+            if cells[0]: cur_n = int(cells[0])
+            meth = re.sub(r'\\textbf\{(.*?)\}', r'\1', cells[1]).strip()
+            meth = re.sub(r'\$\^\{\\dagger\}\$', '', meth).replace('\\ ', ' ').strip()
+
+            def num(x, pct=False):
+                x = re.sub(r'\$_\{\\pm[^}]*\}\$', '', x); x = re.sub(r'\\textbf\{(.*?)\}', r'\1', x).strip()
+                try: v = float(x)
+                except ValueError: return None
+                return v / 100.0 if pct else v
+            vals = cells[2:]
+            for ri, (rname, vcap) in enumerate(regimes):
+                sr, cs, coll, da, tm = vals[ri * 5:(ri + 1) * 5]
+                out.append(dict(map=lab, N=cur_n, method=meth, regime=rname, vcap=vcap, SR=num(sr, True), CS=num(cs, True),
+                                Coll=num(coll), DA=num(da), T=num(tm)))
+    return out
+
+
+VSTYLE = {'CD': dict(color='#8c564b', marker='P', ls='-'), 'PCD-LB': dict(color='#ff7f0e', marker='v', ls=':'),
+          'PCD-SHD': dict(color='#2ca02c', marker='s', ls='--'), 'Final Proj.': dict(color='#e377c2', marker='D', ls='-.'),
+          'DiRecT': dict(color='#d62728', marker='o', ls='-'), 'MDOC': dict(color='#9467bd', marker='h', ls=(0, (4, 1.5))),
+          OURS: dict(color=BLUE, marker='o', ls='-')}
+
+
+def velocity(tables, out, n=9):
+    """Success and constraint safety against the velocity cap at N = n, two panels per map.  Methods with no
+    velocity mechanism (MMD-CBS, Diffusion Policy) are omitted."""
+    rows = parse_tables(tables)
+    NO_AXIS = {'MMD-CBS', 'Diff. Policy', 'Diffusion Policy'}
+    ORDER = [OURS, 'CD', 'PCD-LB', 'PCD-SHD', 'MDOC', 'Final Proj.', 'DiRecT']
+    METRICS = [('SR', 'success'), ('CS', 'safety')]
+    series = {}
+    for r in rows:
+        if r['N'] != n or r['method'] in NO_AXIS: continue
+        series.setdefault((r['map'], r['method']), []).append((r['vcap'], r['SR'], r['CS']))
+    for k in series: series[k].sort(key=lambda t: t[0])
+    PER_ROW = 2; ncol = PER_ROW * len(METRICS); nrow = (len(SUITE_MAPS) + PER_ROW - 1) // PER_ROW
+    fig, axes = plt.subplots(nrow, ncol, figsize=(TEXTWIDTH, 1.42 * nrow + 0.45), squeeze=False, sharey=True)
+    plt.subplots_adjust(left=0.06, right=0.995, top=0.945, bottom=0.07, hspace=0.62, wspace=0.12)
+    drawn = []
+    for mi, (key, disp) in enumerate(SUITE_MAPS):
+        ri, slot = divmod(mi, PER_ROW)
+        for mj, (metric, mlabel) in enumerate(METRICS):
+            ax = axes[ri][slot * len(METRICS) + mj]
+            for meth in ORDER:
+                pts = series.get((key, meth))
+                if not pts: continue
+                keep = [(p[0], p[1 + mj]) for p in pts if p[1 + mj] is not None]
+                if len(keep) < 2: continue
+                st = VSTYLE.get(meth, {}); ours = meth == OURS
+                ax.plot([k[0] for k in keep], [k[1] for k in keep], color=st.get('color'), marker=st.get('marker', 'o'),
+                        ls=st.get('ls', '-'), lw=1.6 if ours else 0.8, ms=2.8 if ours else 2.2, mew=0.5,
+                        mfc=st.get('color') if ours else 'white', zorder=6 if ours else 4, label=meth)
+                if meth not in drawn: drawn.append(meth)
+            ax.set_ylim(-0.04, 1.04); ax.set_yticks([0, 0.5, 1.0])
+            ax.tick_params(labelsize=6, length=2.2, pad=1.5)
+            ax.grid(axis='y', color=GRID, lw=0.5); ax.set_axisbelow(True)
+            ax.set_title('%s, %s' % (disp, mlabel), fontsize=7, pad=2.5, color=INK)
+            ax.set_xlabel(r'$v_{\max}$', fontsize=6.5, labelpad=1)
+    handles = [plt.Line2D([], [], color=VSTYLE[m]['color'], marker=VSTYLE[m]['marker'], ls=VSTYLE[m]['ls'],
+                          lw=1.6 if m == OURS else 0.8, ms=2.8, mew=0.5, mfc=VSTYLE[m]['color'] if m == OURS else 'white', label=m)
+               for m in ORDER if m in drawn]
+    fig.supylabel('success / safety', fontsize=7, x=0.005)
+    fig.legend(handles=handles, loc='upper center', ncol=len(handles), fontsize=7, frameon=False, handlelength=1.6,
+               handletextpad=0.4, columnspacing=1.2, borderaxespad=0.1, bbox_to_anchor=(0.5, 1.0))
+    _save(fig, out); print('wrote', out, f'(N={n}, {len(drawn)} methods)')
+
+
+if __name__ == '__main__':
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    sub = ap.add_subparsers(dest='cmd', required=True)
+    r = sub.add_parser('robotarium'); r.add_argument('--cells', required=True); r.add_argument('--gpa', required=True); r.add_argument('--out', required=True)
+    s = sub.add_parser('suite'); s.add_argument('kind', choices=['single', 'best64', 'compact'])
+    s.add_argument('--baselines', required=True); s.add_argument('--gpa', required=True); s.add_argument('--out', required=True)
+    v = sub.add_parser('velocity'); v.add_argument('--tables', required=True); v.add_argument('--out', required=True); v.add_argument('--n', type=int, default=9)
+    a = ap.parse_args()
+    if a.cmd == 'robotarium': robotarium(a.cells, a.gpa, a.out)
+    elif a.cmd == 'suite': suite(a.kind, a.baselines, a.gpa, a.out)
+    else: velocity(a.tables, a.out, a.n)

applications/multi_robot/robotarium/instances.py

@@ -0,0 +1,185 @@
+"""Start-goal instances and map geometry for the multi-robot experiments.
+
+Two settings:
+
+* the simulated suite (appendix, app:exp-multirobot): 8 maps x N in {3,...,21}, 20 shared start-goal
+  configurations per (map, N), 3 velocity caps per map.  One file per (map, N) in data/suite/, written by
+  `export_suite` from the benchmark's shared-instance file and its own environment classes.
+* the Robotarium setting (Sec. 4.2): 5 maps fitted to the 3.2 x 2.0 m arena, N in {3, 6, 9, 12}, 10 instances
+  per (map, N) (nine sampled, one antipodal ring), robot radius 0.13 m.  data/robotarium/instances.npz (keys
+  '<map>rbt/N<n>/{starts,goals}', world metres), maps.json and geometry.json (axis-aligned boxes in metres).
+
+Map keys follow the benchmark's code names; the paper's names are in PAPER_NAME.
+"""
+import json
+import os
+
+import numpy as np
+
+DATA = os.environ.get('GLOBMP_MR_DATA', os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data'))
+
+PAPER_NAME = {'empty': 'Empty', 'basic': 'Basic', 'dense': 'Dense', 'corridor': 'Conveyor', 'room': 'Highways',
+              'shelf': 'Drop-Region', 'corridorbay': 'Corridor-Bay', 'crossroom': 'Cross-Room'}
+
+# ---- simulated suite
+SUITE_MAPS = ['empty', 'basic', 'dense', 'corridor', 'room', 'shelf', 'corridorbay', 'crossroom']
+SUITE_NS = [3, 6, 9, 12, 15, 18, 21]
+REGIMES = ['low', 'medium', 'high']
+OBS_SCALE = {'empty': 1.0, 'basic': 2.2, 'dense': 1.0, 'corridor': 1.5, 'room': 1.45, 'shelf': 1.6,
+             'corridorbay': 1.15, 'crossroom': 1.1}
+ENV_CLASS = {'empty': 'EnvEmptyNoWait2D', 'basic': 'EnvEmptyNoWait2D', 'dense': 'EnvEmptyNoWait2D',
+             'corridor': 'EnvConveyor2D', 'room': 'EnvHighways2D', 'shelf': 'EnvDropRegion2D',
+             'corridorbay': 'EnvCorridorSMD2D', 'crossroom': 'EnvCrossRoom2D'}
+# velocity caps (low, medium, high); must equal VCAP in tables.py
+VCAP = {'empty': (0.483, 0.571, 0.675), 'basic': (0.477, 0.567, 0.675), 'dense': (0.565, 0.617, 0.675),
+        'room': (0.647, 0.783, 0.878), 'corridor': (1.21, 1.46, 1.76), 'corridorbay': (1.21, 1.46, 1.76),
+        'shelf': (0.928, 1.13, 1.34), 'crossroom': (0.928, 1.13, 1.34)}
+
+
+def suite_robot_radius(mp):
+    return 0.1 if mp == 'empty' else 0.05
+
+
+def load_suite(mp, N, inst_dir=None):
+    """dict with centers (K,2), radii (K,), ws (2,2) [lo; hi], box_c, box_h, starts (I,N,2), goals (I,N,2),
+    horizon (64), robot_r and the three velocity caps vmax_low / vmax_med / vmax_high."""
+    f = os.path.join(inst_dir or os.path.join(DATA, 'suite'), f'{mp}_N{N}.npz')
+    D = np.load(f, allow_pickle=True)
+    return {k: D[k] for k in D.files}
+
+
+# ---- Robotarium setting
+RBT_MAPS = ['empty', 'corridor', 'shelf', 'corridorbay', 'crossroom']
+RBT_NS = [3, 6, 9, 12]
+RBT_ROBOT_R = 0.13
+RBT_ARENA = np.array([[-1.6, -1.0], [1.6, 1.0]])
+
+
+def load_robotarium(mp, N, data_dir=None):
+    """dict with starts, goals (10, N, 2) in metres, boxes (centers (K,2), half sizes (K,2)), ws, robot_r."""
+    d = data_dir or os.path.join(DATA, 'robotarium')
+    I = np.load(os.path.join(d, 'instances.npz'))
+    geo = json.load(open(os.path.join(d, 'geometry.json')))['maps'][mp]
+    BC = np.array([o['center'] for o in geo['obstacles']], float).reshape(-1, 2)
+    BH = np.array([o['size'] for o in geo['obstacles']], float).reshape(-1, 2) / 2.0
+    return dict(starts=I[f'{mp}rbt/N{N}/starts'], goals=I[f'{mp}rbt/N{N}/goals'], box_c=BC, box_h=BH,
+                ws=RBT_ARENA.copy(), robot_r=RBT_ROBOT_R, obs_scale=geo['obs_scale'],
+                wide_corridor=bool(geo.get('wide_corridor', False)), env_class=geo['env_class'])
+
+
+# ---- geometry helpers
+def boxes_to_circles(centers, half_sizes, halo=0.02, max_circles_per_box=400):
+    """Cover each axis-aligned box by discs whose union contains it, so avoiding every disc implies avoiding
+    the box.  The union overshoots the box by (sqrt(2)-1) times the cell half-size, so the box is subdivided
+    until that halo is at most `halo`."""
+    out = []
+    C = np.asarray(centers, dtype=float).reshape(-1, 2)
+    Hs = np.asarray(half_sizes, dtype=float).reshape(-1, 2)
+    for c, h in zip(C, Hs):
+        hx, hy = float(h[0]), float(h[1])
+        if hx <= 1e-9 or hy <= 1e-9:
+            continue
+        cell = halo / (np.sqrt(2.0) - 1.0)
+        nx = max(1, int(np.ceil(hx / cell)))
+        ny = max(1, int(np.ceil(hy / cell)))
+        while nx * ny > max_circles_per_box:
+            nx = max(1, nx - 1); ny = max(1, ny - 1)
+        rx, ry = hx / nx, hy / ny
+        r = float(np.hypot(rx, ry))
+        xs = (np.arange(nx) + 0.5) * 2 * rx - hx
+        ys = (np.arange(ny) + 0.5) * 2 * ry - hy
+        for xq in xs:
+            for yq in ys:
+                out.append((c + np.array([xq, yq]), r))
+    return out
+
+
+def box_discs(BC, BH, halo, halo_wall=0.003, aspect=3.0):
+    """Discs whose union contains each box (the Robotarium projector's obstacle model).  Square-ish boxes: the
+    grid cover above with overshoot <= `halo`.  Thin walls (aspect >= 3): ONE row of discs of radius about the
+    wall's half-thickness, spaced so the long faces overshoot by <= `halo_wall` (a 2 cm halo on a wall would
+    narrow the 0.29 m Conveyor channel below the 0.26 m a robot needs)."""
+    C, R = [], []
+    for c, h in zip(np.asarray(BC, float).reshape(-1, 2), np.asarray(BH, float).reshape(-1, 2)):
+        L, t = (0, 1) if h[0] >= h[1] else (1, 0)          # long axis, thin axis
+        if h[L] >= aspect * h[t]:
+            rl = np.sqrt((h[t] + halo_wall) ** 2 - h[t] ** 2)  # half-spacing along the wall
+            nl = int(np.ceil(h[L] / rl)); rl = h[L] / nl
+            rad = float(np.hypot(rl, h[t]))
+            for k in range(nl):
+                p = np.array(c, float); p[L] += -h[L] + (2 * k + 1) * rl
+                C.append(p); R.append(rad)
+        else:
+            for cc, rr in boxes_to_circles(c[None], h[None], halo=halo):
+                C.append(np.asarray(cc, float)); R.append(float(rr))
+    return np.array(C).reshape(-1, 2), np.array(R).reshape(-1)
+
+
+def benchmark_geometry(task, halo):
+    """Obstacle discs and workspace in world coordinates, read off the benchmark environment of `task`:
+    sphere fields as they are, box fields covered by boxes_to_circles(halo).  Returns (circles, ws, BC, BH)."""
+    env = task.env
+    lim = env.limits.cpu().numpy()
+    ws = np.stack([lim[0, :2], lim[1, :2]])
+    C, Hs, Sph = [], [], []
+    for obj in (env.obj_fixed_list or []):
+        for fld in (getattr(obj, 'fields', None) or []):
+            cen = getattr(fld, 'centers', None)
+            if cen is None or len(cen) == 0:
+                continue
+            cen = np.asarray(cen.cpu(), dtype=float)[:, :2]
+            sz = getattr(fld, 'sizes', None)
+            if sz is not None:                                   # box field: sizes are full extents
+                C.append(cen); Hs.append(np.asarray(sz.cpu(), dtype=float)[:, :2] / 2.0)
+            else:
+                rad = np.asarray(getattr(fld, 'radii', getattr(fld, 'radius')).cpu(), dtype=float).reshape(-1)
+                Sph += [(c, float(r)) for c, r in zip(cen, rad)]
+    circles = list(Sph)
+    if C:
+        circles += boxes_to_circles(np.concatenate(C), np.concatenate(Hs), halo=halo)
+    BC = np.concatenate(C) if C else np.zeros((0, 2)); BH = np.concatenate(Hs) if Hs else np.zeros((0, 2))
+    return circles, ws, BC, BH
+
+
+def export_suite(shared_npz, trained_models_dir, out_dir, maps=None):
+    """Write data/suite/<map>_N<n>.npz from the benchmark's shared instances (the same configurations every
+    method plans).  Runs in the benchmark's environment (needs mmd and torch_robotics on the path).  The disc
+    cover of the box obstacles uses a halo of half the endpoints' clearance, clipped to [0.004, 0.030]."""
+    from evaluate import build_task
+    os.makedirs(out_dir, exist_ok=True)
+    S = np.load(shared_npz, allow_pickle=True)
+    for mp in (maps or SUITE_MAPS):
+        os.environ['OBS_SCALE'] = str(OBS_SCALE[mp])        # read by the benchmark at environment construction
+        rr = suite_robot_radius(mp)
+        t = build_task(ENV_CLASS[mp] + '-RobotPlanarDisk', trained_models_dir)
+        P0 = np.concatenate([S[f'{mp}/N3/starts'].reshape(-1, 2), S[f'{mp}/N3/goals'].reshape(-1, 2)])
+        c_fine, _, _, _ = benchmark_geometry(t, 0.002)
+        if c_fine:
+            Cf = np.array([c for c, _ in c_fine]); Rf = np.array([r for _, r in c_fine])
+            clear = float((np.linalg.norm(P0[:, None, :] - Cf[None], axis=2) - (Rf[None] + rr)).min())
+        else:
+            clear = 1.0
+        halo = float(min(0.030, max(0.004, 0.5 * clear)))
+        circ, ws, BC, BH = benchmark_geometry(t, halo)
+        C = np.array([c for c, _ in circ]) if circ else np.zeros((0, 2))
+        Rd = np.array([r for _, r in circ]) if circ else np.zeros((0,))
+        for n in SUITE_NS:
+            k = f'{mp}/N{n}'
+            if k + '/starts' not in S.files:
+                print('# missing', k, flush=True); continue
+            np.savez(os.path.join(out_dir, f'{mp}_N{n}.npz'), centers=C, radii=Rd, ws=ws, box_c=BC, box_h=BH,
+                     starts=S[k + '/starts'], goals=S[k + '/goals'], agents=n, horizon=64,
+                     robot_r=rr, obs_scale=OBS_SCALE[mp], vmax_low=VCAP[mp][0], vmax_med=VCAP[mp][1],
+                     vmax_high=VCAP[mp][2], map=mp, env_class=ENV_CLASS[mp])
+        print(f'# {mp}: {len(circ)} discs, {len(BC)} boxes, r={rr}, halo={halo:.4f}, endpoint clearance {clear:+.4f}', flush=True)
+
+
+if __name__ == '__main__':
+    import argparse
+    ap = argparse.ArgumentParser(description='Export the simulated-suite instances from the benchmark (benchmark env).')
+    ap.add_argument('--shared-npz', required=True, help="the benchmark's shared instance file (instances_shared.npz)")
+    ap.add_argument('--trained-models', required=True, help="the benchmark's data_trained_models directory")
+    ap.add_argument('--out', default=os.path.join(DATA, 'suite'))
+    ap.add_argument('--maps', nargs='*')
+    a = ap.parse_args()
+    export_suite(a.shared_npz, a.trained_models, a.out, a.maps)

applications/multi_robot/robotarium/projector.py

@@ -0,0 +1,535 @@
+"""GlobMP as the projector inside a diffusion sampler (the BaseProjector interface of DiRecT).
+
+This is the configuration that produced the GlobMP Robotarium results (figures/robotarium_gpa_sr_one.json):
+the sampler's single-agent diffusion models propose a batch of team plans, and at every projection call the
+proposal of each candidate team is the piecewise-linear reference of a tracking cost.  GlobMP restores the
+team (Algorithm 1) and descends on the tracking cost (Algorithm 2) under continuous-time blocks: pairwise
+separation 2r, a uniform speed bound, the workspace box and the obstacle discs.  The waypoints handed back
+are samples of a curve that satisfies every block at every t.
+
+Batch layout (DiRecT): B = n_agents * G samples, agent-major, so sample a*G+j is robot a in team j.
+Index 0 of each plan is pinned by params['pos_init'].
+
+Schedule over the denoising steps (GPA_NCALLS calls per instance): intermediate calls only steer (short
+restoration, no descent, partially restored teams are handed back); the last call is the full solve and
+returns certified teams only.  A team whose piecewise-linear motion already satisfies every constraint
+exactly is left alone.  With GPA_PRIORITIZED=1, if no team certifies jointly at the last call, the teams are
+certified one robot at a time (each robot against the robots already placed, as moving obstacles), with
+wall, timing and detour repairs of the reference; with GPA_FIRST_CERT=1 the first certified team ends the
+call, since the method returns one plan.
+
+Every knob is an environment variable (read inside the sampler process); the values used for the paper are
+in run/run_cell.sh and the README.
+"""
+import os
+import time
+
+import numpy as np
+import torch
+
+from globmp import Problem, init_waypoints, spline as bern
+
+DT = torch.float64
+
+
+def seg_min_dist(A, B):
+    """Exact minimum distance between two robots moving linearly between waypoints. A, B: (..., H, 2)."""
+    P = A[..., :-1, :] - B[..., :-1, :]
+    D = (A[..., 1:, :] - A[..., :-1, :]) - (B[..., 1:, :] - B[..., :-1, :])
+    dd = (D * D).sum(-1)
+    s = np.clip(np.where(dd > 1e-18, -(P * D).sum(-1) / np.maximum(dd, 1e-18), 0.0), 0.0, 1.0)
+    Q = P + s[..., None] * D
+    return np.sqrt((Q * Q).sum(-1)).min(-1)
+
+
+def _static_route(w, discs, clr, lo, hi, res=0.01, lam=4.0):
+    """Wall-free route for ONE robot from w[0] to w[-1] on a grid over [lo, hi]: A* with step cost = length +
+    lam * distance to the robot's own reference path, so the route follows the reference where the reference
+    is feasible.  A cell is free if every obstacle disc is at least `clr` away.  Returns the reference re-timed
+    onto the route (same progress-versus-time profile), or None if no route exists."""
+    import heapq
+    nx = int(np.ceil((hi[0] - lo[0]) / res)) + 1; ny = int(np.ceil((hi[1] - lo[1]) / res)) + 1
+    xs = lo[0] + res * np.arange(nx); ys = lo[1] + res * np.arange(ny)
+    X, Y = np.meshgrid(xs, ys, indexing='ij'); Pg = np.stack([X, Y], -1)
+    free = np.ones((nx, ny), bool)
+    for c, r in discs:
+        free &= np.linalg.norm(Pg - np.asarray(c, float), axis=-1) >= float(r) + clr
+    dref = np.full((nx, ny), np.inf)
+    for q in w[::2]: dref = np.minimum(dref, np.linalg.norm(Pg - q, axis=-1))
+    cell = lambda p: (int(np.clip(round((p[0] - lo[0]) / res), 0, nx - 1)), int(np.clip(round((p[1] - lo[1]) / res), 0, ny - 1)))
+    s, g = cell(w[0]), cell(w[-1])
+    free[s] = free[g] = True                                   # the benchmark's endpoints are feasible
+    nb = [(1, 0), (-1, 0), (0, 1), (0, -1), (1, 1), (1, -1), (-1, 1), (-1, -1)]
+    dist = {s: 0.0}; prev = {}; pq = [(0.0, s)]; gx, gy = g
+    while pq:
+        f_, u = heapq.heappop(pq)
+        if u == g: break
+        du = dist[u]
+        for dx_, dy_ in nb:
+            v = (u[0] + dx_, u[1] + dy_)
+            if not (0 <= v[0] < nx and 0 <= v[1] < ny) or not free[v]: continue
+            step = res * (1.4142 if dx_ and dy_ else 1.0)
+            nd = du + step * (1.0 + lam * dref[v])
+            if nd < dist.get(v, np.inf):
+                dist[v] = nd; prev[v] = u
+                heapq.heappush(pq, (nd + res * np.hypot(v[0] - gx, v[1] - gy), v))
+    if g not in dist: return None
+    path = [g]
+    while path[-1] != s: path.append(prev[path[-1]])
+    path = np.array([[xs[i], ys[j]] for i, j in path[::-1]]); path[0] = w[0]; path[-1] = w[-1]
+    arc = lambda P_: np.r_[0, np.cumsum(np.linalg.norm(np.diff(P_, axis=0), axis=1))]
+    ar, ap = arc(w), arc(path)
+    frac = ar / max(ar[-1], 1e-12)
+    return np.stack([np.interp(frac * ap[-1], ap, path[:, k]) for k in range(2)], 1)
+
+
+class GPAProjector:
+    """Implements BaseProjector.project(x, params) -> x_proj."""
+
+    def __init__(self, n_agents, robot_radius, ws_bounds=None, obstacles=None,
+                 degree=None, segments=None, ell_j=2, ell=0, elev=None, margin=0.0,
+                 alg1_cap=None, alg2_blocks=None, alg2_steps=3, max_teams=None,
+                 team_budget=None, verbose=False):
+        env = lambda k, d, f=int: f(os.environ.get(k) or d)
+        self.n_agents = int(n_agents)
+        self.robot_radius = float(robot_radius)
+        self.ws_bounds = None if ws_bounds is None else np.asarray(ws_bounds, dtype=float)
+        self.obstacles = obstacles or []
+        self.d = env('GPA_DEG', degree or 5)
+        self.M = env('GPA_M', segments or 21)
+        self.De = env('GPA_ELEV', elev or 30)
+        self.ell_j, self.ell = int(ell_j), int(ell)
+        self.margin = env('GPA_MARGIN', margin or 0.005, float)
+        self.alg1_cap = env('GPA_CAP', alg1_cap or 3000)
+        self.alg2_blocks = env('GPA_BLOCKS', alg2_blocks if alg2_blocks is not None else 4)
+        self.alg2_steps = int(alg2_steps)
+        self.max_teams = env('GPA_MAXTEAMS', max_teams or 10 ** 6)
+        self.team_budget = env('GPA_TEAMBUDGET', team_budget or 20, float)
+        self.reach = env('GPA_REACH', 0.25, float)
+        self.verbose = verbose or bool(int(os.environ.get('GPA_VERBOSE', '0')))
+        self.smoothness_weight = 0.0                      # interface parity with the other projectors, unused
+        self.stats = dict(calls=0, teams_seen=0, teams_feasible=0, teams_solved=0,
+                          teams_ok=0, teams_failed=0, time=0.0)
+
+    def _problem(self, n, w0, w1, M):
+        return Problem(n, self.d, 4, self.ell, M, w0, w1, ell_j=self.ell_j, delta=self.margin, mode='orthant', De=self.De)
+
+    # ---------------- one team ----------------
+    def _keep_right(self, Wf):
+        """Steer the reference before a solve.  Keep-right (GPA_RIGHTHAND): wherever a robot has an ONCOMING
+        robot within GPA_RH_RANGE robot radii, shift its reference to its own right, weighted by proximity, so
+        head-on pairs pass side by side.  Circulation (GPA_CIRC): a robot whose heading would pass the
+        workspace centre against the team's net rotation is nudged to the side that carries its direction of
+        travel (GPA_CIRC_MODE=reflect: moved toward its mirror position instead).  GPA_SHIFT_SMOOTH spreads each
+        shift over that many waypoints; GPA_RH_WALLSAFE drops shifts that would approach a wall.  Endpoints stay
+        fixed; not applied at the final call."""
+        w = float(os.environ.get('GPA_RIGHTHAND', '0')); wc = float(os.environ.get('GPA_CIRC', '0'))
+        if (w <= 0 and wc <= 0) or getattr(self, 'final_call', False):
+            return Wf
+        L1, n = Wf.shape; na = n // 2
+        P = Wf.reshape(L1, na, 2)
+        V = np.gradient(P, axis=0)
+        sp = np.linalg.norm(V, axis=-1, keepdims=True)
+        T = V / np.maximum(sp, 1e-9)
+        Rt = np.stack([T[..., 1], -T[..., 0]], -1)                   # right-hand normal
+        R0 = float(os.environ.get('GPA_RH_RANGE', '4')) * self.robot_radius
+        D = np.linalg.norm(P[:, :, None] - P[:, None], axis=-1)      # (L1, na, na)
+        onc = (np.einsum('tai,tbi->tab', T, T) < -0.3) & (D < R0) & (D > 0)
+        wt = (np.clip(1 - D / R0, 0, 1) * onc).sum(-1)               # (L1, na)
+        wt *= (sp[..., 0] > 1e-6)
+        shift = w * np.tanh(wt)[..., None] * Rt
+        if wc > 0:
+            c = np.zeros(2) if self.ws_bounds is None else 0.5 * (np.asarray(self.ws_bounds[0]) + np.asarray(self.ws_bounds[1]))
+            # sign of the team's net rotation about the centre, from each robot's start->goal displacement
+            # (ties, e.g. antipodal swaps, default to counter-clockwise)
+            r0 = P[0] - c[None]; dv = P[-1] - P[0]
+            Lz = float((r0[:, 0] * dv[:, 1] - r0[:, 1] * dv[:, 0]).sum())
+            scale = float((np.linalg.norm(r0, axis=1) * np.linalg.norm(dv, axis=1)).sum()) + 1e-12
+            sgn = -1.0 if Lz < -0.05 * scale else 1.0
+            rel = c[None, None] - P
+            d = sgn * (T[..., 0] * rel[..., 1] - T[..., 1] * rel[..., 0])   # > 0: moving the team's way round
+            Rt = sgn * Rt
+            if os.environ.get('GPA_CIRC_MODE', 'push') == 'reflect':
+                # one-way lanes: a robot heading the wrong way for its side of the centre (d < 0) is moved toward
+                # its mirror position across the centre line (distance 2|d|, capped); right-way robots stay put
+                cap = float(os.environ.get('GPA_CIRC_CAP', '0.35'))
+                shift = shift + wc * np.minimum(2 * np.clip(-d, 0, None), cap)[..., None] * Rt * (sp > 1e-6)
+            else:
+                dmin = float(os.environ.get('GPA_CIRC_D', '4')) * self.robot_radius
+                shift = shift + wc * np.clip((dmin - d) / dmin, 0, 1)[..., None] * Rt * (sp > 1e-6)
+        sw = int(os.environ.get('GPA_SHIFT_SMOOTH', '0'))
+        if sw > 1:
+            # spread every lane change over about sw waypoints (moving average along time), so a shift never
+            # opens a sideways jump that breaks the speed bound
+            k = np.ones(sw) / sw; pad = sw // 2
+            Sp = np.pad(shift, ((pad, sw - 1 - pad), (0, 0), (0, 0)), mode='edge')
+            shift = np.stack([np.stack([np.convolve(Sp[:, a, c], k, mode='valid') for c in range(2)], -1)
+                              for a in range(shift.shape[1])], 1)
+        shift[0] = 0; shift[-1] = 0
+        if self.obstacles and os.environ.get('GPA_RH_WALLSAFE', '0') == '1':
+            # keep a shift only where the shifted point stays clear of every obstacle disc by r + margin
+            Cc = np.array([c for c, _ in self.obstacles], float); Rc = np.array([rr for _, rr in self.obstacles], float)
+            Q = P + shift
+            clr = (np.linalg.norm(Q[..., None, :] - Cc[None, None], axis=-1) - Rc[None, None]).min(-1)
+            need = self.robot_radius + max(self.margin, 0.0) + float(os.environ.get('GPA_RH_CLEAR', '0.02'))
+            shift = shift * (clr >= need)[..., None]
+        return (P + shift).reshape(L1, n)
+
+    def _certify_prioritized(self, Wf, dx_max):
+        """Final-call certification one robot at a time: each robot is restored alone against the walls, the
+        arena, the speed bound and every robot already certified (their spline curves as moving obstacles,
+        exact continuous-time separation 2r).  Order: fewest conflicts in the reference first; robots that do
+        not fit are retried once after the rest, then the ordering restarts with the unplaced robots first.
+        The team is certified iff every robot is.  Same blocks and tolerances as the joint solve."""
+        L = Wf.shape[0] - 1; n = Wf.shape[1]; na = n // 2; M = min(self.M, L); r2 = 2.0 * self.robot_radius
+        P = Wf.reshape(L + 1, na, 2)
+        conf = np.zeros(na)
+        for a in range(na):
+            for b in range(na):
+                if a != b: conf[a] += float(seg_min_dist(P[None, :, a], P[None, :, b])[0] < r2)
+        order = list(np.argsort(conf, kind='stable')); t_end = time.time() + self.team_budget
+        t_cap = time.time() + float(os.environ.get('GPA_PRIO_TEAM_MAX', '900'))   # safety cap per team
+        safety = float(os.environ.get('GPA_PRIO_SAFETY', '600'))
+        for rnd in range(int(os.environ.get('GPA_PRIO_RESTARTS', '3'))):
+            placed = {}                                              # a -> (segments (M,2,d+1), curve (L+1,2))
+            for attempt in range(2):
+                pending = [a for a in order if a not in placed]
+                if not pending: break
+                for a in pending:
+                    if time.time() > t_cap: return None
+                    w = P[:, a]; w2 = None
+
+                    def build(refs):
+                        # one robot: walls near ANY of the given reference paths, speed, arena (hard), and every
+                        # already certified robot as a moving obstacle (exact continuous-time 2r)
+                        keep = [(c, r) for c, r in self.obstacles
+                                if min(np.linalg.norm(q - np.asarray(c, float)[None], axis=1).min() for q in refs)
+                                <= float(r) + self.robot_radius + self.margin + self.reach]
+                        p = self._problem(2, w[0], w[-1], M)
+                        if dx_max and dx_max > 0: p.add_speed(float(dx_max) * L / M)
+                        for c, r in keep: p.add_sphere(np.asarray(c, float), float(r) + self.robot_radius)
+                        for _, (Sb, _c) in placed.items(): p.add_moving(Sb, r2)
+                        if self.ws_bounds is not None:
+                            p.add_box(self.ws_bounds[0] + self.robot_radius - self.margin + 1e-6, self.ws_bounds[1] - self.robot_radius + self.margin - 1e-6)
+                        p.set_tracking_fine(refs[-1])
+                        tv = p.tolvec()
+                        if self.ws_bounds is not None: tv[torch.tensor([t[0] == 'box' for t in p.types])] = 1e-12
+                        return p, tv
+
+                    def fits(p_, Z_):
+                        # the robots execute (and the benchmark scores) the straight lines between waypoints:
+                        # accept only if that motion clears walls, arena, speed and every placed robot exactly
+                        S = np.stack([v[1] for v in placed.values()] + [self._curve(p_, Z_, L, M).numpy()])
+                        return bool(self._pl_slack(S[:, None], dx_max)[0] >= 0)
+
+                    p, tv = build([w])
+                    Z = init_waypoints(p, w[np.round(np.linspace(0, L, M + 1)).astype(int)])
+                    Z, eta, ok = p.project(Z, tv, cap=self.alg1_cap, deadline=time.time() + safety); ok = ok and fits(p, Z)
+                    if not ok and os.environ.get('GPA_REPAIR', '1') == '1' and self.ws_bounds is not None:
+                        # wall repair: the reference crosses walls; re-route to the nearest wall-free route,
+                        # keep its timing, certify that instead
+                        w2 = _static_route(w, self.obstacles, self.robot_radius + self.margin + 0.002,
+                                           self.ws_bounds[0] + self.robot_radius, self.ws_bounds[1] - self.robot_radius)
+                        if w2 is not None:
+                            p, tv = build([w, w2])
+                            Z = init_waypoints(p, w2[np.round(np.linspace(0, L, M + 1)).astype(int)])
+                            Z, eta, ok = p.project(Z, tv, cap=self.alg1_cap, deadline=time.time() + safety); ok = ok and fits(p, Z)
+                    if not ok and os.environ.get('GPA_RETIME', '1') == '1':
+                        # timing repair: the remaining conflicts are with robots already placed; keep the route
+                        # and try re-timings of it (start later or arrive earlier) within the speed bound
+                        base = w2 if w2 is not None else w
+                        arc = np.r_[0, np.cumsum(np.linalg.norm(np.diff(base, axis=0), axis=1))]; tt = np.linspace(0, 1, L + 1)
+                        frac0 = np.interp(np.arange(L + 1), np.arange(L + 1), arc / max(arc[-1], 1e-12))
+                        for kind, k in [('delay', .1), ('early', .1), ('delay', .2), ('early', .2), ('delay', .3), ('delay', .4)]:
+                            if kind == 'delay': u_ = np.clip((tt - k) / (1 - k), 0, 1)
+                            else:               u_ = np.clip(tt / (1 - k), 0, 1)
+                            fr = np.interp(u_, tt, frac0) * arc[-1]
+                            wk = np.stack([np.interp(fr, arc, base[:, j]) for j in range(2)], 1)
+                            if np.linalg.norm(np.diff(wk, axis=0), axis=1).max() > float(dx_max) * 0.98: continue
+                            p, tv = build([w, wk])
+                            Z = init_waypoints(p, wk[np.round(np.linspace(0, L, M + 1)).astype(int)])
+                            Z, eta, ok = p.project(Z, tv, cap=self.alg1_cap, deadline=time.time() + safety); ok = ok and fits(p, Z)
+                            if ok: break
+                    if not ok and os.environ.get('GPA_DETOUR', '1') == '1' and self.ws_bounds is not None and placed:
+                        # detour repair: blocked by robots already placed (e.g. head-on in a one-lane passage);
+                        # re-route around the walls AND the whole paths of the placed robots it conflicts with
+                        wr = w2 if w2 is not None else w
+                        cl = [cb for _, (Sb, cb) in placed.items() if seg_min_dist(wr[None], cb[None])[0] < r2 + self.margin]
+                        ext = [(q, self.robot_radius) for cb in cl for q in cb
+                               if min(np.linalg.norm(q - w[0]), np.linalg.norm(q - w[-1])) > r2 + self.margin]
+                        w3 = None
+                        if ext:
+                            w3 = _static_route(w, list(self.obstacles) + ext, self.robot_radius + self.margin + 0.002,
+                                               self.ws_bounds[0] + self.robot_radius + self.margin, self.ws_bounds[1] - self.robot_radius - self.margin)
+                        if w3 is not None and np.linalg.norm(np.diff(w3, axis=0), axis=1).max() > float(dx_max) * 0.98:
+                            # longer route: traverse it at constant speed instead of the reference's profile
+                            a3 = np.r_[0, np.cumsum(np.linalg.norm(np.diff(w3, axis=0), axis=1))]
+                            w3 = np.stack([np.interp(np.linspace(0, a3[-1], L + 1), a3, w3[:, j]) for j in range(2)], 1)
+                            if a3[-1] / L > float(dx_max) * 0.9: w3 = None       # keep slack for the spline
+                        if w3 is not None:
+                            p, tv = build([w, w3])
+                            Z = init_waypoints(p, w3[np.round(np.linspace(0, L, M + 1)).astype(int)])
+                            Z, eta, ok = p.project(Z, tv, cap=self.alg1_cap, deadline=time.time() + safety); ok = ok and fits(p, Z)
+                    if ok: placed[a] = (p.segs(Z).detach().clone(), self._curve(p, Z, L, M).numpy())
+            if len(placed) == na: break
+            order = [a for a in order if a not in placed] + [a for a in order if a in placed]
+        if len(placed) < na: return None
+        return torch.as_tensor(np.concatenate([placed[a][1] for a in range(na)], axis=1))
+
+    def _pl_slack(self, Wt, dx_max):
+        """Exact slack of the piecewise-linear motion through the waypoints (what the robots execute and the
+        benchmark scores): robot spacing 2r, speed, walls and the arena (a segment stays inside the box iff
+        its endpoints do).  Wt: (na, G, L+1, 2) -> (G,); feasible iff >= 0."""
+        na, G = Wt.shape[0], Wt.shape[1]
+        need = 2.0 * self.robot_radius
+        slack = np.full(G, np.inf)
+        for a in range(na):
+            for b in range(a + 1, na):
+                slack = np.minimum(slack, seg_min_dist(Wt[a], Wt[b]) - need)
+        if dx_max:
+            step = np.linalg.norm(np.diff(Wt, axis=2), axis=-1).max(-1).max(0)   # (G,)
+            slack = np.minimum(slack, dx_max * (1 + 1e-9) - step)
+        if self.obstacles:
+            Cc = np.array([c for c, _ in self.obstacles], float); Rc = np.array([r for _, r in self.obstacles], float)
+            P0, D_ = Wt[..., :-1, :], np.diff(Wt, axis=2)                         # (na,G,L,2)
+            w = Cc[None, None, None] - P0[..., None, :]                             # (na,G,L,K,2)
+            dd = (D_ * D_).sum(-1)[..., None]
+            tt = np.clip((w * D_[..., None, :]).sum(-1) / np.maximum(dd, 1e-18), 0, 1)
+            dist = np.linalg.norm(w - tt[..., None] * D_[..., None, :], axis=-1)   # (na,G,L,K)
+            oslack = (dist - Rc - self.robot_radius).min(-1).min(-1).min(0)        # (G,)
+            slack = np.minimum(slack, oslack)
+        if self.ws_bounds is not None:
+            lo, hi = np.asarray(self.ws_bounds[0], float), np.asarray(self.ws_bounds[1], float)
+            e = np.minimum(Wt - lo, hi - Wt).min(-1).min(-1).min(0) - self.robot_radius
+            slack = np.minimum(slack, e)
+        return slack
+
+    def _pl_ok(self, Wf, dx_max):
+        """Wf: (L+1, n) team plan -> True iff its piecewise-linear motion is feasible (see _pl_slack)."""
+        Wf = np.asarray(Wf, float); na = Wf.shape[1] // 2
+        return bool(self._pl_slack(Wf.reshape(Wf.shape[0], na, 2).transpose(1, 0, 2)[:, None], dx_max)[0] >= 0)
+
+    def _solve_team(self, Wf, dx_max):
+        """Wf: (L+1, n) stacked waypoints of one team.  Returns (L+1, n) or None."""
+        if getattr(self, 'prio', False):
+            return self._certify_prioritized(Wf, dx_max)
+        Wf = self._keep_right(Wf)
+        L = Wf.shape[0] - 1
+        n = Wf.shape[1]
+        na = n // 2
+        M = min(self.M, L)
+        t0 = time.time()
+        # only the obstacle discs a robot could reach: the solve stays near the reference, and the returned
+        # plan is checked against the FULL geometry (_pl_ok, and the benchmark's own scorer)
+        obs_near = []
+        for a in range(na):
+            P_ = Wf[:, 2 * a:2 * a + 2]
+            keep = []
+            for c, r in self.obstacles:
+                d = np.linalg.norm(P_ - np.asarray(c, dtype=float)[None, :], axis=1).min()
+                if d <= float(r) + self.robot_radius + self.margin + self.reach:
+                    keep.append((c, r))
+            obs_near.append(keep)
+        p = self._problem(n, Wf[0], Wf[-1], M)
+        for a in range(na):
+            ra = [2 * a, 2 * a + 1]
+            for b in range(a + 1, na):
+                p.add_sep(ra, [2 * b, 2 * b + 1], 2.0 * self.robot_radius)
+            if dx_max and dx_max > 0:
+                p.add_speed(float(dx_max) * L / M, rows=ra)     # uniform speed bound == every chord <= dx_max
+            for c, r in obs_near[a]:
+                p.add_sphere(np.asarray(c, dtype=float), float(r) + self.robot_radius, rows=ra)
+        if self.ws_bounds is not None:
+            # the box is shrunk by delta (= margin) inside the problem: hand it bounds + r - delta so the target is
+            # the robot disc inside the arena, and make the box blocks hard
+            p.add_box(np.tile(self.ws_bounds[0] + self.robot_radius - self.margin + 1e-6, na),
+                      np.tile(self.ws_bounds[1] - self.robot_radius + self.margin - 1e-6, na))
+        p.set_tracking_fine(Wf)
+        tv = p.tolvec()
+        if self.ws_bounds is not None:
+            tv[torch.tensor([t[0] == 'box' for t in p.types])] = 1e-12
+        Z = init_waypoints(p, Wf[np.round(np.linspace(0, L, M + 1)).astype(int)])
+        dl = t0 + self.team_budget
+        Z, eta, ok = p.project(Z, tv, cap=self.alg1_cap, deadline=dl)
+        if not ok:
+            # steering call: hand back the partially restored iterate; the final call returns certified plans only
+            return self._curve(p, Z, L, M) if getattr(self, 'partial', False) else None
+        if self.alg2_blocks > 0:
+            Z, eta, _ = p.solve(Z, tv, blocks=self.alg2_blocks, steps=self.alg2_steps, eta=eta, deadline=dl)
+            if not bool((p.resid(p.pi_all(Z))[0] <= tv).all()):
+                Z, eta, ok = p.project(Z, tv, cap=self.alg1_cap, deadline=dl)
+                if not ok:
+                    return self._curve(p, Z, L, M) if getattr(self, 'partial', False) else None
+        return self._curve(p, Z, L, M)
+
+    @staticmethod
+    def _curve(p, Z, L, M):
+        """The spline at the L+1 waypoint times -> (L+1, n)."""
+        tt = torch.linspace(0, 1, L + 1, dtype=DT)
+        seg = torch.clamp((tt * M).floor().long(), 0, M - 1)
+        u = tt * M - seg
+        bu = bern.bernstein_basis_eval(p.d, u)                   # (d+1, L+1)
+        S = p.segs(Z)                                            # (M, n, d+1)
+        return torch.einsum('inp,pi->in', S[seg], bu)            # (L+1, n)
+
+    # ---------------- BaseProjector API ----------------
+    def project(self, x, params=None):
+        dev, dt_ = x.device, x.dtype
+        B, Hm1, _ = x.shape
+        na = self.n_agents
+        G = B // na
+        pos_init = params['pos_init'].detach().to('cpu', DT)
+        dx_max = float(params['dx_max']) if params.get('dx_max', None) is not None else None
+        xc = x.detach().to('cpu', DT)
+        out = xc.clone()
+        self.stats['calls'] += 1
+        t_call = time.time()
+
+        idx = np.arange(B).reshape(na, G)                        # idx[a, j] = a*G + j
+        W = torch.cat([pos_init.unsqueeze(1), xc], 1).numpy()    # (B, L+1, 2)
+        Wt = W[idx]                                              # (na, G, L+1, 2)
+        slack = self._pl_slack(Wt, dx_max)
+        bad = np.where(slack < 0)[0]
+        self.stats['teams_seen'] += G
+        self.stats['teams_feasible'] += int(G - len(bad))
+        order = bad[np.argsort(-slack[bad])]                     # least infeasible first
+        order = order[:self.max_teams]
+        # the projector object lives across instances (n-inits > 1): every GPA_NCALLS-th call is the final one
+        ncalls = int(os.environ.get('GPA_NCALLS', '13'))
+        final = self.stats['calls'] % ncalls == 0
+        saved = (self.team_budget, self.alg1_cap, self.alg2_blocks)
+        self.partial = (not final) and os.environ.get('GPA_PARTIAL', '1') == '1'
+        self.final_call = final
+        if not final:
+            self.team_budget = float(os.environ.get('GPA_INTER_BUDGET', '4'))
+            self.alg1_cap = int(os.environ.get('GPA_INTER_CAP', '400'))
+            self.alg2_blocks = 0
+        jobs = [(int(j), np.ascontiguousarray(Wt[:, j].transpose(1, 0, 2).reshape(Hm1 + 1, 2 * na)))
+                for j in order]
+        self.stats['teams_solved'] += len(jobs)
+        if os.environ.get('GPA_BATCH', '0') == '1' and jobs:
+            results = _run_teams_batched(self, jobs, dx_max)
+        else:
+            results = _run_teams(self, jobs, dx_max)
+        self.team_budget, self.alg1_cap, self.alg2_blocks = saved
+        for j, Wp in results:
+            if Wp is None:
+                self.stats['teams_failed'] += 1
+                continue
+            self.stats['teams_ok'] += 1
+            Wp = torch.as_tensor(Wp).reshape(Hm1 + 1, na, 2).permute(1, 0, 2)     # (na, L+1, 2)
+            out[idx[:, j]] = Wp[:, 1:, :].to(DT)
+        self.stats['time'] += time.time() - t_call
+        if final and os.environ.get('GPA_CERT_LOG'):
+            # GlobMP's own certificate for this instance: teams the exact screen already passed plus teams the
+            # final solve certified.  The returned plan is the first certified candidate.
+            cert = sorted(set(int(j) for j in range(G) if j not in set(bad.tolist())) | {int(j) for j, Wp in results if Wp is not None})
+            import json as _json
+            with open(os.environ['GPA_CERT_LOG'], 'a') as fh:
+                fh.write(_json.dumps(dict(call=int(self.stats['calls']), certified=cert)) + '\n')
+        if self.verbose:
+            print(f"  [gpa] call {self.stats['calls']}: {G} teams, {G - len(bad)} already feasible, "
+                  f"{len(order)} solved, {self.stats['teams_failed']} failed, "
+                  f"{time.time() - t_call:.1f}s", flush=True)
+        return out.to(dev, dt_)
+
+
+_POOL_SELF = None
+_POOL = None
+
+
+def _pool_init(proj):
+    global _POOL_SELF
+    torch.set_num_threads(1)
+    _POOL_SELF = proj
+
+
+def _team_worker(arg):
+    j, Wf, dx_max, budget = arg
+    if budget is not None:
+        _POOL_SELF.team_budget, _POOL_SELF.alg1_cap, _POOL_SELF.alg2_blocks, _POOL_SELF.partial, _POOL_SELF.final_call, _POOL_SELF.prio = budget
+    try:
+        Wp = _POOL_SELF._solve_team(Wf, dx_max)
+        if Wp is not None and getattr(_POOL_SELF, 'final_call', False) and not _POOL_SELF._pl_ok(Wp.numpy(), dx_max):
+            Wp = None                                    # the returned plan must also be exact on its waypoint motion
+        return j, (None if Wp is None else Wp.numpy())
+    except Exception:
+        if os.environ.get('GPA_VERBOSE', '0') == '1':
+            import traceback; traceback.print_exc()
+        return j, None
+
+
+def _run_teams(proj, jobs, dx_max):
+    """Solve one call's teams on a persistent pool of SPAWNED workers (GPA_WORKERS; CPU only).  fork is
+    unusable here (the sampler has already run autograd, which torch refuses in a forked child); a spawned
+    child would re-run the driver script as __mp_main__, so the pool is created with __main__.__file__
+    hidden and the workers only import this module."""
+    global _POOL, _POOL_SELF
+    nw = int(os.environ.get('GPA_WORKERS', '16'))
+    budget = (proj.team_budget, proj.alg1_cap, proj.alg2_blocks, getattr(proj, 'partial', False), getattr(proj, 'final_call', False), getattr(proj, 'prio', False))
+    if nw <= 1 or len(jobs) <= 1:
+        _POOL_SELF = proj
+        return [_team_worker((j, W, dx_max, budget)) for j, W in jobs]
+    if _POOL is None:
+        import multiprocessing as mp, sys as _s
+        main = _s.modules.get('__main__'); saved = getattr(main, '__file__', None)
+        if saved is not None:
+            del main.__file__
+        try:
+            _POOL = mp.get_context('spawn').Pool(nw, initializer=_pool_init, initargs=(proj,))
+        finally:
+            if saved is not None:
+                main.__file__ = saved
+    args = [(j, W, dx_max, budget) for j, W in jobs]
+    if getattr(proj, 'prio', False) and os.environ.get('GPA_FIRST_CERT', '1') == '1':
+        # the method returns one plan: take the first team that certifies and drop the rest
+        out = []
+        for j, Wp in _POOL.imap_unordered(_team_worker, args, chunksize=1):
+            out.append((j, Wp))
+            if Wp is not None:
+                _POOL.terminate(); _POOL = None
+                break
+        return out
+    return _POOL.map(_team_worker, args, chunksize=1)
+
+
+def _run_teams_batched(proj, jobs, dx_max):
+    """All teams of one call as ONE batched problem (batch.py); same constraints, steering and
+    certified-only final output as the per-team path."""
+    from batch import solve_teams
+    W = np.stack([proj._keep_right(Wf) for _, Wf in jobs])                   # (G, L+1, n)
+    dev = os.environ.get('GPA_DEVICE', 'cuda')
+    if dev == 'cpu':
+        torch.set_num_threads(int(os.environ.get('GPA_BATCH_THREADS', '4')))
+    out, ok = solve_teams(W, proj.n_agents, proj.robot_radius, proj.obstacles, proj.ws_bounds, dx_max,
+                          d=proj.d, M=proj.M, delta=proj.margin, budget=proj.team_budget, cap=proj.alg1_cap,
+                          alg2_blocks=proj.alg2_blocks, partial=getattr(proj, 'partial', False), device=dev)
+    part = getattr(proj, 'partial', False)
+    ok = np.asarray(ok, bool).copy()
+    if getattr(proj, 'final_call', False):
+        for g in np.where(ok)[0]:
+            ok[g] = proj._pl_ok(np.asarray(out[g]), dx_max)          # exact on the executed waypoint motion too
+    res = [(j, (out[g] if (ok[g] or part) else None)) for g, (j, _) in enumerate(jobs)]
+    if getattr(proj, 'final_call', False) and os.environ.get('GPA_PRIORITIZED', '0') == '1':
+        failed = [(j, Wf) for g, (j, Wf) in enumerate(jobs) if not ok[g]]
+        first_only = os.environ.get('GPA_FIRST_CERT', '1') == '1'
+        if failed and not (first_only and bool(np.any(ok))):
+            # prioritized certification on the CPU pool, fewest-conflict candidates first; with GPA_FIRST_CERT
+            # it stops at the first certified team (the returned plan is the first certified candidate)
+            na_ = proj.n_agents
+            def conflicts(Wf):
+                Pw = Wf.reshape(Wf.shape[0], na_, 2); r2 = 2 * proj.robot_radius
+                return sum(float(seg_min_dist(Pw[None, :, a], Pw[None, :, b])[0] < r2) for a in range(na_) for b in range(a + 1, na_))
+            failed.sort(key=lambda jw: conflicts(jw[1]))
+            nw = max(1, int(os.environ.get('GPA_WORKERS', '16')))
+            proj.prio = True; pr = {}
+            try:
+                if first_only:
+                    pr.update(dict(_run_teams(proj, failed, dx_max)))
+                else:
+                    for k in range(0, len(failed), nw):
+                        pr.update(dict(_run_teams(proj, failed[k:k + nw], dx_max)))
+            finally:
+                proj.prio = False
+            res = [(j, (pr[j] if (Wp is None and pr.get(j) is not None) else Wp)) for j, Wp in res]
+    return res

applications/multi_robot/robotarium/robotarium_baselines/.gitignore

@@ -0,0 +1,6 @@
+config/paths.env
+runs/
+**/data_trajectories/
+**/data_trained_models/
+**/results/
+**/__pycache__/

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/.gitignore

@@ -0,0 +1,43 @@
+# Caches
+__pycache__/
+*.pyc
+*.pyo
+
+# Python packaging
+*.egg-info/
+build/
+dist/
+
+# Experiments, logs, and run outputs
+results/
+logs/
+scripts/inference/results/
+scripts/inference/data_trained_models/
+
+# Downloaded / generated data
+data_trajectories/
+data_trained_models/
+*.tar.xz
+
+# Mujoco
+MUJOCO_LOG.TXT
+
+# Notebooks
+*.ipynb
+
+# Claude
+.claude/
+
+# Slurm
+*.out
+*.err
+
+# OS
+.DS_Store
+
+# IDE
+.vscode/
+.idea/
+
+# Wandb
+wandb/

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/LICENSE

@@ -0,0 +1,28 @@
+MIT License
+
+Copyright (c) 2022 Julen Urain
+Copyright (c) 2023 Columbia Artificial Intelligence and Robotics Lab
+Copyright (c) 2024 Zeyu Feng and Hao Luan
+Copyright (c) 2024 Yorai Shaoul, Itamar Mishani, and Shivam Vats 
+Copyright (c) 2025 Hao Luan and Yi Xian Goh
+Copyright (c) 2025 Wei Xiao
+Copyright (c) 2026 Laboratory for Information and Decision Systems, Azizan Group, MIT
+Copyright (c) 2026 Paolo Giaretta, Zeyang Li, Navid Azizan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/README.md

@@ -0,0 +1,209 @@
+# Safe multi-robot motion planning in MMD
+
+This folder presents the benchmarks for _DiRecT_, _PCD_, and relevant baselines on the multi-robot motion planning task. The codebase is forked from the [**Projected Coupled Diffusion (PCD)**](https://github.com/EdmundLuan/pcd) [1], which itself extends [MMD](https://github.com/yoraish/mmd) [2]: the planner stack, projection operators, coupling costs, and most of the inference / analysis scripts already exist in the PCD parent. What is contributed here is the **DiRecT** planner on top.
+
+![Multi-robot grid visualization](assets/mrmp-grid.gif)
+
+**Table of Contents**
+- [Installation](#installation)
+- [Get Datasets](#get-datasets)
+- [Generating missing Drop Region dataset](#generating-missing-drop-region-dataset)
+- [Planning Inference](#planning-inference)
+- [Analyzing Results](#analyzing-results)
+- [What we modified/added](#what-we-modifiedadded)
+- [Optimizers](#optimizers)
+- [Code structure](#code-structure)
+- [References](#references)
+
+---
+
+## Installation
+
+Follow the setup steps from the PCD parent ([EdmundLuan/pcd](https://github.com/EdmundLuan/pcd)) as below. In summary:
+
+### Requirements
+- [miniconda](https://docs.conda.io/projects/miniconda/en/latest/index.html)
+
+### Steps
+
+1. Create a conda environment and activate it (the env is named `direct-mmd` in `environment.yml`; this installs Python, PyTorch, JAX, and all pip dependencies):
+    ```sh
+    conda env create -f environment.yml
+    conda activate direct-mmd
+    ```
+2. Install the local (patched) versions of [2, 3]:
+    ```sh
+    cd deps/torch_robotics      && pip install -e . && cd ../..
+    cd deps/experiment_launcher && pip install -e . && cd ../..
+    cd deps/motion_planning_baselines && pip install -e . && cd ../..
+    ```
+3. Install the `mmd` package:
+    ```sh
+    pip install -e .
+    ```
+
+## Get Datasets
+
+Download the sample datasets and trained models provided by the PCD parent:
+
+```sh
+conda activate direct-mmd
+
+gdown --id 1Onw0s1pDsMLDfJVOAqmNme4eVVoAkkjz
+tar -xJvf data_trajectories.tar.xz
+
+gdown --id 1WO3tpvg-HU0m9RyDvGyfDamo7roBYMud
+tar -xJvf data_trained_models.tar.xz
+```
+
+> If the gdown IDs above have expired, the current dataset locations are tracked in the PCD parent repo: https://github.com/EdmundLuan/pcd.
+
+## Generating missing Drop Region dataset
+
+The gdown bundle ships the **DropRegion trained model** but not its **trajectory dataset**. The dataset is still required at inference time — only for the normalizer stats (min/max for `[-1, 1]` scaling), not for retraining.
+
+> [!WARNING]
+> To run any DropRegion experiment you must first generate this dataset, otherwise inference will fail to load the normalizer. Before running, edit the SLURM partition, GPU resources, and other system-specific settings in the script to match your cluster. Generation may take a while depending on your hardware.
+
+```sh
+conda activate direct-mmd
+python scripts/generate_data/launch_generate_trajectories.py
+```
+
+The launcher has no CLI; edit `env_id` at the top of the file to switch environments (default `'EnvDropRegion2D'`). Output goes to `logs/<timestamp>/...` — move or symlink it into `data_trajectories/EnvDropRegion2D-RobotPlanarDisk/` to match the layout of the other environments.
+
+---
+
+## Planning Inference
+
+All single-agent methods (PCD, DiRecT, baselines) go through `scripts/inference/launch_pcdiff.py`; MMD-CBS uses its own `launch_mmd.py`. Pick a config from `configs/` (`pcdiff_*.yaml` for PCD / baselines, `direct_*.yaml` for DiRecT) and one of the six commands below.
+
+Before the first run, create the results directory:
+```sh
+mkdir -p results/000event_logs
+```
+
+```sh
+conda activate direct-mmd
+
+# Unconstrained — vanilla diffusion, no projection
+python scripts/inference/launch_pcdiff.py --config pcdiff_empty.yaml -g vanilla -p none -m unconstrained
+
+# PCD-SHD — coupling guidance + ADMM velocity projection, squared-hinge cost
+python scripts/inference/launch_pcdiff.py --config pcdiff_empty.yaml -g coupling -p max_vel_admm -c hinge_sqr_l2 -r -m pcd_shd
+
+# PCD-LB — coupling guidance + ADMM velocity projection, log-barrier cost
+python scripts/inference/launch_pcdiff.py --config pcdiff_empty.yaml -g coupling -p max_vel_admm -c sum_log_l2 -r -m pcd_lb
+
+# Final projection — unguided diffusion, project only at the last sampling step
+python scripts/inference/launch_pcdiff.py --config direct_empty.yaml -g final -p scp_admm -m final_projection
+
+# DiRecT — unguided diffusion + projection every step (use -p scp_admm or -p ipopt)
+python scripts/inference/launch_pcdiff.py --config direct_empty.yaml -g base -p scp_admm -m direct
+
+# MMD-CBS — multi-agent baseline (no CLI; edit MultiAgentPlanningExperimentConfig in the file)
+python scripts/inference/launch_mmd.py
+```
+
+**Optional CLI overrides** (apply to any method above):
+- `-n <int>` — number of initial states (default 100), e.g. `-n 2` for a quick test.
+- `-v <float>` — velocity cap for projection. Repeatable for sweeps, e.g. `-v 0.675 -v 0.703`.
+- `-a <int>` — number of agents (default 4).
+
+Swap `pcdiff_empty.yaml` / `direct_empty.yaml` for the `highway`, `conveyor`, or `dropregion` variants to reproduce the other environments. To launch the full sweep across all methods × environments × agent counts in one go, use `python scripts/inference/run_benchmark.py --help`, or one of the shell wrappers in `scripts/inference/` (`run_full_sweep.sh`, `run_single_method.sh`, `run_all_methods.sh`, `sbatch_benchmark.sh`).
+
+Results land under `results/<EnvRobot>Random/...` (PCD / DiRecT under `agent___N/method___.../cost___...`, MMD-CBS under `<timestamp>/num_agents___N/planner___CBS/...`); event logs are written to `results/000event_logs/`.
+
+---
+
+## Analyzing Results
+
+PCD/DiRecT and MMD-CBS produce different output layouts, so they have separate analyzers:
+
+```sh
+# PCD / DiRecT results (any run produced by launch_pcdiff.py)
+python scripts/analyze/analyzing_pcd.py --config analysis_config.yaml
+
+# MMD-CBS results (auto-discover + post-process + analyze)
+python scripts/analyze/run_all_mmd_postprocessing.py
+```
+
+Point the analyzer at the runs you want by editing `log_paths` in `configs/analysis_config.yaml` (PCD/DiRecT) or `configs/analyze_mmd.yaml` (MMD-CBS).
+
+---
+
+## What we modified/added
+
+**Inherited from PCD (and from MMD before it):**
+- `mmd/models/projection/` — ADMM-based state-space projection operators.
+- `mmd/coupling_costs/` — pair-wise inter-robot cost functions.
+- `mmd/models/diffusion_models/sample_functions.py::projected_coupling_ddpm_sample_fn` — projection-aware DDPM sampling hook.
+- `mmd/models/diffusion_models/guides.py` — `GuideManager` framework for gradient-based coupling guidance.
+- `mmd/planners/single_agent/` — `MPD`, `MPDEnsemble`, `PCD` planners (upstream MMD + PCD).
+- `mmd/planners/multi_agent/` — `CBS`, `PrioritizedPlanning` multi-agent planners (upstream MMD).
+- `mmd/common/` — conflict bookkeeping (`conflict_conversion.py`, `multi_agent_utils.py`, `constraints.py`, `experiences.py`).
+- `configs/pcdiff_*.yaml`, `launch_pcdiff.py` / `launch_mmd.py`, analysis scripts under `scripts/analyze/`, and data-generation scripts under `scripts/generate_data/`.
+
+**Added by this fork (DiRecT):**
+- `mmd/planners/single_agent/direct.py` — new `DiRecT` (unguided diffusion + projection) planner class.
+- `mmd/planners/single_agent/` — new `IpoptProjector` and `JaxScpAdmmProjector` projector classes.
+- `mmd/models/diffusion_models/` — additional sampling-loop wiring to call the new projectors from the `-g base/final` branches.
+- `mmd/config/mmd_params.py` — new entries for the DiRecT hyperparameters.
+- New DiRecT branches in `scripts/inference/launch_pcdiff.py` — the `-g {base, final}` guiders and the `-p {ipopt, scp_admm}` projectors.
+- New configs `configs/direct_{empty,highway,conveyor,dropregion}.yaml`.
+- New benchmark methods `direct` and `final_projection` in `scripts/inference/run_benchmark.py`, plus shell helpers (`run_all_methods.sh`, `run_full_sweep.sh`, `run_single_method.sh`, and SLURM wrappers).
+- Local patches to `deps/torch_robotics/` and other vendored deps for JAX SCP-ADMM and IPOPT projector compatibility.
+
+## Optimizers
+
+The planning stack exposes three families of hard-projection optimizers, picked at the command line via `-p, --projector`:
+
+- **ADMM velocity projection** *(from the PCD parent)* — implemented in `mmd/models/projection/projectors.py` as `ADMMProjectionOperator` and the GPU-accelerated `ADMMProjectionOperatorCUDAGraph`. These solve the convex problem of projecting a candidate trajectory onto the bounded-velocity set (`‖x_{i+1} − x_i‖₂ ≤ v_max·Δt`) by alternating an x-update (linear solve with cached LU / Cholesky factorization), a z-update (Euclidean L2-ball projection), and a u-update (dual-variable ascent), with primal/dual residual convergence checks. Selected via `-p max_vel_admm`.
+- **JAX Sequential Convex Programming + ADMM** *(DiRecT)* — implemented in `JaxScpAdmmProjector` (`mmd/planners/single_agent/`). It linearizes harder constraints (obstacles / dynamics) via SCP and runs ADMM on each linearized subproblem in JAX, JIT-compiled for batched trajectories. Selected via `-p scp_admm`; this is the default projector for DiRecT and for the `final_projection` method.
+- **IPOPT interior-point** *(DiRecT)* — implemented in `IpoptProjector` as a general-purpose alternative when ADMM-based projectors are not flexible enough for a given constraint set. Selected via `-p ipopt`.
+
+---
+
+## Code structure
+
+Only DiRecT-added and implementation-critical files are listed; everything else inherited from PCD/MMD is omitted for brevity.
+
+```
+multi_robot/
+├── configs/
+│   ├── direct_{empty,highway,conveyor,dropregion}.yaml      # NEW: DiRecT configs
+│   └── pcdiff_{empty,highway,conveyor,dropregion}.yaml      # PCD / CD configs (inherited)
+├── mmd/
+│   ├── planners/single_agent/
+│   │   └── direct.py                                        # NEW: DiRecT, IpoptProjector, JaxScpAdmmProjector
+│   ├── models/projection/projectors.py                      # ADMM projection operators (inherited)
+│   └── coupling_costs/                                      # inter-robot coupling costs (inherited)
+├── scripts/
+│   ├── inference/
+│   │   ├── launch_pcdiff.py                                 # main entrypoint (PCD / DiRecT / baselines)
+│   │   ├── launch_mmd.py                                    # MMD-CBS entrypoint (edit-the-file)
+│   │   └── run_benchmark.py                                 # full sweep across methods × envs × #agents
+│   ├── generate_data/launch_generate_trajectories.py        # generate DropRegion (or other) datasets
+│   └── analyze/{analyzing_pcd.py, run_all_mmd_postprocessing.py}
+├── deps/torch_robotics/                                     # local fork (patched for SCP-ADMM / IPOPT)
+├── environment.yml                                          # conda env: direct-mmd
+└── setup.py                                                 # `pip install -e .`  →  package `mmd`
+```
+
+## ✏️ Citation
+If you find our code or paper useful for your research, please consider citing our work:
+```
+@misc{giaretta2026directsafediffusionbasedplanning,
+      title={DiRecT: Safe Diffusion-Based Planning via Receding-Horizon Denoising}, 
+      author={Paolo Giaretta and Zeyang Li and Navid Azizan},
+      year={2026},
+      eprint={2606.15359},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2606.15359}, 
+}
+```
+
+## References
+- [1] Hao Luan. Projected Coupled Diffusion (PCD) for multi-robot motion planning. Implementation extending MMD. https://github.com/EdmundLuan/pcd
+- [2] Y. Shaoul, I. Mishani, S. Vats, J. Li, M. Likhachev. Multi-Robot Motion Planning with Diffusion Models. ICLR 2025. https://github.com/yoraish/mmd

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/configs/analysis_config.yaml

@@ -0,0 +1,63 @@
+# -------------------
+# ANALYSIS CONFIG
+# -------------------
+
+# List of data directories to process. The script will loop through each directory.
+# Add the paths to your own benchmark logs here, e.g.:
+#   - "results/000event_logs/benchmark-YYYYMMDD_HHMMSS-my_tag.log"
+log_paths: # []
+  # Examples (uncomment / adapt to your runs):
+  - "results/000event_logs/eval_runs-260519_152524-unconstrained.log"
+  # - "results/000event_logs/benchmark-260423_232958-bench_j1of6_conveyor_v1.21.log"
+
+
+# Switches for plotting and saving figures/data.
+switches:
+  # plot_trajectories: true
+  plot_trajectories: false
+  plot_velocity_stats: true
+  save_plots: true
+  save_stats_json: true
+
+# General constants and parameters for the analysis.
+constants:
+  EPS: 1.0e-6
+  VEL_TOL: 5.0e-5  # Tolerance for velocity checks
+  ## --------------- Highway -------------------
+  # DEFAULT_MAX_VEL: 0.647  # Default maximum velocity threshold (80% of max action)
+  # DEFAULT_MAX_VEL: 0.781  # (90% of max action)
+  # DEFAULT_MAX_VEL: 0.878 # (95% of max action)
+  ## --------------- Empty -------------------
+  ##[0.675, 0.692, 0.703]
+  # DEFAULT_MAX_VEL: 0.675  # (75% of max action)
+  # DEFAULT_MAX_VEL: 0.692  # (90% of max action)
+  # DEFAULT_MAX_VEL: 0.703  # (95% of max action)
+  ## --------------- Conveyor -------------------
+  ## [1.08, 1.46, 1.76]
+  # DEFAULT_MAX_VEL: 1.21  # (85% of max action)
+  # DEFAULT_MAX_VEL: 1.46  # (90% of max action)
+  # DEFAULT_MAX_VEL: 1.76  # (95% of max action)
+  ## --------------- DropRegion -------------------
+  ## [0.928, 1.13, 1.34]
+  DEFAULT_MAX_VEL: 0.928  # (80% of max action)
+  # DEFAULT_MAX_VEL: 1.13  # (90% of max action)
+  # DEFAULT_MAX_VEL: 1.34  # (95% of max action)
+
+
+  VEL_HIST_QUANTILES: [0.75, 0.8, 0.85, 0.9, 0.95, 0.99]
+
+# Plotting-specific parameters.
+plotting:
+  # Number of initial states figures to plot
+  num_inits_plot: 10
+  # Number of trajectory tuples to plot
+  num_trj_tups_plot: 64
+  # Trajectory plotting settings
+  traj_plot_xlim: [-1, 1]
+  traj_plot_ylim: [-1, 1]
+  traj_plot_h_end: 64
+  # Velocity histogram settings
+  vel_hist_bins: 100
+
+data:
+  default_trained_models_dir: "default_trained_models_dir"

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/configs/analyze_mmd.yaml

@@ -0,0 +1,59 @@
+# -------------------
+# ANALYSIS CONFIG
+# -------------------
+
+# List of data directories to process. The script will loop through each directory.
+log_paths:
+  ##---------------------- Highway ----------------------
+  # - "results/EnvHighways2DRobotPlanarDiskRandom/dummy_logs.log"
+  ### -------------- Empty -------------------
+  # - "results/EnvEmpty2DRobotPlanarDiskRandom/dummy_logs-04142026_202235.log"
+  ### ----------------- Conveyor -------------------
+  # - "results/EnvConveyor2DRobotPlanarDiskRandom/dummy_logs-04142026_204123.log"
+
+switches:
+  plot_trajectories: false
+  plot_velocity_stats: false
+  save_plots: false
+  save_stats_json: true
+
+constants:
+  EPS: 1.0e-6
+  VEL_TOL: 5.0e-5  # Tolerance for velocity checks
+  ## --------------- Highway -------------------
+  # DEFAULT_MAX_VEL: 0.647  # Default maximum velocity threshold (80% of max action)
+  # DEFAULT_MAX_VEL: 0.781  # (90% of max action)
+  # DEFAULT_MAX_VEL: 0.878 # (95% of max action)
+  ## --------------- Empty -------------------
+  ##[0.675, 0.692, 0.703]
+  # DEFAULT_MAX_VEL: 0.675  # (75% of max action)
+  # DEFAULT_MAX_VEL: 0.692  # (90% of max action)
+  # DEFAULT_MAX_VEL: 0.703  # (95% of max action)
+  ## --------------- Conveyor -------------------
+  ## [1.21, 1.46, 1.76]
+  # DEFAULT_MAX_VEL: 1.21  # (85% of max action)
+  # DEFAULT_MAX_VEL: 1.46  # (90% of max action)
+  DEFAULT_MAX_VEL: 1.76  # (95% of max action)
+  ## --------------- DropRegion -------------------
+  ## [0.928, 1.13, 1.34]
+  # DEFAULT_MAX_VEL: 0.928  # (85% of max action)
+  # DEFAULT_MAX_VEL: 1.13  # (90% of max action)
+  # DEFAULT_MAX_VEL: 1.34  # (95% of max action)
+
+  VEL_HIST_QUANTILES: [0.75, 0.8, 0.85, 0.9, 0.95, 0.99]
+
+# Plotting-specific parameters.
+plotting:
+  # Number of initial states figures to plot
+  num_inits_plot: 10
+  # Number of trajectory tuples to plot
+  num_trj_tups_plot: 64
+  # Trajectory plotting settings
+  traj_plot_xlim: [-1, 1]
+  traj_plot_ylim: [-1, 1]
+  traj_plot_h_end: 64
+  # Velocity histogram settings
+  vel_hist_bins: 100
+
+data:
+  default_trained_models_dir: "default_trained_models_dir"

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/configs/direct_conveyor.yaml

@@ -0,0 +1,79 @@
+experiment:
+  name: "single_tile_conveyor"
+  num_agents: 2
+  n_inits: 3
+  runtime_limit: 60  # seconds
+  seed: 18
+  time_str: "auto"
+  render_animation: true
+
+system:
+  device: "cuda:0"
+  base_dir: "auto"
+  results_dir_name: "logs"
+  trained_models_dir_name: "data_trained_models"
+  log_pth: null
+
+robot:
+  planar_disk_radius: 0.05
+
+environment:
+  type: "single_tile"
+  global_model_ids: [['EnvConveyor2D-RobotPlanarDisk']]
+  agent_skeleton_l: "auto"
+  start_goal_config:
+    method: "random_in_env"
+    env_class_for_random: "EnvConveyor2D"
+    margin: 0.2
+    obstacle_margin: 0.2
+
+planner:
+  multi_agent:
+    class: "DummyPlanning"
+    stagger_start_time_dt: 0
+    is_xcbs: false
+    is_ecbs: false
+    conflict_type_to_constraint_types:
+      PointConflict: [ "MultiPointConstraint" ]
+
+  single_agent:
+    class: "DiRecT"
+    planner_alg: "base"
+    horizon: 64
+    trajectory_duration: 5.0
+    n_samples: 128
+    debug: true
+    timeit: false
+    use_guide_on_extra_objects_only: false
+    n_local_inference_noising_steps: 3
+    n_local_inference_denoising_steps: 3
+    start_guide_steps_fraction: 1.0
+    n_guide_steps: 1
+    n_diffusion_steps_without_noise: 0
+    cost_func_key: "sum_log_l2"
+    cost_func_param:
+      margin_coeff: 1.9
+      radius_coeff: 6.0
+    weight_grad_cost_collision: 0.05
+    weight_grad_cost_smoothness: 0.0
+    weight_grad_cost_constraints: 0.2
+    weight_grad_cost_soft_constraints: 0.0
+    factor_num_interpolated_points_for_collision: 0
+    projector_type: "JaxScpADMM"
+    start_projection_guidance: 0.5 # guidance start
+    low_level_choose_path_from_batch_strategy: "least_collisions"
+
+  projection:
+    projector:
+      _target_: "ADMMProjectionOperatorCUDAGraph"
+      rho: 6.0
+      max_iter: 500
+      tol: 1.0e-5
+      decomp: "lu"
+      convergence_check: true
+      check_period: 200
+      verbose: false
+
+    project_params:
+      vel_max: 0.8
+      smoothness_weight: 0.0

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/configs/direct_dropregion.yaml

@@ -0,0 +1,79 @@
+experiment:
+  name: "single_tile_dropregion"
+  num_agents: 2
+  n_inits: 3
+  runtime_limit: 60  # seconds
+  seed: 18
+  time_str: "auto"
+  render_animation: true
+
+system:
+  device: "cuda:0"
+  base_dir: "auto"
+  results_dir_name: "logs"
+  trained_models_dir_name: "data_trained_models"
+  log_pth: null
+
+robot:
+  planar_disk_radius: 0.05
+
+environment:
+  type: "single_tile"
+  global_model_ids: [['EnvDropRegion2D-RobotPlanarDisk']]
+  agent_skeleton_l: "auto"
+  start_goal_config:
+    method: "random_in_env"
+    env_class_for_random: "EnvDropRegion2D"
+    margin: 0.2
+    obstacle_margin: 0.2
+
+planner:
+  multi_agent:
+    class: "DummyPlanning"
+    stagger_start_time_dt: 0
+    is_xcbs: false
+    is_ecbs: false
+    conflict_type_to_constraint_types:
+      PointConflict: [ "MultiPointConstraint" ]
+
+  single_agent:
+    class: "DiRecT"
+    planner_alg: "base"
+    horizon: 64
+    trajectory_duration: 5.0
+    n_samples: 128
+    debug: true
+    timeit: false
+    use_guide_on_extra_objects_only: false
+    n_local_inference_noising_steps: 3
+    n_local_inference_denoising_steps: 3
+    start_guide_steps_fraction: 1.0
+    n_guide_steps: 1
+    n_diffusion_steps_without_noise: 0
+    cost_func_key: "sum_log_l2"
+    cost_func_param:
+      margin_coeff: 1.9
+      radius_coeff: 6.0
+    weight_grad_cost_collision: 0.05
+    weight_grad_cost_smoothness: 0.0
+    weight_grad_cost_constraints: 0.2
+    weight_grad_cost_soft_constraints: 0.0
+    factor_num_interpolated_points_for_collision: 0
+    projector_type: "JaxScpADMM"
+    start_projection_guidance: 0.5
+    low_level_choose_path_from_batch_strategy: "least_collisions"
+
+  projection:
+    projector:
+      _target_: "ADMMProjectionOperatorCUDAGraph"
+      rho: 6.0
+      max_iter: 500
+      tol: 1.0e-5
+      decomp: "lu"
+      convergence_check: true
+      check_period: 200
+      verbose: false
+
+    project_params:
+      vel_max: 0.8
+      smoothness_weight: 0.0

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/configs/direct_empty.yaml

@@ -0,0 +1,81 @@
+
+experiment:
+  name: "single_tile_empty"
+  num_agents: 8
+  n_inits: 10
+  runtime_limit: 60  # seconds
+  seed: 18
+  time_str: "auto"  # auto timestamp at runtime
+  render_animation: true
+
+system:
+  device: "cuda:0"
+  base_dir: "auto"  # auto resolve base project directory
+  results_dir_name: "logs"
+  trained_models_dir_name: "data_trained_models"
+  log_pth: null  # use default logger path behavior
+
+robot:
+  planar_disk_radius: 0.05
+
+environment:
+  type: "single_tile"
+  global_model_ids: [['EnvEmpty2D-RobotPlanarDisk']]
+  agent_skeleton_l: "auto"  # auto-generate skeleton from environment
+  start_goal_config:
+    method: "random_in_env"
+    env_class_for_random: "EnvEmpty2D"
+    margin: 0.2
+    obstacle_margin: 0.11
+
+planner:
+  multi_agent:
+    class: "DummyPlanning"
+    stagger_start_time_dt: 0
+    is_xcbs: false
+    is_ecbs: false
+    conflict_type_to_constraint_types:
+      PointConflict: [ "MultiPointConstraint" ]
+
+  single_agent:
+    class: "DiRecT"
+    planner_alg: "base"
+    horizon: 64
+    trajectory_duration: 5.0
+    n_samples: 128  # Number of samples **per agent**
+    debug: true
+    timeit: false # Enable timing for performance measurement
+    # --- Guidance and Diffusion Parameters ---
+    use_guide_on_extra_objects_only: false
+    n_local_inference_noising_steps: 3
+    n_local_inference_denoising_steps: 3
+    start_guide_steps_fraction: 1.0
+    n_guide_steps: 1
+    n_diffusion_steps_without_noise: 0
+    cost_func_key: "sum_log_l2"
+    cost_func_param:
+      margin_coeff: 1.9
+      radius_coeff: 6.0
+    weight_grad_cost_collision: 0.1
+    weight_grad_cost_smoothness: 0.0
+    weight_grad_cost_constraints: 0.2
+    weight_grad_cost_soft_constraints: 0.0
+    factor_num_interpolated_points_for_collision: 0
+    projector_type: "JaxScpADMM"
+    start_projection_guidance: 0.5 # guidance start
+    low_level_choose_path_from_batch_strategy: "least_collisions"
+
+  projection:
+    projector:
+      _target_: "ADMMProjectionOperatorCUDAGraph"
+      rho: 10.0
+      max_iter: 500
+      tol: 1.0e-5
+      decomp: "lu"
+      convergence_check: true
+      check_period: 200
+      verbose: false
+
+    project_params:
+      vel_max: 0.675
+      smoothness_weight: 0.0

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/configs/direct_highway.yaml

@@ -0,0 +1,80 @@
+experiment:
+  name: "single_tile_highway"
+  num_agents: 2
+  n_inits: 3
+  runtime_limit: 30  # seconds
+  seed: 18
+  time_str: "auto"  # auto timestamp at runtime
+  render_animation: true
+
+system:
+  device: "cuda:0"
+  base_dir: "auto"  # auto resolve base project directory
+  results_dir_name: "logs"
+  trained_models_dir_name: "data_trained_models"
+  log_pth: null  # use default logger path behavior
+
+robot:
+  planar_disk_radius: 0.05
+
+environment:
+  type: "single_tile"
+  global_model_ids: [['EnvHighways2D-RobotPlanarDisk']]
+  agent_skeleton_l: "auto"  # auto-generate skeleton from environment
+  start_goal_config:
+    method: "random_in_env"
+    env_class_for_random: "EnvHighways2D"
+    margin: 0.2
+    obstacle_margin: 0.11
+
+planner:
+  multi_agent:
+    class: "DummyPlanning"
+    stagger_start_time_dt: 0
+    is_xcbs: false
+    is_ecbs: false
+    conflict_type_to_constraint_types:
+      PointConflict: [ "MultiPointConstraint" ]
+
+  single_agent:
+    class: "DiRecT"
+    planner_alg: "base"
+    horizon: 64
+    trajectory_duration: 5.0
+    n_samples: 128
+    debug: true
+    timeit: false
+    use_guide_on_extra_objects_only: false
+    n_local_inference_noising_steps: 3
+    n_local_inference_denoising_steps: 3
+    start_guide_steps_fraction: 1.0
+    n_guide_steps: 1
+    n_diffusion_steps_without_noise: 0
+    cost_func_key: "sum_log_l2"
+    cost_func_param:
+      margin_coeff: 1.9
+      radius_coeff: 6.0
+    weight_grad_cost_collision: 0.1
+    weight_grad_cost_smoothness: 0.0
+    weight_grad_cost_constraints: 0.2
+    weight_grad_cost_soft_constraints: 0.0
+    factor_num_interpolated_points_for_collision: 0
+    projector_type: "JaxScpADMM"
+    start_projection_guidance: 0.5 # 0.25 older but overwritte by run_benchmark.py
+    low_level_choose_path_from_batch_strategy: "least_collisions"
+
+  projection:
+    projector:
+      _target_: "ADMMProjectionOperatorCUDAGraph"
+      rho: 6.0
+      max_iter: 500
+      tol: 1.0e-5
+      decomp: "lu"
+      convergence_check: true
+      check_period: 200
+      verbose: false
+
+    project_params:
+      vel_max: 0.8
+      smoothness_weight: 0.0
+

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/configs/pcdiff_conveyor.yaml

@@ -0,0 +1,90 @@
+# config.yaml
+#
+# This YAML file centralizes all configurations for the multi-robot planning experiment.
+# To run an experiment, execute:
+# python script_name.py --config path/to/this/config.yaml
+
+# General experiment settings
+experiment:
+  name: "single_tile_conveyor"
+  num_agents: 2
+  n_inits: 3
+  runtime_limit: 60  # in seconds
+  seed: 18
+  time_str: "auto" # "auto" will generate from current time, or provide a string like "07132025_100000"
+  render_animation: true
+
+# System and directory settings
+system:
+  device: "cuda:0"
+  base_dir: "auto" # "auto" resolves to the project root, or provide a full path
+  results_dir_name: "logs"
+  trained_models_dir_name: "data_trained_models"
+  log_pth: null # "auto" resolves to the results directory, or provide a full path
+
+# Robot-specific parameters
+robot:
+  planar_disk_radius: 0.05
+
+# Environment and scenario setup
+environment:
+  type: "single_tile"
+  global_model_ids: [['EnvConveyor2D-RobotPlanarDisk']]
+  agent_skeleton_l: "auto"  # auto-generate skeleton from environment
+  start_goal_config:
+    method: "random_in_env"
+    env_class_for_random: "EnvConveyor2D"
+    margin: 0.2
+    obstacle_margin: 0.2
+
+# Planner configurations
+planner:
+  multi_agent:
+    class: "DummyPlanning"
+    stagger_start_time_dt: 0
+    is_xcbs: false
+    is_ecbs: false
+    conflict_type_to_constraint_types:
+      PointConflict: [ "MultiPointConstraint" ]
+
+  # Low-level single-agent trajectory generation planner (the diffusion model)
+  single_agent:
+    class: "PCD"
+    planner_alg: "pcd"
+    horizon: 64
+    trajectory_duration: 5.0
+    n_samples: 128  # Number of samples **per agent** 
+    debug: true
+    timeit: false # Enable timing for performance measurement
+    # --- Guidance and Diffusion Parameters ---
+    use_guide_on_extra_objects_only: false
+    n_local_inference_noising_steps: 3
+    n_local_inference_denoising_steps: 3
+    start_guide_steps_fraction: 1.0
+    n_guide_steps: 1
+    n_diffusion_steps_without_noise: 0
+    cost_func_key: "sum_log_l2"  # Key for the cost function used in guidance
+    cost_func_param: 
+      margin_coeff: 1.9 
+      radius_coeff: 6.0  # Coefficient for the collision cost radius
+    # --- Cost Weights for Guidance ---
+    weight_grad_cost_collision: 0.05
+    weight_grad_cost_smoothness: 0.0
+    weight_grad_cost_constraints: 0.2
+    weight_grad_cost_soft_constraints: 0.0
+    factor_num_interpolated_points_for_collision: 0
+    low_level_choose_path_from_batch_strategy: "least_collisions"
+
+  projection:
+    projector:
+      _target_: "ADMMProjectionOperatorCUDAGraph"
+      rho: 6.0
+      max_iter: 500
+      tol: 1.0e-5
+      decomp: "lu"
+      convergence_check: true
+      check_period: 200
+      verbose: false
+
+    project_params:
+      vel_max: 0.8

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/configs/pcdiff_dropregion.yaml

@@ -0,0 +1,76 @@
+experiment:
+  name: "single_tile_dropregion"
+  num_agents: 2
+  n_inits: 3
+  runtime_limit: 60  # seconds
+  seed: 18
+  time_str: "auto"  # auto timestamp at runtime
+  render_animation: true
+
+system:
+  device: "cuda:0"
+  base_dir: "auto"  # auto resolve base project directory
+  results_dir_name: "logs"
+  trained_models_dir_name: "data_trained_models"
+  log_pth: null  # use default logger path behavior
+
+robot:
+  planar_disk_radius: 0.05
+
+environment:
+  type: "single_tile"
+  global_model_ids: [['EnvDropRegion2D-RobotPlanarDisk']]
+  agent_skeleton_l: "auto"  # auto-generate skeleton from environment
+  start_goal_config:
+    method: "random_in_env"
+    env_class_for_random: "EnvDropRegion2D"
+    margin: 0.2
+    obstacle_margin: 0.2
+
+planner:
+  multi_agent:
+    class: "DummyPlanning"
+    stagger_start_time_dt: 0
+    is_xcbs: false
+    is_ecbs: false
+    conflict_type_to_constraint_types:
+      PointConflict: [ "MultiPointConstraint" ]
+
+  single_agent:
+    class: "PCD"
+    planner_alg: "pcd"
+    horizon: 64
+    trajectory_duration: 5.0
+    n_samples: 128
+    debug: true
+    timeit: false
+    use_guide_on_extra_objects_only: false
+    n_local_inference_noising_steps: 3
+    n_local_inference_denoising_steps: 3
+    start_guide_steps_fraction: 1.0
+    n_guide_steps: 1
+    n_diffusion_steps_without_noise: 0
+    cost_func_key: "sum_log_l2"
+    cost_func_param:
+      margin_coeff: 1.9
+      radius_coeff: 6.0
+    weight_grad_cost_collision: 0.05
+    weight_grad_cost_smoothness: 0.0
+    weight_grad_cost_constraints: 0.2
+    weight_grad_cost_soft_constraints: 0.0
+    factor_num_interpolated_points_for_collision: 0
+    low_level_choose_path_from_batch_strategy: "least_collisions"
+
+  projection:
+    projector:
+      _target_: "ADMMProjectionOperatorCUDAGraph"
+      rho: 6.0
+      max_iter: 500
+      tol: 1.0e-5
+      decomp: "lu"
+      convergence_check: true
+      check_period: 200
+      verbose: false
+
+    project_params:
+      vel_max: 0.8

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/configs/pcdiff_empty.yaml

@@ -0,0 +1,77 @@
+experiment:
+  name: "single_tile_empty"
+  num_agents: 8
+  n_inits: 10
+  runtime_limit: 60  # seconds
+  seed: 18
+  time_str: "auto"  # auto timestamp at runtime
+  render_animation: true
+
+system:
+  device: "cuda:0"
+  base_dir: "auto"  # auto resolve base project directory
+  results_dir_name: "logs"
+  trained_models_dir_name: "data_trained_models"
+  log_pth: null  # use default logger path behavior
+
+robot:
+  planar_disk_radius: 0.05
+
+environment:
+  type: "single_tile"
+  global_model_ids: [['EnvEmpty2D-RobotPlanarDisk']]
+  agent_skeleton_l: "auto"  # auto-generate skeleton from environment
+  start_goal_config:
+    method: "random_in_env"
+    env_class_for_random: "EnvEmpty2D"
+    margin: 0.2
+    obstacle_margin: 0.11
+
+planner:
+  multi_agent:
+    class: "DummyPlanning"
+    stagger_start_time_dt: 0
+    is_xcbs: false
+    is_ecbs: false
+    conflict_type_to_constraint_types:
+      PointConflict: [ "MultiPointConstraint" ]
+
+  single_agent:
+    class: "PCD"
+    planner_alg: "pcd"
+    horizon: 64
+    trajectory_duration: 5.0
+    n_samples: 128  # Number of samples **per agent** 
+    debug: true
+    timeit: false # Enable timing for performance measurement
+    # --- Guidance and Diffusion Parameters ---
+    use_guide_on_extra_objects_only: false
+    n_local_inference_noising_steps: 3
+    n_local_inference_denoising_steps: 3
+    start_guide_steps_fraction: 1.0
+    n_guide_steps: 1
+    n_diffusion_steps_without_noise: 0
+    cost_func_key: "sum_log_l2"
+    cost_func_param:
+      margin_coeff: 1.9
+      radius_coeff: 6.0
+    weight_grad_cost_collision: 0.1
+    weight_grad_cost_smoothness: 0.0
+    weight_grad_cost_constraints: 0.2
+    weight_grad_cost_soft_constraints: 0.0
+    factor_num_interpolated_points_for_collision: 0
+    low_level_choose_path_from_batch_strategy: "least_collisions"
+
+  projection:
+    projector:
+      _target_: "ADMMProjectionOperatorCUDAGraph"
+      rho: 10.0
+      max_iter: 500
+      tol: 1.0e-5
+      decomp: "lu"
+      convergence_check: true
+      check_period: 200
+      verbose: false
+
+    project_params:
+      vel_max: 0.675

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/configs/pcdiff_highway.yaml

@@ -0,0 +1,76 @@
+experiment:
+  name: "single_tile_highway"
+  num_agents: 2
+  n_inits: 3
+  runtime_limit: 30  # seconds
+  seed: 18
+  time_str: "auto"  # auto timestamp at runtime
+  render_animation: true
+
+system:
+  device: "cuda:0"
+  base_dir: "auto"  # auto resolve base project directory
+  results_dir_name: "logs"
+  trained_models_dir_name: "data_trained_models"
+  log_pth: null  # use default logger path behavior
+
+robot:
+  planar_disk_radius: 0.05
+
+environment:
+  type: "single_tile"
+  global_model_ids: [['EnvHighways2D-RobotPlanarDisk']]
+  agent_skeleton_l: "auto"  # auto-generate skeleton from environment
+  start_goal_config:
+    method: "random_in_env"
+    env_class_for_random: "EnvHighways2D"
+    margin: 0.2
+    obstacle_margin: 0.11
+
+planner:
+  multi_agent:
+    class: "DummyPlanning"
+    stagger_start_time_dt: 0
+    is_xcbs: false
+    is_ecbs: false
+    conflict_type_to_constraint_types:
+      PointConflict: [ "MultiPointConstraint" ]
+
+  single_agent:
+    class: "PCD"
+    planner_alg: "pcd"
+    horizon: 64
+    trajectory_duration: 5.0
+    n_samples: 128
+    debug: true
+    timeit: false
+    use_guide_on_extra_objects_only: false
+    n_local_inference_noising_steps: 3
+    n_local_inference_denoising_steps: 3
+    start_guide_steps_fraction: 1.0
+    n_guide_steps: 1
+    n_diffusion_steps_without_noise: 0
+    cost_func_key: "sum_log_l2"
+    cost_func_param:
+      margin_coeff: 1.9
+      radius_coeff: 6.0
+    weight_grad_cost_collision: 0.1
+    weight_grad_cost_smoothness: 0.0
+    weight_grad_cost_constraints: 0.2
+    weight_grad_cost_soft_constraints: 0.0
+    factor_num_interpolated_points_for_collision: 0
+    low_level_choose_path_from_batch_strategy: "least_collisions"
+
+  projection:
+    projector:
+      _target_: "ADMMProjectionOperatorCUDAGraph"
+      rho: 6.0
+      max_iter: 500
+      tol: 1.0e-5
+      decomp: "lu"
+      convergence_check: true
+      check_period: 200
+      verbose: false
+
+    project_params:
+      vel_max: 0.8

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020 Common
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/MANIFEST.in

@@ -0,0 +1,2 @@
+include requirements.txt
+include setup.py

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/Makefile

@@ -0,0 +1,14 @@
+package: 
+	python setup.py sdist
+
+install: package
+	pip install $(shell ls dist/*.tar.gz)
+
+all: clean install
+
+upload:
+	python -m twine upload dist/*
+
+clean:
+	rm -rf dist
+	rm -rf build

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/README.md

@@ -0,0 +1,111 @@
+# Experiment Launcher
+
+Launch experiments locally or on a cluster running SLURM in a single file.  
+
+## Description 
+
+The ``experiment_launcher`` package provides a simple way to run multiple experiments using SLURM or Joblib with minimum 
+effort - you just have to set the `local` parameter to `True` to run locally,
+or to `False` to run with a cluster using SLURM. 
+
+It is particularly useful to run multiple seeds and/or test multiple configurations of hyperparameters such as learning rates, batch size.
+
+## Installation
+
+You can do a minimal installation of ``experiment_launcher`` with:
+```
+pip install  -e .
+```
+
+## How to Use
+
+### Basic Usage
+
+The best way to understand experiment launcher is to look at the basic example in [examples/basic](examples/basic).
+
+- The file [examples/basic/test.py](examples/basic/test.py) contains a single experiment
+- The file [examples/basic/launch_test.py](examples/basic/launch_test.py) contains the launcher file configurations and call
+
+**Single experiment**
+- [examples/basic/test.py](examples/basic/test.py) consists of:
+  - The function `experiment` is the entry point of your experiment
+    - It takes as arguments your experiment settings (e.g., the number of layers in a neural 
+        network, the learning rate, ...)
+    - The arguments need to be assigned a type and default value in the function definition
+      - Current accepted types are `int`, `float`, `str`, `bool`, `list`
+      - The arguments `seed` and `results_dir` **must** always be included
+      - Python kwargs can also be added as `**kwargs` (accepted types are the same as above)
+    - This function **must** be decorated with a decorator, e.g. `@single_experiment`. 
+      This will take care of creating proper results directories.
+ 
+  - The `if __name__ == '__main__'` block
+      - This **must** contain one single line: `run_experiment(experiment)`
+- You can test your code by running
+    ```bash
+    cd examples/basic
+    python test.py
+    ```
+
+**Launch file**
+- [examples/basic/launch_test.py](examples/basic/launch_test.py) consists of:
+  - Creating an instance of the `Launcher` object, that contains the SLURM or Joblib (if run locally) parameters. These are some of the important parameters. For more consult the class definition.
+    - `exp_name` is the experiment name, under which a results directory will be created
+    - `exp_file` is the path to the python file where the `experiment` is implemented (without the extension `.py`)
+    - `project_name` is the project name in the cluster
+    - `n_seeds` is the number of random seeds for each single experiment configuration
+    - `n_exps_in_parallel` is the number of experiments to be run in parallel. This is useful to run multiple jobs in a single GPU in the cluster
+    - `n_cores` is the number of cores for each experiment.
+       Note that if `n_exps_in_parallel > 1`, then `n_exps_in_parallel` jobs will share `n_cores`.
+       This parameter does not apply if running locally. For that you should use the linux comamnd [taskset](https://man7.org/linux/man-pages/man1/taskset.1.html)
+    - `memory_per_core` is amount of memory in MB requested for each core in SLURM. If you specify this too low, SLURM might crash.
+    - `partition` is the SLURM partition, which is cluster dependent
+    - `gres` are special resources asked for a SLURM experiment
+    - `conda_env` if you are using a conda environment, specify its name here
+  - Adding experiments with `launcher.add_experiment`
+    - Use `launcher.add_experiment` to create an experiment for a particular configuration (e.g., different learning rates)
+    - Parameters that are followed by _**two trailing underscores**_ will be used to create experiment directories
+      - Results directories are going to be created by default as, e.g. with `param__` and `param2__`
+        - `./logs/exp_name_DATE/param1__value/param2__value/SEED/` 
+      - **Be careful** when specifying the parameters and double underscores. 
+        The launcher will throw an error if two experiments lead to the same results directory (which would be overwritten).
+      - To see this problem, uncomment the line `# dict(floating_arg=11.0, string_arg='test')` in [examples/basic/launch_test.py](examples/basic/launch_test.py)
+      - Parameters will with trailing underscores have priority over their counterparts without underscores. E.g., `env='some_env'` will be overwritten by `env__='my_env'`
+  - Running the experiments with `launcher.run(LOCAL)`
+    - This runs your experiment either locally (LOCAL: `True`) or in the cluster (LOCAL: `False`) 
+
+
+**Running the experiment**
+- To run the launcher simply call 
+  ```bash
+  cd examples/basic
+  python launch_test.py
+  ```
+- Log files will be placed in 
+  - `./logs` if running locally or in the IAS cluster
+  - `/work/scratch/$USERNAME` if ran in the `Lichtenberg-Hochleistungsrechner of the TU Darmstadt` 
+
+
+### Integration with Weights and Biases
+
+The experiment launcher provides an easy way to integrate with Weights and Biases.
+- In the experiment file add
+  - `**kwargs` in the `experiment` function definition
+- In the launcher file, create wandb options and pass them to the `launcher.add_experiment()`
+  ```python
+  wandb_options = dict(
+    wandb_enabled=False,  # If True, runs and logs to wandb.
+    wandb_entity='joaocorreiacarvalho',
+    wandb_project='experiment_launcher_test',
+    wandb_group='group_test'
+  )
+  ```
+  
+
+### Running experiments with parameters in configuration files
+- If you have many parameters that you need to change, a good idea can be to use configuration files. 
+- If you want to specify your parameters in a configuration file, look at the example under `examples/config_files`
+  
+
+## Notes
+- For reproducibility, the seeds are created sequentially from `0` to `n_exps-1`.
+

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/examples/basic/launch_test.py

@@ -0,0 +1,75 @@
+from itertools import product
+
+from experiment_launcher import Launcher, is_local
+
+LOCAL = is_local()
+TEST = False
+USE_CUDA = False
+
+N_SEEDS = 3
+
+if LOCAL:
+    N_EXPS_IN_PARALLEL = 5
+else:
+    N_EXPS_IN_PARALLEL = 3
+
+N_CORES = N_EXPS_IN_PARALLEL
+MEMORY_SINGLE_JOB = 1000
+MEMORY_PER_CORE = N_EXPS_IN_PARALLEL * MEMORY_SINGLE_JOB // N_CORES
+PARTITION = 'amd2,amd'  # 'amd', 'rtx'
+GRES = 'gpu:1' if USE_CUDA else None  # gpu:rtx2080:1, gpu:rtx3080:1
+CONDA_ENV = 'el'  # None
+
+launcher = Launcher(
+    exp_name='test_launcher',
+    exp_file='test',
+    # project_name='project01234',  # for hrz cluster
+    n_seeds=N_SEEDS,
+    n_exps_in_parallel=N_EXPS_IN_PARALLEL,
+    n_cores=N_CORES,
+    memory_per_core=MEMORY_PER_CORE,
+    days=2,
+    hours=23,
+    minutes=59,
+    seconds=0,
+    partition=PARTITION,
+    gres=GRES,
+    conda_env=CONDA_ENV,
+    use_timestamp=True,
+    compact_dirs=False
+)
+
+envs = {
+    'env_00': {'env_param': 'aa'},
+    'env_01': {'env_param': 'bb'}
+}
+a_l = [1, 2, 3]
+boolean_param_l = [True, False]
+some_default_param = 'b'
+
+# These arguments are kwargs of the experiment function
+unknown_args_list = [
+    dict(integer_arg=10),
+    # dict(floating_arg=11.0, string_arg='test')
+]
+
+for env in envs:
+    d = envs[env]
+    for a, boolean_param in product(a_l, boolean_param_l):
+        for unknown_args in unknown_args_list:
+            launcher.add_experiment(
+                # A subdirectory will be created for parameters with a trailing double underscore.
+                env__=env,
+                a__=a,
+                boolean_param__=boolean_param,
+
+                env='some_env',  # This value will be overwritten by env__
+
+                **d,
+                some_default_param=some_default_param,
+
+                **unknown_args,
+                debug=False,
+            )
+
+launcher.run(LOCAL, TEST)

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/examples/basic/test.py

@@ -0,0 +1,47 @@
+import os
+
+import wandb
+
+from experiment_launcher import run_experiment, single_experiment
+
+
+# This decorator creates results_dir as results_dir/seed, and saves the experiment arguments into a file.
+@single_experiment
+def experiment(
+    #######################################
+    env: str = 'env',  # You need to specify the argument type if you use the automatic parser.
+    env_param: str = 'aa',
+    a: int = 1,
+    boolean_param: bool = True,
+    some_default_param: str = 'b',
+
+    debug: bool = True,
+
+    #######################################
+    # MANDATORY
+    seed: int = 0,
+    results_dir: str = 'logs',
+
+    #######################################
+    # OPTIONAL
+    # accept unknown arguments
+    **kwargs
+):
+    # EXPERIMENT
+    print(f'DEBUG MODE: {debug}')
+
+    print(f'kwargs: {kwargs}')
+
+    filename = os.path.join(results_dir, 'log_' + str(seed) + '.txt')
+    out_str = f'Running experiment with seed {seed}, env {env}, ' \
+              f'env_param {env_param}, a {a}, ' \
+              f'boolean_param {boolean_param}, some_default_param {some_default_param}'
+    print(out_str)
+    with open(filename, 'w') as file:
+        file.write('Some logs in a log file.\n')
+        file.write(out_str)
+
+
+if __name__ == '__main__':
+    # Leave unchanged
+    run_experiment(experiment)

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/examples/config_files/configs/config00.yaml

@@ -0,0 +1,7 @@
+network: mynetwork
+layers: 3
+activations:
+  - relu
+  - relu
+  - tanh
+lr: 1e-3

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/examples/config_files/configs/config01.yaml

@@ -0,0 +1,7 @@
+network: myothernetwork
+layers: 3
+activations:
+  - relu
+  - relu
+  - tanh
+lr: 1e-3

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/examples/config_files/launch_test.py

@@ -0,0 +1,65 @@
+from experiment_launcher import Launcher, is_local
+
+LOCAL = is_local()
+TEST = False
+USE_CUDA = False
+
+N_SEEDS = 3
+
+if LOCAL:
+    N_EXPS_IN_PARALLEL = 5
+else:
+    N_EXPS_IN_PARALLEL = 3
+
+N_CORES = N_EXPS_IN_PARALLEL
+MEMORY_SINGLE_JOB = 1000
+MEMORY_PER_CORE = N_EXPS_IN_PARALLEL * MEMORY_SINGLE_JOB // N_CORES
+PARTITION = 'amd2,amd'  # 'amd', 'rtx'
+GRES = 'gpu:1' if USE_CUDA else None  # gpu:rtx2080:1, gpu:rtx3080:1
+CONDA_ENV = 'el'  # None
+
+launcher = Launcher(
+    exp_name='test_launcher',
+    exp_file='test',
+    # project_name='project01234',  # for hrz cluster
+    n_seeds=N_SEEDS,
+    n_exps_in_parallel=N_EXPS_IN_PARALLEL,
+    n_cores=N_CORES,
+    memory_per_core=MEMORY_PER_CORE,
+    days=2,
+    hours=23,
+    minutes=59,
+    seconds=0,
+    partition=PARTITION,
+    conda_env=CONDA_ENV,
+    gres=GRES,
+    use_timestamp=True,
+    compact_dirs=False
+)
+
+config_files_l = [
+    'configs/config00.yaml',
+    'configs/config01.yaml',
+]
+
+# Optional arguments for Weights and Biases
+wandb_options = dict(
+    wandb_mode='disabled',  # "online", "offline" or "disabled"
+    wandb_entity='joaocorreiacarvalho',
+    wandb_project='test_experiment_launcher_config_files'
+)
+
+for i, config_file in enumerate(config_files_l):
+    launcher.add_experiment(
+        # A subdirectory will be created for parameters with a trailing double underscore.
+        config__=f'config-{str(i).zfill(len(config_files_l))}',
+
+        config_file_path=config_file,
+
+        debug=False,
+
+        **wandb_options,
+        wandb_group=f'test_group-el-{config_file}'
+    )
+
+launcher.run(LOCAL, TEST)

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/examples/config_files/test.py

@@ -0,0 +1,49 @@
+import os
+
+import wandb
+import yaml
+
+from experiment_launcher import run_experiment, single_experiment_yaml
+
+
+# This decorator creates results_dir as results_dir/seed, and saves the experiment arguments into a file.
+@single_experiment_yaml
+def experiment(
+    #######################################
+    config_file_path: str = './configs/config00.yaml',
+
+    some_default_param: str = 'b',
+
+    debug: bool = True,
+
+    #######################################
+    # MANDATORY
+    seed: int = 0,
+    results_dir: str = 'logs',
+
+    #######################################
+    # OPTIONAL
+    # accept unknown arguments
+    **kwargs
+):
+    # EXPERIMENT
+    print(f'DEBUG MODE: {debug}')
+
+    with open(config_file_path, 'r') as f:
+        configs = yaml.load(f, yaml.Loader)
+
+    print('Config file content:')
+    print(configs)
+
+    filename = os.path.join(results_dir, 'log_' + str(seed) + '.txt')
+    out_str = f'Running experiment with seed {seed}'
+    with open(filename, 'w') as file:
+        file.write('Some logs in a log file.\n')
+        file.write(out_str)
+
+    wandb.log({'seed': seed}, step=1)
+
+
+if __name__ == '__main__':
+    # Leave unchanged
+    run_experiment(experiment)

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/examples/multi_exp_same_gpu/launch_test.py

@@ -0,0 +1,60 @@
+from experiment_launcher import Launcher, is_local
+
+LOCAL = is_local()
+TEST = False
+USE_CUDA = True
+
+N_SEEDS = 10
+
+if LOCAL:
+    N_EXPS_IN_PARALLEL = 2
+else:
+    N_EXPS_IN_PARALLEL = 3
+
+N_CORES = N_EXPS_IN_PARALLEL
+MEMORY_SINGLE_JOB = 1000
+MEMORY_PER_CORE = N_EXPS_IN_PARALLEL * MEMORY_SINGLE_JOB // N_CORES
+PARTITION = 'rtx2,rtx' if USE_CUDA else 'amd2,amd'
+GRES = 'gpu:1' if USE_CUDA else None  # gpu:rtx2080:1, gpu:rtx3080:1
+CONDA_ENV = 'el'  # None
+
+launcher = Launcher(
+    exp_name='test_launcher',
+    exp_file='test',
+    # project_name='project01234',  # for hrz cluster
+    n_seeds=N_SEEDS,
+    n_exps_in_parallel=N_EXPS_IN_PARALLEL,
+    n_cores=N_CORES,
+    memory_per_core=MEMORY_PER_CORE,
+    days=2,
+    hours=23,
+    minutes=59,
+    seconds=0,
+    partition=PARTITION,
+    conda_env=CONDA_ENV,
+    gres=GRES,
+    use_timestamp=True,
+    compact_dirs=False
+)
+
+tensor_sizes_l = [1000000 + i for i in range(5)]
+
+# Optional arguments for Weights and Biases
+wandb_options = dict(
+    wandb_mode='disabled',  # "online", "offline" or "disabled"
+    wandb_entity='joaocorreiacarvalho',
+    wandb_project='test_experiment_launcher_config_files'
+)
+
+for tensor_size in tensor_sizes_l:
+    launcher.add_experiment(
+        # A subdirectory will be created for parameters with a trailing double underscore.
+        tensor_size__=tensor_size,
+
+        debug=False,
+
+        **wandb_options,
+        wandb_group=f'test_group-el-{tensor_size}',
+    )
+
+launcher.run(LOCAL, TEST)

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/examples/multi_exp_same_gpu/test.py

@@ -0,0 +1,54 @@
+import os
+import time
+
+import torch.cuda
+import wandb
+
+from experiment_launcher import run_experiment, single_experiment_yaml
+
+
+# This decorator creates results_dir as results_dir/seed, and saves the experiment arguments into a file.
+@single_experiment_yaml
+def experiment(
+    #######################################
+    env: str = 'env-name',  # You need to specify the argument type if you use the automatic parser.
+    tensor_size: int = 10000,
+
+    debug: bool = True,
+
+    #######################################
+    # MANDATORY
+    seed: int = 0,
+    results_dir: str = 'logs',
+
+    #######################################
+    # OPTIONAL
+    # accept unknown arguments
+    **kwargs
+):
+    # EXPERIMENT
+    print(f'DEBUG MODE: {debug}')
+
+    tensor1 = torch.ones(tensor_size, device='cuda' if torch.cuda.is_available() else 'cpu')
+    tensor2 = tensor1.clone()
+
+    print(f'Env: {env}')
+    print(f'Tensor shape: {tensor1.shape}')
+
+    filename = os.path.join(results_dir, 'log_' + str(seed) + '.txt')
+    out_str = f'Running experiment with seed {seed}'
+    with open(filename, 'w') as file:
+        file.write('Some logs in a log file.\n')
+        file.write(out_str)
+
+    wandb.log({'seed': seed}, step=1)
+
+    end_time = time.time() + 0.01
+    while time.time() < end_time:
+        tensor = tensor1 * tensor2
+    print('Finished experiment.')
+
+
+if __name__ == '__main__':
+    # Leave unchanged
+    run_experiment(experiment)

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/examples/start_seed/launch_test.py

@@ -0,0 +1,98 @@
+import time
+from itertools import product
+
+from experiment_launcher import Launcher, is_local
+
+LOCAL = is_local()
+TEST = False
+USE_CUDA = False
+
+N_SEEDS = 103
+
+if LOCAL:
+    N_EXPS_IN_PARALLEL = 5
+else:
+    N_EXPS_IN_PARALLEL = 3
+
+N_CORES = N_EXPS_IN_PARALLEL
+MEMORY_SINGLE_JOB = 1000
+MEMORY_PER_CORE = N_EXPS_IN_PARALLEL * MEMORY_SINGLE_JOB // N_CORES
+PARTITION = 'amd2,amd'  # 'amd', 'rtx'
+GRES = 'gpu:1' if USE_CUDA else None  # gpu:rtx2080:1, gpu:rtx3080:1
+CONDA_ENV = 'el'  # None
+
+
+# E.g. SLURM only allows arrays with N_SEEDS_MAX_SLURM size
+N_SEEDS_MAX_SLURM = 100
+start_seed = 0
+if LOCAL:
+    n_seeds_split = [N_SEEDS]
+else:
+    if N_SEEDS <= N_SEEDS_MAX_SLURM:
+        n_seeds_split = [N_SEEDS]
+    elif N_SEEDS % N_SEEDS_MAX_SLURM == 0:
+        n_seeds_split = [N_SEEDS_MAX_SLURM for _ in range(N_SEEDS // N_SEEDS_MAX_SLURM)]
+    else:
+        n_seeds_split = [N_SEEDS_MAX_SLURM for _ in range(N_SEEDS // N_SEEDS_MAX_SLURM)] + [N_SEEDS % N_SEEDS_MAX_SLURM]
+
+for _n_seeds in n_seeds_split:
+    print(f"---------> Running seeds {start_seed}-{start_seed + _n_seeds}")
+
+    time.sleep(1)
+    launcher = Launcher(
+        exp_name='test_launcher',
+        exp_file='test',
+        # project_name='project01234',  # for hrz cluster
+        n_seeds=_n_seeds,
+        start_seed=start_seed,
+        n_exps_in_parallel=N_EXPS_IN_PARALLEL,
+        n_cores=N_CORES,
+        memory_per_core=MEMORY_PER_CORE,
+        days=2,
+        hours=23,
+        minutes=59,
+        seconds=0,
+        partition=PARTITION,
+        gres=GRES,
+        conda_env=CONDA_ENV,
+        use_timestamp=True,
+        compact_dirs=False
+    )
+
+    envs = {
+        'env_00': {'env_param': 'aa'},
+        'env_01': {'env_param': 'bb'}
+    }
+    a_l = [1, 2, 3]
+    boolean_param_l = [True, False]
+    some_default_param = 'b'
+
+    # These arguments are kwargs of the experiment function
+    unknown_args_list = [
+        dict(integer_arg=10),
+        # dict(floating_arg=11.0, string_arg='test')
+    ]
+
+    for env in envs:
+        d = envs[env]
+        for a, boolean_param in product(a_l, boolean_param_l):
+            for unknown_args in unknown_args_list:
+                launcher.add_experiment(
+                    # A subdirectory will be created for parameters with a trailing double underscore.
+                    env__=env,
+                    a__=a,
+                    boolean_param__=boolean_param,
+
+                    env='some_env',  # This value will be overwritten by env__
+
+                    **d,
+                    some_default_param=some_default_param,
+
+                    **unknown_args,
+                    debug=False,
+                )
+
+    launcher.run(LOCAL, TEST)
+
+    ################################################################################################################
+    start_seed += _n_seeds

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/examples/start_seed/test.py

@@ -0,0 +1,47 @@
+import os
+
+import wandb
+
+from experiment_launcher import run_experiment, single_experiment
+
+
+# This decorator creates results_dir as results_dir/seed, and saves the experiment arguments into a file.
+@single_experiment
+def experiment(
+    #######################################
+    env: str = 'env',  # You need to specify the argument type if you use the automatic parser.
+    env_param: str = 'aa',
+    a: int = 1,
+    boolean_param: bool = True,
+    some_default_param: str = 'b',
+
+    debug: bool = True,
+
+    #######################################
+    # MANDATORY
+    seed: int = 0,
+    results_dir: str = 'logs',
+
+    #######################################
+    # OPTIONAL
+    # accept unknown arguments
+    **kwargs
+):
+    # EXPERIMENT
+    print(f'DEBUG MODE: {debug}')
+
+    print(f'kwargs: {kwargs}')
+
+    filename = os.path.join(results_dir, 'log_' + str(seed) + '.txt')
+    out_str = f'Running experiment with seed {seed}, env {env}, ' \
+              f'env_param {env_param}, a {a}, ' \
+              f'boolean_param {boolean_param}, some_default_param {some_default_param}'
+    print(out_str)
+    with open(filename, 'w') as file:
+        file.write('Some logs in a log file.\n')
+        file.write(out_str)
+
+
+if __name__ == '__main__':
+    # Leave unchanged
+    run_experiment(experiment)

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/experiment_launcher/__init__.py

@@ -0,0 +1,8 @@
+try:
+    from .launcher import Launcher, run_experiment
+    from .decorators import single_experiment, single_experiment_yaml, single_experiment_flat
+    from .utils import is_local
+except ImportError:
+    pass
+
+__version__ = '2.3'

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/experiment_launcher/decorators.py

@@ -0,0 +1,49 @@
+import logging
+import os
+from functools import wraps, partial
+
+from experiment_launcher.utils import save_args, start_wandb, create_results_dir
+
+
+def wrapper_single_experiment(exp_func, save_args_yaml=False, use_logging=False,
+                              make_dirs_with_seed=True, print_exp_args=False):
+    @wraps(exp_func)
+    def wrapper(
+        # Function arguments
+        *args,
+        **kwargs
+    ):
+        # Make results directory
+        create_results_dir(kwargs, make_dirs_with_seed)
+
+        # Setup logging
+        if use_logging:
+            logging.basicConfig(level=logging.INFO,
+                                filename=os.path.join(kwargs['results_dir'], "logfile"),
+                                # stream=sys.stdout,
+                                filemode="a+",
+                                format="%(asctime)-15s %(levelname)-8s %(message)s")
+
+        # Save arguments
+        save_args(kwargs['results_dir'], kwargs,
+                  git_repo_path='./', save_args_as_yaml=save_args_yaml, print_exp_args=print_exp_args)
+
+        # Start WandB
+        wandb_activated = True if 'wandb_mode' in kwargs else False
+        if wandb_activated:
+            wandb_run = start_wandb(**kwargs)
+
+        # Run the experiment
+        exp_func(*args, **kwargs)
+
+        # Clean up
+        if wandb_activated:
+            wandb_run.finish()
+
+    return wrapper
+
+
+single_experiment = partial(wrapper_single_experiment)
+single_experiment_flat = partial(wrapper_single_experiment, make_dirs_with_seed=False)
+single_experiment_yaml = partial(wrapper_single_experiment, save_args_yaml=True, use_logging=True, print_exp_args=False)
+single_experiment_flat_yaml = partial(wrapper_single_experiment, make_dirs_with_seed=False, save_args_yaml=True)

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/experiment_launcher/exceptions.py

@@ -0,0 +1,10 @@
+class ResultsDirException(Exception):
+    """
+    Raised when the results directory already exists
+    """
+    def __init__(self, exp, results_dir='./logs'):
+        message = f"\n" \
+                  f"When trying to add the experiment: {exp}\n" \
+                  f"Results directory {results_dir} already exists.\n" \
+                  f"Make sure that varying parameters have a trailing double underscore."
+        super().__init__(message)

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/experiment_launcher/launcher.py

@@ -0,0 +1,495 @@
+import argparse
+import datetime
+import inspect
+import os
+import traceback
+from copy import copy, deepcopy
+from distutils.util import strtobool
+from importlib import import_module
+
+import numpy as np
+from joblib import Parallel, delayed
+
+from experiment_launcher.exceptions import ResultsDirException
+
+
+class Launcher(object):
+    """
+    Creates and starts jobs with Joblib or SLURM.
+
+    """
+
+    def __init__(self, exp_name, exp_file, n_seeds, start_seed=0, n_cores=1, memory_per_core=2000,
+                 days=0, hours=24, minutes=0, seconds=0,
+                 project_name=None, base_dir=None,
+                 n_exps_in_parallel=1,
+                 conda_env=None, gres=None, constraint=None, partition=None,
+                 begin=None, use_timestamp=True, compact_dirs=False):
+        """
+        Constructor.
+
+        Args:
+            exp_name (str): name of the experiment
+            exp_file (str): name of the python module running a single experiment (relative path)
+            n_seeds (int): number of seeds for each experiment configuration
+            start_seed (int): first seed
+            n_cores (int): number of cpu cores
+            memory_per_core (int): maximum memory per core (slurm will kill the job if this is reached)
+            days (int): number of days the experiment can last (in slurm)
+            hours (int): number of hours the experiment can last (in slurm)
+            minutes (int): number of minutes the experiment can last (in slurm)
+            seconds (int): number of seconds the experiment can last (in slurm)
+            project_name (str): name of the project for slurm. This is important if you have
+                different projects (e.g. in the hhlr cluster)
+            base_dir (str): path to directory to save the results (in hhlr results are saved to /work/scratch/$USER)
+            n_exps_in_parallel (int): number of experiment configurations to run in parallel.
+                If running in the cluster, and the gpu is selected, then it is the number of jobs in each slurm file
+                (e.g. for multiple experiments in the same gpu)
+            conda_env (str): name of the conda environment to run the experiments in
+            gres (str): request cluster resources. E.g. to add a GPU in the IAS cluster specify gres='gpu:rtx2080:1'
+            constraint (str): constraint for the slurm job. E.g. to add a GPU in the IAS cluster: constraint='rtx2080'
+            partition (str, None): the partition to use in case of slurm execution. If None, no partition is specified.
+            begin (str): start the slurm experiment at a given time (see --begin in slurm docs)
+            use_timestamp (bool): add a timestamp to the experiment name
+            compact_dirs (bool): If true, only the parameter value is used for the directory name.
+
+        """
+        self._exp_name = exp_name
+        self._exp_file = exp_file
+        self._start_seed = start_seed
+        self._n_seeds = n_seeds
+        self._n_cores = n_cores
+        self._memory_per_core = memory_per_core
+        self._duration = Launcher._to_duration(days, hours, minutes, seconds)
+        self._project_name = project_name
+        self._n_exps_in_parallel = n_exps_in_parallel
+        self._conda_env = conda_env
+        self._gres = gres
+        self._constraint = constraint
+        self._partition = partition
+        self._begin = begin
+
+        self._experiment_list = list()
+
+        if use_timestamp:
+            self._exp_name += datetime.datetime.now().strftime('_%Y-%m-%d_%H-%M-%S')
+
+        base_dir = './logs' if base_dir is None else base_dir
+        self._exp_dir_local = os.path.join(base_dir, self._exp_name)
+
+        # tracks the results directories
+        self._results_dir_l = []
+
+        self._exp_dir_slurm = self._exp_dir_local
+        if os.getenv("USER"):
+            scratch_dir = os.path.join('/work', 'scratch', os.getenv('USER'))
+            if os.path.isdir(scratch_dir):
+                self._exp_dir_slurm = os.path.join(scratch_dir, self._exp_name)
+        os.makedirs(self._exp_dir_slurm)
+
+        # directories for slurm sbatch files and logs
+        self._exp_dir_slurm_files = os.path.join(self._exp_dir_slurm, "slurm_files")
+        self._exp_dir_slurm_logs = os.path.join(self._exp_dir_slurm, "slurm_logs")
+
+        self._compact_dirs = compact_dirs
+
+    def add_experiment(self, **kwargs):
+        self._experiment_list.append(deepcopy(kwargs))
+
+    def run(self, local, test=False, sequential=False):
+        self._check_experiments_results_directories()
+        if local:
+            if sequential:
+                self._run_sequential(test)
+            else:
+                self._run_local(test)
+        else:
+            self._run_slurm(test)
+
+        self._experiment_list = list()
+
+    def generate_slurm(self, command_line_list=None):
+        project_name_option = ''
+        partition_option = ''
+        begin_option = ''
+        gres_option = ''
+        constraint_option = ''
+
+        if self._project_name:
+            project_name_option = '#SBATCH -A ' + self._project_name + '\n'
+        if self._partition:
+            partition_option += f'#SBATCH -p {self._partition}\n'
+        if self._begin:
+            begin_option += f'#SBATCH --begin={self._begin}\n'
+        if self._gres:
+            print(self._gres)
+            gres_option += '#SBATCH --gres=' + str(self._gres) + '\n'
+        if self._constraint:
+            print(self._constraint)
+            constraint_option += '#SBATCH --constraint=' + str(self._constraint) + '\n'
+
+        conda_code = ''
+        if self._conda_env:
+            if os.path.exists(f'{os.getenv("HOME")}/miniconda3'):
+                conda_code += f'eval \"$({os.getenv("HOME")}/miniconda3/bin/conda shell.bash hook)\"\n'
+            elif os.path.exists(f'{os.getenv("HOME")}/anaconda3'):
+                conda_code += f'eval \"$({os.getenv("HOME")}/anaconda/bin/conda shell.bash hook)\"\n'
+            else:
+                raise Exception('You do not have a /home/USER/miniconda3 or /home/USER/anaconda3 directories')
+            conda_code += f'conda activate {self._conda_env}\n\n'
+            python_code = f'python {self._exp_file_path} \\'
+        else:
+            python_code = f'python3  {self._exp_file_path} \\'
+
+        experiment_args = '\t\t'
+        experiment_args += r'${@: 2}'
+        experiment_args += ' \\'
+
+        seed_code = f'\t\t--seed $(({self._start_seed} + $SLURM_ARRAY_TASK_ID)) \\'
+        result_dir_code = '\t\t--results_dir $1'
+
+        code = f"""\
+#!/usr/bin/env bash
+
+###############################################################################
+# SLURM Configurations
+
+# Optional parameters
+{project_name_option}{partition_option}{begin_option}{gres_option}{constraint_option}
+# Mandatory parameters
+#SBATCH -J {self._exp_name}
+#SBATCH -a 0-{self._n_seeds - 1}
+#SBATCH -t {self._duration}
+#SBATCH --ntasks 1
+#SBATCH --cpus-per-task {self._n_cores}
+#SBATCH --mem-per-cpu={self._memory_per_core}
+#SBATCH -o {self._exp_dir_slurm_logs}/%A_%a.out
+#SBATCH -e {self._exp_dir_slurm_logs}/%A_%a.err
+
+###############################################################################
+# Your PROGRAM call starts here
+echo "Starting Job $SLURM_JOB_ID, Index $SLURM_ARRAY_TASK_ID"
+
+# Program specific arguments
+{conda_code}
+
+"""
+        code += f"""\
+# Program specific arguments
+
+echo "Running scripts in parallel..."
+echo "########################################################################"
+            
+"""
+        for command_line in command_line_list:
+            code += f"""\
+                
+{python_code}
+{seed_code}
+\t\t{command_line}  &
+
+"""
+
+        code += f"""\
+            
+wait # This will wait until both scripts finish
+echo "########################################################################"
+echo "...done."
+"""
+        return code
+
+    def save_slurm(self, command_line_list=None, idx: str = None):
+        code = self.generate_slurm(command_line_list)
+
+        label = f"_{idx}" if idx is not None else ""
+        script_name = f'slurm_{self._exp_name}{label}.sh'
+        full_path = os.path.join(self._exp_dir_slurm_files, script_name)
+
+        with open(full_path, "w") as file:
+            file.write(code)
+
+        return full_path
+
+    def _run_slurm(self, test):
+        # Create slurm directories for sbatch and log files
+        os.makedirs(self._exp_dir_slurm_files)
+        os.makedirs(self._exp_dir_slurm_logs)
+
+        # Generate and save slurm files
+        slurm_files_path_l = []
+        experiment_list_chunked = []
+        for i in range(0, len(self._experiment_list), self._n_exps_in_parallel):
+            experiment_list_chunked.append(self._experiment_list[i:i + self._n_exps_in_parallel])
+
+        for i, exps in enumerate(experiment_list_chunked):
+            command_line_l = []
+            for exp in exps:
+                exp_new_without_underscore = self.remove_last_underscores_dict(exp)
+                command_line_arguments = self._convert_to_command_line(exp_new_without_underscore)
+                results_dir = self._generate_results_dir(self._exp_dir_slurm, exp)
+                command_line_l.append(f'--results_dir {results_dir} {command_line_arguments}')
+            slurm_files_path_l.append(self.save_slurm(command_line_l,
+                                                      str(i).zfill(len(str(len(experiment_list_chunked))))))
+
+        # Launch slurm files in parallel
+        for slurm_file_path in slurm_files_path_l:
+            command = f"sbatch {slurm_file_path}"
+            if test:
+                print(command)
+            else:
+                os.system(command)
+
+    def _run_local(self, test):
+        if not test:
+            os.makedirs(self._exp_dir_local, exist_ok=True)
+
+        module = import_module(self._exp_file)
+        experiment = module.experiment
+
+        if test:
+            self._test_experiment_local()
+        else:
+            def experiment_wrapper(params):
+                try:
+                    experiment(**params)
+                except Exception:
+                    print("Experiment failed with parameters:")
+                    print(params)
+                    traceback.print_exc()
+
+            params_dict = get_experiment_default_params(experiment)
+
+            Parallel(n_jobs=self._n_exps_in_parallel)(delayed(experiment_wrapper)(deepcopy(params))
+                                                      for params in self._generate_exp_params(params_dict))
+
+    def _run_sequential(self, test):
+        if not test:
+            os.makedirs(self._exp_dir_local, exist_ok=True)
+
+        module = import_module(self._exp_file)
+        experiment = module.experiment
+
+        if test:
+            self._test_experiment_local()
+        else:
+            default_params_dict = get_experiment_default_params(experiment)
+
+            for params in self._generate_exp_params(default_params_dict):
+                try:
+                    experiment(**params)
+                except Exception:
+                    print("Experiment failed with parameters:")
+                    print(params)
+                    traceback.print_exc()
+
+    def _check_experiments_results_directories(self):
+        """
+        Check if the results directory produced for each experiment clash.
+        """
+        for exp in self._experiment_list:
+            results_dir = self._generate_results_dir(self._exp_dir_local, exp)
+            # Check if the results directory already exists.
+            if results_dir in self._results_dir_l:
+                # Terminate to prevent overwriting the results directory.
+                raise ResultsDirException(exp, results_dir)
+            self._results_dir_l.append(results_dir)
+
+    def _test_experiment_local(self):
+        for exp, results_dir in zip(self._experiment_list, self._results_dir_l):
+            for i in range(self._start_seed, self._n_seeds):
+                params = str(exp).replace('{', '(').replace('}', '').replace(': ', '=').replace('\'', '')
+                if params:
+                    params += ', '
+                print('experiment' + params + 'seed=' + str(i) + ', results_dir=' + results_dir + ')')
+
+    def _generate_results_dir(self, results_dir, exp, n=6):
+        for key, value in exp.items():
+            if key.endswith('__'):
+                if self._compact_dirs:
+                    subfolder = str(value)
+                else:
+                    subfolder = key + '_' + str(value).replace(' ', '')
+                subfolder = subfolder.replace('/', '-')  # avoids creating subfolders if there is a slash in the name
+                results_dir = os.path.join(results_dir, subfolder)
+        return results_dir
+
+    def _generate_exp_params(self, default_params_dict):
+        seeds = np.arange(self._start_seed, self._start_seed + self._n_seeds)
+        for exp in self._experiment_list:
+            params_dict = deepcopy(default_params_dict)
+            exp_new_without_underscore = self.remove_last_underscores_dict(exp)
+            params_dict.update(exp_new_without_underscore)
+            params_dict['results_dir'] = self._generate_results_dir(self._exp_dir_local, exp)
+            for seed in seeds:
+                params_dict['seed'] = int(seed)
+                yield params_dict
+
+    @staticmethod
+    def remove_last_underscores_dict(exp_dict):
+        exp_dict_new = copy(exp_dict)
+        for key, value in exp_dict.items():
+            if key.endswith('__'):
+                exp_dict_new[key[:-2]] = value
+                del exp_dict_new[key]
+        return exp_dict_new
+
+    @staticmethod
+    def _convert_to_command_line(exp):
+        command_line = ''
+        for key, value in exp.items():
+            new_command = '--' + key + ' '
+
+            if isinstance(value, list):
+                new_command += ' '.join(map(str, value)) + ' '
+            else:
+                new_command += str(value) + ' '
+
+            command_line += new_command
+
+        # remove last space
+        command_line = command_line[:-1]
+
+        return command_line
+
+    @staticmethod
+    def _to_duration(days, hours, minutes, seconds):
+        h = "0" + str(hours) if hours < 10 else str(hours)
+        m = "0" + str(minutes) if minutes < 10 else str(minutes)
+        s = "0" + str(seconds) if seconds < 10 else str(seconds)
+
+        return str(days) + '-' + h + ":" + m + ":" + s
+
+    @property
+    def exp_name(self):
+        return self._exp_name
+
+    def log_dir(self, local=True):
+        if local:
+            return self._exp_dir_local
+        else:
+            return self._exp_dir_slurm
+
+    @property
+    def _exp_file_path(self):
+        module = import_module(self._exp_file)
+        return module.__file__
+
+
+def get_experiment_default_params(func):
+    signature = inspect.signature(func)
+    defaults = {}
+    for k, v in signature.parameters.items():
+        if v.default is not inspect.Parameter.empty:
+            defaults[k] = v.default
+    return defaults
+
+
+def translate_experiment_params_to_argparse(parser, func):
+    annotation_to_argparse = {
+        'str': str,
+        'int': int,
+        'float': float,
+        'bool': bool,
+        'list': None,
+        'tuple': None,
+    }
+    arg_experiments = parser.add_argument_group('Experiment')
+    signature = inspect.signature(func)
+    for k, v in signature.parameters.items():
+        if k not in ['seed', 'results_dir']:
+            if v.default is not inspect.Parameter.empty:
+                if v.annotation.__name__ in annotation_to_argparse:
+                    if v.annotation.__name__ == 'bool':
+                        arg_experiments.add_argument(f"--{str(k)}", type=lambda x: bool(strtobool(x)),
+                                                     nargs='?', const=v.default, default=v.default)
+                    elif v.annotation.__name__ == 'list':
+                        arg_experiments.add_argument(f"--{str(k)}", nargs='+')
+                    else:
+                        arg_experiments.add_argument(f"--{str(k)}", type=annotation_to_argparse[v.annotation.__name__])
+                else:
+                    raise NotImplementedError(f'{v.annotation.__name__} not found in annotation_to_argparse.')
+    return parser
+
+
+def add_launcher_base_args(parser):
+    arg_default = parser.add_argument_group('Default')
+    arg_default.add_argument('--seed', type=int)
+    arg_default.add_argument('--results_dir', type=str)
+    return parser
+
+
+def has_kwargs(func):
+    signature = inspect.signature(func)
+    for k, v in signature.parameters.items():
+        if v.kind == v.VAR_KEYWORD:
+            return True
+
+    return False
+
+
+def string_to_primitive(string):
+    try:
+        return int(string)
+    except ValueError:
+        try:
+            return float(string)
+        except ValueError:
+            try:
+                # boolean
+                return bool(strtobool(string))
+            except ValueError:
+                return string
+
+
+def parse_unknown_args(unknown):
+    kwargs = dict()
+
+    key_idxs = [i for i, arg in enumerate(unknown) if arg.startswith('--')]
+
+    if len(key_idxs) > 0:
+        key_n_args = [key_idxs[i+1] - 1 - key_idxs[i] for i in range(len(key_idxs)-1)]
+        key_n_args.append(len(unknown) - 1 - key_idxs[-1])
+
+        for i, idx in enumerate(key_idxs):
+            key = unknown[idx][2:]
+            n_args = key_n_args[i]
+            if n_args > 1:
+                values = list()
+
+                for v in unknown[idx+1:idx + 1 + n_args]:
+                    values.append(string_to_primitive(v))
+
+                kwargs[key] = values
+
+            elif n_args == 1:
+                kwargs[key] = string_to_primitive(unknown[idx+1])
+
+    return kwargs
+
+
+def parse_args(func):
+    parser = argparse.ArgumentParser()
+
+    parser = translate_experiment_params_to_argparse(parser, func)
+
+    parser = add_launcher_base_args(parser)
+    parser.set_defaults(**get_experiment_default_params(func))
+
+    if has_kwargs(func):
+        args, unknown = parser.parse_known_args()
+        kwargs = parse_unknown_args(unknown)
+
+        args = vars(args)
+        args.update(kwargs)
+
+        return args
+    else:
+        args = parser.parse_args()
+        return vars(args)
+
+
+def run_experiment(func, args=None):
+    if not args:
+        args = parse_args(func)
+
+    func(**args)

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/experiment_launcher/utils.py

@@ -0,0 +1,117 @@
+import datetime
+import json
+import os
+import random
+from random import randint
+from shutil import which
+
+import git
+import wandb
+import yaml
+from git import InvalidGitRepositoryError
+
+
+def save_args(results_dir, args, git_repo_path=None, seed=None, save_args_as_yaml=False, print_exp_args=False):
+    try:
+        repo = git.Repo(git_repo_path, search_parent_directories=True)
+        args['git_hash'] = repo.head.object.hexsha
+        args['git_url'] = repo.remotes.origin.url
+    except InvalidGitRepositoryError:
+        args['git_hash'] = ''
+        args['git_url'] = ''
+    except ValueError:
+        args['git_hash'] = ''
+        args['git_url'] = ''
+
+    if save_args_as_yaml:
+        filename = 'args.yaml' if seed is None else f'args-{seed}.yaml'
+        with open(os.path.join(results_dir, filename), 'w') as f:
+            yaml.dump(args, f, Dumper=yaml.Dumper)
+    else:
+        filename = 'args.json' if seed is None else f'args-{seed}.json'
+        with open(os.path.join(results_dir, filename), 'w') as f:
+            json.dump(args, f, indent=2)
+
+    del args['git_hash']
+    del args['git_url']
+
+    if print_exp_args:
+        print('------------------------------------------------------------------------------------')
+        print('--------> Experiment arguments')
+        print(json.dumps(args, indent=2))
+        print('------------------------------------------------------------------------------------')
+
+
+def is_local():
+    return which('sbatch') is None
+
+
+def start_wandb(
+        wandb_mode='disabled',
+        wandb_entity='experiment_launcher',
+        wandb_project='test_experiment_launcher',
+        wandb_group=None,
+        wandb_run_name=None,
+        **kwargs
+):
+    # https://github.com/wandb/wandb/issues/3911#issuecomment-1409769887
+    # workaround to make sure that wandb does not crash
+    os.environ["WANDB__SERVICE_WAIT"] = "600"
+
+    modes = ['online', 'offline', 'disabled']
+    assert wandb_mode in modes, f"wandb_mode must be in {modes}"
+
+    if wandb_mode == 'disabled':
+        return wandb.init(mode="disabled", reinit=True)
+
+    init = {
+        "mode": wandb_mode,
+        "entity": wandb_entity,
+        "project": wandb_project,
+        "group": wandb_group,
+        "name": wandb_run_name,
+        "reinit": True,
+        "notes": datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S'),
+    }
+
+    return wandb.init(**init)
+
+
+def fix_random_seed(seed):
+    random.seed(seed)
+
+    try:
+        import numpy as np
+        import torch
+    except ImportError:
+        pass
+
+    try:
+        np.random.seed(seed)
+    except NameError:
+        pass
+
+    try:
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        # https://pytorch.org/docs/stable/notes/randomness.html#cuda-convolution-benchmarking
+        # torch.backends.cudnn.benchmark = False
+    except NameError:
+        pass
+
+
+def random_int_with_n_digits(n):
+    range_start = 10 ** (n - 1)
+    range_end = (10 ** n) - 1
+    return randint(range_start, range_end)
+
+
+def create_results_dir(kwargs, make_dirs_with_seed=True):
+    seed = kwargs['seed']
+    results_dir = kwargs['results_dir']
+    assert results_dir is not None and seed is not None, "results_dir and seed must exist"
+    results_dir_new = results_dir
+    if make_dirs_with_seed:
+        results_dir_new = os.path.join(results_dir_new, str(seed))
+    os.makedirs(results_dir_new, exist_ok=True)
+    kwargs['results_dir'] = results_dir_new

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/requirements.txt

@@ -0,0 +1,5 @@
+joblib
+gitpython
+numpy
+pyyaml
+wandb

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/experiment_launcher/setup.py

@@ -0,0 +1,35 @@
+from setuptools import setup, find_packages
+from codecs import open
+from os import path
+
+from experiment_launcher import __version__
+
+here = path.abspath(path.dirname(__file__))
+
+requires_list = []
+with open(path.join(here, 'requirements.txt'), encoding='utf-8') as f:
+    for line in f:
+        requires_list.append(str(line))
+
+long_description = 'experiment_launcher is a Python library used to run experiments.' \
+                   ' Supports Joblib and slurm jobs.'
+
+setup(
+    name='experiment-launcher',
+    version=__version__,
+    description='A Python toolkit for launching experiments.',
+    long_description=long_description,
+    url='https://git.ias.informatik.tu-darmstadt.de/common/experiment_launcher',
+    author="Davide Tateo, Joao Carvalho",
+    author_email='davide@robot-learning.de, joao@robot-learning.de',
+    license='MIT',
+    packages=[package for package in find_packages()
+              if package.startswith('experiment_launcher')],
+    zip_safe=False,
+    install_requires=requires_list,
+    extras_require={},
+    classifiers=["Programming Language :: Python :: 3",
+                 "License :: OSI Approved :: MIT License",
+                 "Operating System :: OS Independent",
+                 ]
+)

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/motion_planning_baselines/LICENSE

@@ -0,0 +1,24 @@
+MIT License
+
+Copyright (c) 2024 Yorai Shaoul, slightly modified original code written by An Thai Le, Joao Carvalho.
+Original license below:
+
+Copyright (c) 2023 An Thai Le, Joao Carvalho
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/motion_planning_baselines/README.md

@@ -0,0 +1,54 @@
+The original repository is this one: https://github.com/anindex/motion_planning_baselines.
+Their original README follows.
+
+# Motion Planning Baselines in PyTorch
+
+This library implements various Motion Planning methods.
+
+**NOTE**: `mp_baslines` is under heavy development and highly experimental.
+
+## Installation
+
+Simply activate your conda/Python environment and run
+
+```bash
+pip install -e .
+```
+
+## Examples
+
+```bash
+python examples/pointmass_dense_2d_CHOMP.py
+python examples/pointmass_dense_2d_GPMP.py
+python examples/pointmass_grid_circles_2d_Stoch-GPMP.py
+python examples/pointmass_grid_circles_2d_STOMP.py
+python examples/pointmass_grid_circles_2d_MPPI.py
+python examples/pointmass_dense_2d_RRT_multiprocess.py
+```
+
+## Contact
+
+If you have any questions or find any bugs, please let us know:
+
+- [An Le](https://www.ias.informatik.tu-darmstadt.de/Team/AnThaiLe), [an@robot-learning.de](an@robot-learning.de)
+- [Joao Carvalho](https://www.ias.informatik.tu-darmstadt.de/Team/JoaoCarvalho), [joao@robot-learning.de](joao@robot-learning.de)
+
+## Citation
+
+If you found this repository useful, please consider citing these references:
+
+```bibtex
+@inproceedings{le2023accelerating,
+  title={Accelerating Motion Planning via Optimal Transport},
+  author={Le, An T. and Chalvatzaki, Georgia and Biess, Armin and Peters, Jan},
+  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
+  year={2023}
+}
+
+@article{carvalho2023motion,
+  title={Motion planning diffusion: Learning and planning of robot motions with diffusion models},
+  author={Carvalho, Joao and Le, An T and Baierl, Mark and Koert, Dorothea and Peters, Jan},
+  journal={arXiv preprint arXiv:2308.01557},
+  year={2023}
+}
+```

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/motion_planning_baselines/mp_baselines/planners/base.py

@@ -0,0 +1,226 @@
+import abc
+from abc import ABC, abstractmethod
+from typing import Tuple
+import torch
+
+from mp_baselines.planners.costs.factors.gp_factor import GPFactor
+from mp_baselines.planners.costs.factors.mp_priors_multi import MultiMPPrior
+from mp_baselines.planners.costs.factors.unary_factor import UnaryFactor
+from torch_robotics.trajectory.utils import finite_difference_vector
+
+
+class MPPlanner(ABC):
+    """Base class for all planners."""
+
+    def __init__(
+        self,
+        name: str = None,
+        tensor_args: dict = None,
+        **kwargs
+    ):
+        self.name = name
+        self.tensor_args = tensor_args
+        self._kwargs = kwargs
+
+    @abstractmethod
+    def optimize(self, opt_iters: int = 1, **observation) -> Tuple[bool, torch.Tensor]:
+        """Plan a path from start to goal.
+
+        Args:
+            opt_iters: Number of optimization iters.
+            observation: dict of observations.
+
+        Returns:
+            success: True if a path was found.
+            path: Path from start to goal.
+        """
+        pass
+
+    def __call__(self, opt_iters: int = 1, **observation) -> Tuple[bool, torch.Tensor]:
+        """Plan a path from start to goal.
+
+        Args:
+            start: Start position.
+            goal: Goal position.
+
+        Returns:
+            success: True if a path was found.
+            path: Path from start to goal.
+        """
+        return self.optimize(opt_iters, **observation)
+
+    def __repr__(self):
+        return f"{self.name}({self._kwargs})"
+
+    @abc.abstractmethod
+    def render(self, ax, **kwargs):
+        raise NotImplementedError
+
+
+class OptimizationPlanner(MPPlanner):
+
+    def __init__(
+        self,
+        name: str = 'OptimizationPlanner',
+        n_dof: int = None,
+        n_support_points: int = None,
+        n_interpolated_points: int = None,
+        num_particles_per_goal: int = None,
+        opt_iters: int = None,
+        dt: float = None,
+        start_state: torch.Tensor = None,
+        cost=None,
+        initial_particle_means=None,
+        multi_goal_states: torch.Tensor = None,
+        sigma_start_init: float = 0.001,
+        sigma_goal_init: float = 0.001,
+        sigma_gp_init: float = 10.,
+        pos_only: bool = False,
+        tensor_args: dict = None, **kwargs
+    ):
+        super().__init__(name, tensor_args, **kwargs)
+        self.n_dof = n_dof
+        self.dim = 2 * self.n_dof
+        self.n_support_points = n_support_points
+        self.n_interpolated_points = n_interpolated_points
+        self.num_particles_per_goal = num_particles_per_goal
+        self.opt_iters = opt_iters
+        self.dt = dt
+        self.pos_only = pos_only
+
+        self.start_state = start_state
+        self.multi_goal_states = multi_goal_states
+        if multi_goal_states is None:  # NOTE(an): if there is no goal, we assume xere is at least one solution
+            self.num_goals = 1
+        else:
+            assert multi_goal_states.ndim == 2
+            self.num_goals = multi_goal_states.shape[0]
+        self.num_particles = self.num_goals * self.num_particles_per_goal
+        self.cost = cost
+        self.initial_particle_means = initial_particle_means
+        self._particle_means = None
+        if self.pos_only:
+            self.d_state_opt = self.n_dof
+            self.start_state = self.start_state
+        else:
+            self.d_state_opt = 2 * self.n_dof
+            self.start_state = torch.cat([self.start_state, torch.zeros_like(self.start_state)], dim=-1)
+            if self.multi_goal_states is not None:
+                self.multi_goal_states = torch.cat([self.multi_goal_states, torch.zeros_like(self.multi_goal_states)], dim=-1)
+
+        self.sigma_start_init = sigma_start_init
+        self.sigma_goal_init = sigma_goal_init
+        self.sigma_gp_init = sigma_gp_init
+
+    def get_GP_prior(
+            self,
+            start_K,
+            gp_K,
+            goal_K,
+            state_init,
+            particle_means=None,
+            goal_states=None,
+            tensor_args=None,
+    ):
+        if tensor_args is None:
+            tensor_args = self.tensor_args
+        return MultiMPPrior(
+            self.n_support_points - 1,
+            self.dt,
+            self.dim,
+            self.n_dof,
+            start_K,
+            gp_K,
+            state_init,
+            K_g_inv=goal_K,  # NOTE(sasha) Assume same goal Cov. for now
+            means=particle_means,
+            goal_states=goal_states,
+            tensor_args=tensor_args,
+        )
+
+    def const_vel_trajectories(
+        self,
+        start_state,
+        multi_goal_states,
+    ):
+        traj_dim = (multi_goal_states.shape[0], self.n_support_points, self.dim)
+        state_traj = torch.zeros(traj_dim, **self.tensor_args)
+        mean_vel = (multi_goal_states[:, :self.n_dof] - start_state[:, :self.n_dof]) / (self.n_support_points * self.dt)
+        for i in range(self.n_support_points):
+            state_traj[:, i, :self.n_dof] = start_state[:, :self.n_dof] * (self.n_support_points - i - 1) / (self.n_support_points - 1) \
+                                  + multi_goal_states[:, :self.n_dof] * i / (self.n_support_points - 1)
+        state_traj[:, :, self.n_dof:] = mean_vel.unsqueeze(1).repeat(1, self.n_support_points, 1)
+        return state_traj
+
+    def get_random_trajs(self):
+        # force torch.float64
+        tensor_args = dict(dtype=torch.float64, device=self.tensor_args['device'])
+        # set zero velocity for GP prior
+        start_state = torch.cat((self.start_state, torch.zeros_like(self.start_state)), dim=-1).to(**tensor_args)
+        if self.multi_goal_states is not None:
+            multi_goal_states = torch.cat((self.multi_goal_states, torch.zeros_like(self.multi_goal_states)), dim=-1).to(**tensor_args)
+        else:
+            multi_goal_states = None
+        #========= Initialization factors ===============
+        self.start_prior_init = UnaryFactor(
+            self.n_dof * 2,
+            self.sigma_start_init,
+            start_state,
+            tensor_args,
+        )
+
+        self.gp_prior_init = GPFactor(
+            self.n_dof,
+            self.sigma_gp_init,
+            self.dt,
+            self.n_support_points - 1,
+            tensor_args,
+        )
+
+        self.multi_goal_prior_init = []
+        if multi_goal_states is not None:
+            for i in range(self.num_goals):
+                self.multi_goal_prior_init.append(
+                    UnaryFactor(
+                        self.n_dof * 2,
+                        self.sigma_goal_init,
+                        multi_goal_states[i],
+                        tensor_args,
+                    )
+                )
+        self._traj_dist = self.get_GP_prior(
+                self.start_prior_init.K,
+                self.gp_prior_init.Q_inv[0],
+                self.multi_goal_prior_init[0].K if multi_goal_states is not None else None,
+                start_state,
+                goal_states=multi_goal_states,
+                tensor_args=tensor_args,
+            )
+        particles = self._traj_dist.sample(self.num_particles_per_goal).to(**tensor_args)
+        self.traj_dim = particles.shape
+        del self._traj_dist  # free memory
+        return particles.flatten(0, 1).to(**self.tensor_args)
+
+    def _get_traj(self):
+        """
+            Get position-velocity trajectory from control distribution.
+        """
+        trajs = self._particle_means.clone()
+        if self.pos_only:
+            # Linear velocity by central finite differencing
+            vels = finite_difference_vector(trajs, dt=self.dt)
+            trajs = torch.cat((trajs, vels), dim=1)
+        return trajs
+
+    def get_traj(self):
+        return self._get_traj()
+
+    def _get_costs(self, state_trajectories, **observation):
+        if self.cost is None:
+            costs = torch.zeros(self.num_particles, )
+        else:
+            costs = self.cost(state_trajectories, **observation)
+        return costs
+
+    def render(self, ax, **kwargs):
+        raise NotImplementedError

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/motion_planning_baselines/mp_baselines/planners/chomp.py

@@ -0,0 +1,169 @@
+import torch
+
+from mp_baselines.planners.base import OptimizationPlanner
+from torch_robotics.torch_utils.torch_timer import TimerCUDA
+from torch_robotics.torch_utils.torch_utils import batched_weighted_dot_prod
+
+
+class CHOMP(OptimizationPlanner):
+    # https://www.ri.cmu.edu/pub_files/2009/5/icra09-chomp.pdf
+
+    def __init__(
+            self,
+            n_dof: int,
+            n_support_points: int,
+            num_particles_per_goal: int,
+            opt_iters: int,
+            dt: float,
+            start_state: torch.Tensor,
+            cost=None,
+            weight_prior_cost=0.1,
+            initial_particle_means=None,
+            step_size: float = 1.,
+            grad_clip: float = .01,
+            multi_goal_states: torch.Tensor = None,
+            sigma_start_init=0.001,
+            sigma_goal_init=0.001,
+            sigma_gp_init=10.,
+            pos_only: bool = True,
+            **kwargs
+    ):
+        super(CHOMP, self).__init__(name='CHOMP',
+                                    n_dof=n_dof,
+                                    n_support_points=n_support_points,
+                                    num_particles_per_goal=num_particles_per_goal,
+                                    opt_iters=opt_iters,
+                                    dt=dt,
+                                    start_state=start_state,
+                                    cost=cost,
+                                    initial_particle_means=initial_particle_means,
+                                    multi_goal_states=multi_goal_states,
+                                    sigma_start_init=sigma_start_init,
+                                    sigma_goal_init=sigma_goal_init,
+                                    sigma_gp_init=sigma_gp_init,
+                                    pos_only=pos_only,
+                                    **kwargs)
+
+        # CHOMP params
+        self.lr = step_size
+        self.grad_clip = grad_clip
+
+        self._particle_means = None
+        # Precision matrix, shape: [ctrl_dim, n_support_points, n_support_points]
+        self.Sigma_inv = self._get_R_mat(dt=self.dt, n_support_points=self.n_support_points, tensor_args=self.tensor_args)
+        self.Sigma = torch.inverse(self.Sigma_inv)
+        self.reset(initial_particle_means=initial_particle_means)
+
+        # Weight prior
+        self.weight_prior_cost = weight_prior_cost
+
+    def _get_R_mat2(self):
+        """
+        STOMP time-correlated Precision matrix.
+        Central finite difference velocity.
+        """
+        upper_diag = torch.diag(torch.ones(self.n_support_points - 1), diagonal=1)
+        lower_diag = torch.diag(torch.ones(self.n_support_points - 1), diagonal=-1,)
+        diag = -2 * torch.eye(self.n_support_points)
+        A_mat = upper_diag + diag + lower_diag
+        A_mat = torch.cat(
+            (torch.zeros(1, self.n_support_points),
+             A_mat,
+             torch.zeros(1, self.n_support_points)),
+            dim=0,
+        )
+        A_mat[0, 0] = 1.
+        A_mat[-1, -1] = 1.
+        A_mat = A_mat * 1./self.dt**2
+        R_mat = A_mat.t() @ A_mat
+        return R_mat.to(**self.tensor_args)
+
+    @classmethod
+    def _get_R_mat(
+            cls,
+            dt=0.01,
+            n_support_points=64,
+            tensor_args=None,
+            **kwargs
+    ):
+        """
+        CHOMP time-correlated Precision matrix.
+        Backward finite difference velocity.
+        """
+        lower_diag = -torch.diag(torch.ones(n_support_points - 1), diagonal=-1)
+        diag = 1 * torch.eye(n_support_points)
+        K_mat = diag + lower_diag
+        K_mat = torch.cat((K_mat, torch.zeros(1, n_support_points)), dim=0)
+        K_mat[-1, -1] = -1.
+        K_mat = K_mat * 1. / dt ** 2
+        R_mat = K_mat.t() @ K_mat
+
+        return R_mat.to(**tensor_args)
+
+    def reset(
+            self,
+            initial_particle_means=None,
+    ):
+        # Straight line position-trajectory from start to goal with const vel
+        if initial_particle_means is not None:
+            self._particle_means = initial_particle_means.clone()
+        else:
+            self._particle_means = self.get_random_trajs()
+
+    def optimize(
+            self,
+            opt_iters=None,
+            **observation
+    ):
+        """
+        Optimize for best trajectory at current state
+        """
+
+        self._run_optimization(opt_iters, **observation)
+        # get mean trajectory
+        curr_traj = self._get_traj()
+        return curr_traj
+
+    def _run_optimization(self, opt_iters, **observation):
+        """
+        Run optimization iterations.
+        """
+        optim_vis = observation.get('optim_vis', False)
+        if opt_iters is None:
+            opt_iters = self.opt_iters
+        for opt_step in range(opt_iters):
+            self._particle_means.requires_grad_(True)
+            costs = self._eval(self._particle_means, **observation)
+
+            # Get grad
+            costs.sum().backward(retain_graph=True)
+            # For stabilizing and preventing high gradients
+            self._particle_means.grad.data.clamp_(-self.grad_clip, self.grad_clip)
+            # zeroing grad at start and goal
+            self._particle_means.grad.data[..., 0, :] = 0.
+            self._particle_means.grad.data[..., -1, :] = 0.
+
+            # Update trajectory
+            self._particle_means.data.add_(-self.lr * self._particle_means.grad.data)
+            self._particle_means.grad.detach_()
+            self._particle_means.grad.zero_()
+
+        self._particle_means = self._particle_means.detach()
+
+    def _eval(self, x, **observation):
+        """
+        Evaluate costs.
+        """
+        if x.ndim == 2:
+            x = x.unsqueeze(0)
+
+        # Evaluate (collision, ...) costs
+        costs = self._get_costs(x, **observation)
+
+        # Add smoothness term (prior)
+        R_mat = self.Sigma_inv
+        smooth_cost = batched_weighted_dot_prod(x, R_mat, x).sum()
+
+        costs += self.weight_prior_cost * smooth_cost
+
+        return costs

applications/multi_robot/robotarium/robotarium_baselines/DiRecT/deps/motion_planning_baselines/mp_baselines/planners/costs/cost_functions.py

@@ -0,0 +1,745 @@
+from abc import ABC, abstractmethod
+
+import einops
+import numpy as np
+import torch
+from typing import List, Optional
+
+from mp_baselines.planners.chomp import CHOMP
+from mp_baselines.planners.costs.factors.field_factor import FieldFactor
+from mp_baselines.planners.costs.factors.gp_factor import GPFactor
+from mp_baselines.planners.costs.factors.unary_factor import UnaryFactor
+from torch_robotics.torch_kinematics_tree.geometrics.utils import link_pos_from_link_tensor
+from torch_robotics.torch_planning_objectives.fields.distance_fields import interpolate_points_v1
+from torch_robotics.torch_utils.torch_utils import batched_weighted_dot_prod
+from torch_robotics.trajectory.utils import finite_difference_vector
+
+
+class Cost(ABC):
+    def __init__(self, robot, n_support_points, tensor_args=None, **kwargs):
+        self.robot = robot
+        self.n_dof = robot.q_dim
+        self.dim = 2 * self.n_dof  # position + velocity
+        self.n_support_points = n_support_points
+
+        self.tensor_args = tensor_args
+
+    def set_cost_factors(self):
+        pass
+
+    def __call__(self, trajs, **kwargs):
+        return self.eval(trajs, **kwargs)
+
+    @abstractmethod
+    def eval(self, trajs, **kwargs):
+        pass
+
+    @abstractmethod
+    def get_linear_system(self, trajs, **kwargs):
+        pass
+
+    def get_q_pos_vel_and_fk_map(self, trajs, **kwargs):
+        assert trajs.ndim == 3 or trajs.ndim == 4
+        N = 1
+        if trajs.ndim == 4:
+            N, B, H, D = trajs.shape  # n_goals (or steps), batch of trajectories, length, dim
+            trajs = einops.rearrange(trajs, 'N B H D -> (N B) H D')
+        else:
+            B, H, D = trajs.shape
+
+        q_pos = self.robot.get_position(trajs)
+        q_vel = self.robot.get_velocity(trajs)
+        H_positions = self.robot.fk_map_collision(
+            q_pos)  # I, taskspaces, x_dim+1, x_dim+1 (homogeneous transformation matrices)
+        return trajs, q_pos, q_vel, H_positions
+
+
+class CostComposite(Cost):
+
+    def __init__(
+            self,
+            robot,
+            n_support_points,
+            cost_list,
+            weights_cost_l=None,
+            **kwargs
+    ):
+        super().__init__(robot, n_support_points, **kwargs)
+        self.cost_l = cost_list
+        self.weight_cost_l = weights_cost_l if weights_cost_l is not None else [1.0] * len(cost_list)
+
+    def eval(self, trajs, trajs_interpolated=None, return_invidual_costs_and_weights=False, **kwargs):
+        trajs, q_pos, q_vel, H_positions = self.get_q_pos_vel_and_fk_map(trajs)
+
+        if not return_invidual_costs_and_weights:
+            cost_total = 0
+            for cost, weight_cost in zip(self.cost_l, self.weight_cost_l):
+                if trajs_interpolated is not None:
+                    # Compute only collision costs with interpolated trajectories.
+                    # Other costs are computed with non-interpolated trajectories, e.g. smoothness
+                    if isinstance(cost, CostCollision):
+                        trajs_tmp = trajs_interpolated
+                    else:
+                        trajs_tmp = trajs
+                else:
+                    trajs_tmp = trajs
+                cost_tmp = weight_cost * cost(trajs_tmp, q_pos=q_pos, q_vel=q_vel, H_positions=H_positions, **kwargs)
+                cost_total += cost_tmp
+            return cost_total
+        else:
+            cost_l = []
+            for cost in self.cost_l:
+                if trajs_interpolated is not None:
+                    # Compute only collision costs with interpolated trajectories.
+                    # Other costs are computed with non-interpolated trajectories, e.g. smoothness
+                    if isinstance(cost, CostCollision):
+                        trajs_tmp = trajs_interpolated
+                    else:
+                        trajs_tmp = trajs
+                else:
+                    trajs_tmp = trajs
+
+                cost_tmp = cost(trajs_tmp, q_pos=q_pos, q_vel=q_vel, H_positions=H_positions, **kwargs)
+                cost_l.append(cost_tmp)
+
+            if return_invidual_costs_and_weights:
+                return cost_l, self.weight_cost_l
+
+    def get_linear_system(self, trajs, n_interpolated_points=None, **kwargs):
+        trajs.requires_grad = True
+        # TODO - join fk map into one call for trajs and trajs_interp
+        trajs, q_pos, q_vel, H_positions = self.get_q_pos_vel_and_fk_map(trajs)
+
+        # Upsample trajectory for finer collision checking
+        # Interpolate in joint space
+        # TODO - change from linear interpolation to GP interpolation
+        if n_interpolated_points is None:
+            trajs_interp, q_pos_interp, q_vel_interp, H_positions_interp = None, None, None, None
+        else:
+            trajs_interp = interpolate_points_v1(trajs, n_interpolated_points)
+            trajs_interp, q_pos_interp, q_vel_interp, H_positions_interp = self.get_q_pos_vel_and_fk_map(trajs_interp)
+
+        batch_size = trajs.shape[0]
+        As, bs, Ks = [], [], []
+        optim_dim = 0
+        for cost, weight_cost in zip(self.cost_l, self.weight_cost_l):
+            A, b, K = cost.get_linear_system(
+                trajs, q_pos=q_pos, q_vel=q_vel, H_positions=H_positions,
+                trajs_interp=trajs_interp, q_pos_interp=q_pos_interp, q_vel_interp=q_vel_interp,
+                H_positions_interp=H_positions_interp,
+                **kwargs)
+            if A is None or b is None or K is None:
+                continue
+            optim_dim += A.shape[1]
+            As.append(A.detach())
+            bs.append(b.detach())
+            Ks.append(K.detach())
+
+        A = torch.cat(As, dim=1)
+        b = torch.cat(bs, dim=1)
+        K = torch.zeros(batch_size, optim_dim, optim_dim, **self.tensor_args)
+        offset = 0
+        for i in range(len(Ks)):
+            dim = Ks[i].shape[1]
+            K[:, offset:offset + dim, offset:offset + dim] = Ks[i]
+            offset += dim
+        return A, b, K
+
+
+class CostCollision(Cost):
+
+    def __init__(
+            self,
+            robot,
+            n_support_points,
+            field=None,
+            sigma_coll=None,
+            **kwargs
+    ):
+        super().__init__(robot, n_support_points, **kwargs)
+        self.obst_factor = None
+        self.field = field
+        self.sigma_coll = sigma_coll
+
+        self.set_cost_factors()
+
+    def set_cost_factors(self):
+        # ========= Cost factors ===============
+        self.obst_factor = FieldFactor(
+            self.n_dof,
+            self.sigma_coll,
+            # [45, 50]  # Can consider also only a subset of the trajectory.
+            [1, None]  # Take the whole trajectory except for the first point.
+        )
+
+    def eval(self, trajs, q_pos=None, q_vel=None, H_positions=None, **observation):
+        costs = 0
+        if self.field is not None:
+            # H_pos = link_pos_from_link_tensor(H)  # get translation part from transformation matrices
+            H_pos = H_positions
+            err_obst = self.obst_factor.get_error(
+                trajs,
+                self.field,
+                q_pos=q_pos,
+                q_vel=q_vel,
+                H_pos=H_pos,
+                calc_jacobian=False,  # TODO: NOTE(an): no need for grads in StochGPMP
+                obstacle_spheres=observation.get('obstacle_spheres', None)
+            )
+            w_mat = self.obst_factor.K
+            obst_costs = w_mat * err_obst.sum(1)
+            costs = obst_costs
+
+        return costs
+
+    def get_linear_system(self, trajs, q_pos=None, q_vel=None, H_positions=None,
+                          trajs_interp=None, q_pos_interp=None, q_vel_interp=None, H_positions_interp=None,
+                          **observation):
+        A, b, K = None, None, None
+        if self.field is not None:
+            batch_size = trajs.shape[0]
+            # H_pos = link_pos_from_link_tensor(H)  # get translation part from transformation matrices
+            # H_pos = link_pos_from_link_tensor(H)  # get translation part from transformation matrices
+            H_pos = H_positions
+
+            # Get H_obst wrt to the interpolated trajectory. This is computed inside get_error
+
+            err_obst, H_obst = self.obst_factor.get_error(
+                trajs,
+                self.field,
+                q_pos=q_pos,
+                q_vel=q_vel,
+                H_pos=H_pos,
+                trajs_interp=trajs_interp,
+                q_pos_interp=q_pos_interp,
+                q_vel_interp=q_vel_interp,
+                H_pos_interp=H_positions_interp,
+                calc_jacobian=True,
+                obstacle_spheres=observation.get('obstacle_spheres', None)
+            )
+
+            A = torch.zeros(batch_size, self.n_support_points - 1, self.dim * self.n_support_points, **self.tensor_args)
+            A[:, :, :H_obst.shape[-1]] = H_obst
+            # shift each row by self.dim
+            idxs = torch.arange(A.shape[-1], **self.tensor_args).repeat(A.shape[-2], 1)
+            idxs = (idxs - torch.arange(self.dim, (idxs.shape[0] + 1) * self.dim, self.dim, **self.tensor_args).view(-1, 1)) % idxs.shape[-1]
+            idxs = idxs.to(torch.int64)
+            A = torch.gather(A, -1, idxs.repeat(batch_size, 1, 1))
+
+            # old code not vectorized
+            # https://github.com/anindex/stoch_gpmp/blob/main/stoch_gpmp/costs/cost_functions.py#L275
+
+            b = err_obst.unsqueeze(-1)
+            K = self.obst_factor.K * torch.eye((self.n_support_points - 1), **self.tensor_args).repeat(batch_size, 1, 1)
+
+        return A, b, K
+
+
+# class CostConstraint(Cost):
+#     """
+#     Currently implementing a "vertex constraint."
+#     Avoid this configuration at this time (interval).
+#     """
+#     def __init__(
+#             self,
+#             robot,
+#             n_support_points,
+#             q: torch.Tensor,  # Shape: (n_dof,)
+#             traj_range: tuple,
+#             radius: float,
+#             sigma=None,
+#             **kwargs
+#     ):
+#         super().__init__(robot, n_support_points, **kwargs)
+#         self.q = q.to(**self.tensor_args)
+#         self.traj_range = traj_range
+#         self.radius = radius
+#         self.sigma = sigma
+#
+#     def eval(self, trajs, **observation):
+#         costs = 0
+#         if self.q is not None:
+#             q_pos = self.robot.get_position(trajs)
+#             q_pos = q_pos[:, self.traj_range[0]:self.traj_range[1], :]
+#             dist_constraint = torch.norm(q_pos - self.q, dim=-1)
+#             mask = dist_constraint > self.radius
+#             dist_constraint = torch.where(mask, torch.zeros_like(dist_constraint), dist_constraint)
+#             costs = self.radius - dist_constraint
+#             costs = costs.sum(-1)
+#         return costs
+#
+#     def get_linear_system(self, trajs, **kwargs):
+#         pass
+
+
+class CostConstraint(Cost):
+    """
+    Implements multiple "vertex constraints" in a fully vectorized manner.
+    Each constraint is associated with a specific configuration and time interval.
+    """
+
+    def __init__(
+            self,
+            robot,
+            n_support_points,
+            q_l: List[torch.Tensor],  # Length n.
+            traj_range_l: List[tuple],  # Length n. A tuple (start, end) for each constraint. Inclusive.
+            radius_l: List[float],  # Length n. The radius for each constraint.
+            is_soft: bool = False,  # Sigma stands for the standard deviation of the cost.
+            **kwargs
+    ):
+        super().__init__(robot, n_support_points, **kwargs)
+        self.qs = torch.stack([q.to(**self.tensor_args) for q in q_l], dim=0)  # (n, q_dim)
+        self.traj_ranges = torch.tensor(traj_range_l, **self.tensor_args)  # (n, 2)
+        self.radii = torch.tensor(radius_l, device=self.qs.device)  # (n,)
+        self.is_soft = is_soft
+
+    def eval(self, trajs, **observation):
+        q_pos = self.robot.get_position(trajs)  # (B, H, q_dim)
+        # Get the start and end indices for each range.
+        start_indices = self.traj_ranges[:, 0]  # (n,)
+        end_indices = self.traj_ranges[:, 1]  # (n,)
+
+        # Prepare a mask to extract relevant positions in the trajectory for each constraint.
+        mask = torch.arange(q_pos.shape[1], device=q_pos.device).unsqueeze(0).unsqueeze(0)  # (1, 1, H)
+        mask = (mask >= start_indices.view(-1, 1, 1)) & (mask < end_indices.view(-1, 1, 1))  # (n, 1, H)
+
+        # Repeat and reshape q_pos to match the constraints dimensions.
+        q_pos = q_pos.unsqueeze(0).expand(self.qs.shape[0], -1, -1, -1)  # (n, B, H, q_dim)
+        q_pos_masked = q_pos * mask.unsqueeze(-1)  # (n, B, H, q_dim)
+
+        # Compute the distance between each position in the range and the corresponding q.
+        dist_constraint = torch.norm(q_pos_masked - self.qs.view(-1, 1, 1, q_pos.shape[-1]),
+                                     dim=-1)  # (n, B, H)
+
+        # Apply the radius constraints.
+        mask_radius = dist_constraint > self.radii.view(-1, 1, 1)  # (n, B, H)
+        dist_constraint = torch.where(mask_radius, torch.zeros_like(dist_constraint),
+                                      dist_constraint)  # Only penalize inside the radius
+
+        # Compute the cost.
+        costs = (self.radii.view(-1, 1, 1) - dist_constraint).sum(dim=-1)  # (n, B)
+
+        # Sum the costs across all constraints.
+        total_costs = costs.sum()  # (1,)
+
+        return total_costs
+
+    def get_linear_system(self, trajs, **kwargs):
+        pass
+
+
+class CostMaxVelocity(Cost):
+    def __init__(
+            self,
+            robot,
+            n_support_points,
+            max_vel=None,
+            sigma_max_vel=None,
+            **kwargs
+    ):
+        super().__init__(robot, n_support_points, **kwargs)
+        self.max_vel = max_vel
+        self.sigma_max_vel = sigma_max_vel
+
+    def eval(self, trajs, **observation):
+        costs = 0
+        q_pos_trajs = self.robot.get_position(trajs)
+        if self.max_vel is not None:
+            vel = finite_difference_vector(q_pos_trajs, dt=self.robot.dt, method='central')
+            cost = torch.pow(vel - self.max_vel, 2)
+            costs = cost.sum(-1)
+        return costs
+
+    def get_linear_system(self, trajs, **observation):
+        pass
+
+
+class CostVelocityAndDirectionAlignment(Cost):
+    def __init__(
+            self,
+            robot,
+            n_support_points,
+            **kwargs
+    ):
+        super().__init__(robot, n_support_points, **kwargs)
+
+    def eval(self, trajs, **observation):
+        q_direction_trajs = self.robot.get_direction_vector(trajs)
+        q_spatial = self.robot.get_spatial_component(trajs)
+        q_spatial_vel = finite_difference_vector(q_spatial, dt=self.robot.dt, method='central')
+        q_spatial_vel = q_spatial_vel / torch.norm(q_spatial_vel, dim=-1, keepdim=True)
+        cost = torch.sum(q_direction_trajs * q_spatial_vel, dim=-1)  # Just dot.
+        costs = cost.sum(-1)
+        return costs
+
+    def get_linear_system(self, trajs, **observation):
+        batch_size = trajs.shape[0]
+        dim = trajs.shape[-1]
+
+        # Initialize the linear system components A, b, and K.
+        A = torch.zeros(batch_size, dim * self.n_support_points, dim * self.n_support_points, **self.tensor_args)
+        b = torch.zeros(batch_size, dim * self.n_support_points, 1, **self.tensor_args)
+        K = torch.zeros(batch_size, dim * self.n_support_points, dim * self.n_support_points, **self.tensor_args)
+
+        # Get direction and velocity trajectories
+        q_direction_trajs = self.robot.get_direction_vector(trajs)
+        q_spatial = self.robot.get_spatial_component(trajs)
+        q_spatial_vel = finite_difference_vector(q_spatial, dt=self.robot.dt, method='central')
+        q_spatial_vel = q_spatial_vel / torch.norm(q_spatial_vel, dim=-1, keepdim=True)
+
+        # Calculate the error as the difference between direction and normalized velocity vectors
+        error = q_direction_trajs - q_spatial_vel
+
+        # Now, populate A, b, and K matrices.
+        # A will represent the linearized relationship between the deviation and the state
+        # b will represent the error between the direction and velocity
+        # K will represent a weighting matrix for this cost (often identity or diagonal)
+
+        # For simplicity, we assume a diagonal weighting matrix (you can adjust based on problem setup).
+        K_weight = torch.eye(dim, **self.tensor_args)
+        K_weight = K_weight.unsqueeze(0).repeat(batch_size, self.n_support_points, 1, 1)
+
+        # Fill in the matrices with appropriate dimensions
+        for i in range(self.n_support_points):
+            idx_start = i * dim
+            idx_end = (i + 1) * dim
+
+            # Populate A matrix for this support point (Jacobian of the error wrt the state)
+            A[:, idx_start:idx_end, idx_start:idx_end] = K_weight[:, i]
+
+            # Populate b with the error at this support point
+            b[:, idx_start:idx_end] = error[:, i].unsqueeze(-1)
+
+            # Populate K with the weighting matrix
+            K[:, idx_start:idx_end, idx_start:idx_end] = K_weight[:, i]
+
+        return A, b, K
+
+
+class CostGP(Cost):
+
+    def __init__(
+            self,
+            robot,
+            n_support_points,
+            start_state,
+            dt,
+            sigma_params,
+            **kwargs
+    ):
+        super().__init__(robot, n_support_points, **kwargs)
+        self.start_state = start_state
+        self.dt = dt
+
+        self.sigma_start = sigma_params['sigma_start']
+        self.sigma_gp = sigma_params['sigma_gp']
+
+        self.set_cost_factors()
+
+    def set_cost_factors(self):
+        # ========= Cost factors ===============
+        self.start_prior = UnaryFactor(
+            self.dim,
+            self.sigma_start,
+            self.start_state,
+            self.tensor_args,
+        )
+
+        self.gp_prior = GPFactor(
+            self.n_dof,
+            self.sigma_gp,
+            self.dt,
+            self.n_support_points - 1,
+            self.tensor_args,
+        )
+
+    def eval(self, trajs, **observation):
+        # trajs = trajs.reshape(-1, self.n_support_points, self.dim)
+        # Start cost
+        err_p = self.start_prior.get_error(trajs[:, [0]], calc_jacobian=False)
+        w_mat = self.start_prior.K
+        start_costs = err_p @ w_mat.unsqueeze(0) @ err_p.transpose(1, 2)
+        start_costs = start_costs.squeeze()
+
+        # GP Trajectory cost
+        err_gp = self.gp_prior.get_error(trajs, calc_jacobian=False)
+        w_mat = self.gp_prior.Q_inv[0]  # repeated Q_inv
+        w_mat = w_mat.reshape(1, 1, self.dim, self.dim)
+        gp_costs = err_gp.transpose(2, 3) @ w_mat @ err_gp
+        gp_costs = gp_costs.sum(1)
+        gp_costs = gp_costs.squeeze()
+
+        costs = start_costs + gp_costs
+
+        return costs
+
+    def get_linear_system(self, trajs, **observation):
+        batch_size = trajs.shape[0]
+        A = torch.zeros(batch_size, self.dim * self.n_support_points, self.dim * self.n_support_points,
+                        **self.tensor_args)
+        b = torch.zeros(batch_size, self.dim * self.n_support_points, 1, **self.tensor_args)
+        K = torch.zeros(batch_size, self.dim * self.n_support_points, self.dim * self.n_support_points,
+                        **self.tensor_args)
+
+        # Start prior factor
+        err_p, H_p = self.start_prior.get_error(trajs[:, [0]])
+        A[:, :self.dim, :self.dim] = H_p
+        b[:, :self.dim] = err_p
+        K[:, :self.dim, :self.dim] = self.start_prior.K
+
+        # GP factors
+        err_gp, H1_gp, H2_gp = self.gp_prior.get_error(trajs)
+
+        A[:, self.dim:, :-self.dim] = torch.block_diag(*H1_gp)
+        A[:, self.dim:, self.dim:] += torch.block_diag(*H2_gp)
+        b[:, self.dim:] = einops.rearrange(err_gp, "b h d 1 -> b (h d) 1")
+        K[:, self.dim:, self.dim:] += torch.block_diag(*self.gp_prior.Q_inv)
+
+        # old code not vectorized
+        # https://github.com/anindex/stoch_gpmp/blob/main/stoch_gpmp/costs/cost_functions.py#L161
+
+        return A, b, K
+
+
+class CostGPTrajectory(Cost):
+
+    def __init__(
+            self,
+            robot,
+            n_support_points,
+            dt,
+            sigma_gp=None,
+            **kwargs
+    ):
+        super().__init__(robot, n_support_points, **kwargs)
+        self.dt = dt
+
+        self.sigma_gp = sigma_gp
+
+        self.set_cost_factors()
+
+    def set_cost_factors(self):
+        # ========= Cost factors ===============
+        self.gp_prior = GPFactor(
+            self.n_dof,
+            self.sigma_gp,
+            self.dt,
+            self.n_support_points - 1,
+            self.tensor_args,
+        )
+
+    def eval(self, trajs, **observation):
+        # trajs = trajs.reshape(-1, self.n_support_points, self.dim)
+
+        # GP cost
+        err_gp = self.gp_prior.get_error(trajs, calc_jacobian=False)
+        w_mat = self.gp_prior.Q_inv[0]  # repeated Q_inv
+        w_mat = w_mat.reshape(1, 1, self.dim, self.dim)
+        gp_costs = err_gp.transpose(2, 3) @ w_mat @ err_gp
+        gp_costs = gp_costs.sum(1).squeeze()
+        costs = gp_costs
+        return costs
+
+    def get_linear_system(self, trajs, **observation):
+        pass
+
+
+class CostGPTrajectoryPositionOnlyWrapper(CostGPTrajectory):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def eval(self, trajs, **observation):
+        vel = finite_difference_vector(trajs, dt=self.dt, method='central')
+        trajs_tmp = torch.cat((trajs, vel), dim=-1)
+        return super().eval(trajs_tmp, **observation)
+
+
+class CostSmoothnessCHOMP(Cost):
+
+    def __init__(
+            self,
+            robot,
+            n_support_points,
+            **kwargs
+    ):
+        super().__init__(robot, n_support_points, **kwargs)
+        self.dt = robot.dt
+
+        self.Sigma_inv = CHOMP._get_R_mat(dt=self.dt, n_support_points=n_support_points, **kwargs)
+
+    def eval(self, trajs, **observation):
+        R_mat = self.Sigma_inv
+        cost = batched_weighted_dot_prod(trajs, R_mat, trajs)
+        return cost
+
+    def get_linear_system(self, trajs, **observation):
+        pass
+
+
+class CostJointLimits(Cost):
+
+    def __init__(
+            self,
+            robot,
+            n_support_points,
+            eps=np.deg2rad(3),
+            **kwargs
+    ):
+        super().__init__(robot, n_support_points, **kwargs)
+
+        self.eps = eps
+
+    def eval(self, trajs, **observation):
+        assert trajs.ndim == 3
+
+        # trajs = trajs.reshape(-1, self.n_support_points, self.dim)
+        trajs_pos = self.robot.get_position(trajs)
+
+        idxs_lower = torch.argwhere(trajs_pos < self.robot.q_min + self.eps)
+        cost_lower = torch.pow(
+            self.robot.q_min[idxs_lower[:, 2]] + self.eps - trajs_pos[
+                idxs_lower[:, 0], idxs_lower[:, 1], idxs_lower[:, 2]],
+            2
+        ).sum(-1)
+
+        idxs_upper = torch.argwhere(trajs_pos > self.robot.q_max - self.eps)
+        cost_upper = torch.pow(
+            self.robot.q_max[idxs_upper[:, 2]] - self.eps - trajs_pos[
+                idxs_upper[:, 0], idxs_upper[:, 1], idxs_upper[:, 2]],
+            2
+        ).sum(-1)
+
+        costs = (cost_lower + cost_upper)
+
+        return costs
+
+    def get_linear_system(self, trajs, **observation):
+        pass
+
+
+class CostGoal(Cost):
+
+    def __init__(
+            self,
+            robot,
+            n_support_points,
+            field=None,
+            sigma_goal=None,
+            **kwargs
+    ):
+        super().__init__(robot, n_support_points, **kwargs)
+        self.field = field
+        self.sigma_goal = sigma_goal
+
+        self.set_cost_factors()
+
+    def set_cost_factors(self):
+        #========= Cost factors ===============
+        self.goal_factor = FieldFactor(
+            self.n_dof,
+            self.sigma_goal,
+            [-1, None]  # only take last point
+        )
+
+    def eval(self, trajs, x_trajs=None, **observation):
+        costs = 0
+        if self.field is not None:
+            err_obst = self.goal_factor.get_error(
+                trajs,
+                self.field,
+                x_trajs=x_trajs,
+                calc_jacobian=False,  # NOTE(an): no need for grads in StochGPMP
+            )
+            w_mat = self.goal_factor.K
+            obst_costs = w_mat * err_obst.sum(1)
+            costs = obst_costs
+
+        return costs
+
+    def get_linear_system(self, trajs, x_trajs=None, **observation):
+        A, b, K = None, None, None
+        if self.field is not None:
+            batch_size = trajs.shape[0]
+            A = torch.zeros(batch_size, 1, self.dim * self.n_support_points, **self.tensor_args)
+            err_goal, H_goal = self.goal_factor.get_error(
+                trajs,
+                self.field,
+                x_trajs=x_trajs,
+                calc_jacobian=True,
+            )
+            A[:, :, -self.dim:(-self.dim + self.n_dof)] = H_goal
+            b = err_goal.unsqueeze(-1)
+            K = self.goal_factor.K * torch.eye(1, **self.tensor_args).repeat(batch_size, 1, 1)
+        return A, b, K
+
+
+class CostGoalPrior(Cost):
+
+    def __init__(
+            self,
+            robot,
+            n_support_points,
+            multi_goal_states=None,  # num_goal x n_dim (pos + vel)
+            num_particles_per_goal=None,
+            num_samples=None,
+            sigma_goal_prior=None,
+            **kwargs
+    ):
+        super().__init__(robot, n_support_points, **kwargs)
+        self.multi_goal_states = multi_goal_states
+        self.num_goals = multi_goal_states.shape[0]
+        self.num_particles_per_goal = num_particles_per_goal
+        self.num_particles = num_particles_per_goal * self.num_goals
+        self.num_samples = num_samples
+        self.sigma_goal_prior = sigma_goal_prior
+
+        self.set_cost_factors()
+
+    def set_cost_factors(self):
+        self.multi_goal_prior = []
+        # TODO: remove this for loop
+        for i in range(self.num_goals):
+            self.multi_goal_prior.append(
+                UnaryFactor(
+                    self.dim,
+                    self.sigma_goal_prior,
+                    self.multi_goal_states[i],
+                    self.tensor_args,
+                )
+            )
+
+    def eval(self, trajs, **observation):
+        costs = 0
+        if self.multi_goal_states is not None:
+            x = trajs.reshape(self.num_goals, self.num_particles_per_goal * self.num_samples, self.n_support_points,
+                              self.dim)
+            costs = torch.zeros(self.num_goals, self.num_particles_per_goal * self.num_samples, **self.tensor_args)
+            # TODO: remove this for loop
+            for i in range(self.num_goals):
+                err_g = self.multi_goal_prior[i].get_error(x[i, :, [-1]], calc_jacobian=False)
+                w_mat = self.multi_goal_prior[i].K
+                goal_costs = err_g @ w_mat.unsqueeze(0) @ err_g.transpose(1, 2)
+                goal_costs = goal_costs.squeeze()
+                costs[i] += goal_costs
+            costs = costs.flatten()
+        return costs
+
+    def get_linear_system(self, trajs, **observation):
+        A, b, K = None, None, None
+        if self.multi_goal_states is not None:
+            npg = self.num_particles_per_goal
+            batch_size = npg * self.num_goals
+            x = trajs.reshape(self.num_goals, self.num_particles_per_goal, self.n_support_points, self.dim)
+            A = torch.zeros(batch_size, self.dim, self.dim * self.n_support_points, **self.tensor_args)
+            b = torch.zeros(batch_size, self.dim, 1, **self.tensor_args)
+            K = torch.zeros(batch_size, self.dim, self.dim, **self.tensor_args)
+            # TODO: remove this for loop
+            for i in range(self.num_goals):
+                err_g, H_g = self.multi_goal_prior[i].get_error(x[i, :, [-1]])
+                A[i * npg: (i + 1) * npg, :, -self.dim:] = H_g
+                b[i * npg: (i + 1) * npg] = err_g
+                K[i * npg: (i + 1) * npg] = self.multi_goal_prior[i].K
+
+        return A, b, K