diff --git a/README.md b/README.md index db01b13..28db3f0 100644 --- a/README.md +++ b/README.md @@ -56,106 +56,87 @@ python -m venv .venv && source .venv/bin/activate pip install -e . ``` -## Demos & entry points +## Command-line interface (`grl-snam`) -GRL-SNAM installs **two console scripts** plus a **volrover3 REPL** entry point. -GRL-SNAM and volrover3 are installed **separately** and do **not** depend on each -other — both depend only on `pycvc` + `pycvc_gl` (from libcvc / cvcpkg). The lab -detects volrover3 at runtime via the host-injected `vrhost` module. - -### 1. Numerical correctness demo (console) — `grl-snam-selftest` -Exercises the surrogate navigation dynamics (`integrate_surrogate_v2`) on small, -seeded, self-contained cases; needs `torch`, runs in ~1 s, no dataset. +Installing GRL-SNAM puts one console script on your `PATH` — **`grl-snam`** — the single +entry point for every workflow, from a raw world model to a running demo. (`grl-snam`, +`grl-snam-selftest`, and `grl-snam-lab-demo` are the installed scripts; the last two are +aliases for `grl-snam selftest` / `grl-snam lab-demo`.) ```bash -grl-snam-selftest # or: python -m grl_snam.selftest -``` -**Expected output** — four `PASS` checks and `OK`: +grl-snam --help ``` -grl-snam-selftest: exercising integrate_surrogate_v2 (surrogate navigation dynamics) - [PASS] shapes (o,v,min_clear) o(1, 2) v(1, 2) clr(1,) - [PASS] goal-seeking reduces distance |o0-goal|=20.000 -> |oT-goal|=0.00x - [PASS] goal-seeking output finite ... - [PASS] IPC barrier path stays finite min_clear=... - [PASS] deterministic (reproducible) two runs bit-identical -grl-snam-selftest: OK — surrogate dynamics correct -``` -It proves: goal-seeking (the damped spring reduces distance to goal), the IPC -collision barrier stays numerically finite, and the integrator is deterministic. -### 2. Standalone visual lab demo (console) — `grl-snam-lab-demo` -Builds a demo scene (terrain + an agent track + a marker) in a **self-owned** -window (needs `pycvc`/`pycvc_gl` + a GL display), or renders it offscreen: +| Command | What it does | +|---|---| +| `grl-snam selftest` | numerical correctness check (surrogate dynamics; no data) | +| `grl-snam obstacles BUNDLE` | world model → circular-obstacle set (`.npz`) | +| `grl-snam build-sdf BUNDLE` | world model → navigation SDF (`.npz`) | +| `grl-snam train SDF.npz` | self-supervised SDF-coefficient training | +| `grl-snam capture drive\|multigoal …` | drive the policy → mp4 offscreen, with a live HUD | +| `grl-snam pipeline BUNDLE` | the whole chain: world model → SDF → train → video | +| `grl-snam demo NAME` | run a demo live inside VolRover3 | +| `grl-snam lab-demo [PNG]` | standalone lab visualization | +| `grl-snam eval` / `train-coef` | the CoefEnergyNet evaluator / trainer (legacy track) | + +A `BUNDLE` is a **world model** — a directory with `terrain.json` (heightfield) + +`buildings.glb` (city mesh), e.g. an OpenStreetMap + SRTM export. The heavy work lives +in `grl_snam.tools`; the demos in `grl_snam.demos`; the graphics deps (`pycvc_gl`, VTK, +`vrhost`) load lazily, so `grl-snam --help` runs anywhere. + +### The learned-nav pipeline, end to end + +From nothing but a world model to a rendered demo, in one command: ```bash -grl-snam-lab-demo # opens an interactive window -grl-snam-lab-demo out.png # headless: writes a PNG snapshot +grl-snam pipeline path/to/austin_south # build-sdf → train → multigoal video ``` -**Expected to see** — a green Gaussian-bump **terrain** surface, a yellow looping -**agent track** draped above it, and a red **marker** at the agent's start. -### 3. Live demo inside volrover3 (REPL) — `grl_snam_lab.run_in_volrover()` -The same scene, but rendered into a **running volrover3's live viewport** — the -whole point of the embedded interpreter. In volrover3's Python Console (REPL tab): +…or step by step (each artifact is reusable): -```python ->>> import grl_snam_lab ->>> grl_snam_lab.run_in_volrover() -grl_snam_lab: live volrover3 scene now has 3 node(s) (terrain + agent0_track + agent0) — look at the viewport. +```bash +grl-snam build-sdf path/to/austin_south # → nav_sdf.npz +grl-snam train path/to/austin_south/nav_sdf.npz -o coef_sdf.pt --steps 1500 +grl-snam capture multigoal path/to/austin_south coef_sdf.pt --minutes 3 -o drive.mp4 ``` -**Expected to see** — the terrain + track + marker appear **in the volrover3 3D -window you already have open** (not a separate window); you can orbit/zoom them -with the app's camera, toggle them in the scene tree, and drive them further from -the REPL. Outside volrover3 (`vrhost` absent) it prints a clear message pointing -you at `run_standalone()` / `grl-snam-lab-demo`. -### 4. Learned SDF navigation on real geometry — the live drive demos +`build-sdf`/`train` run on CPU in a few minutes (GPU auto-used if present). +`capture` renders the real 3-D scene **offscreen** through VTK (the VolRover3 engine — +not a screen grab) and draws a **live metrics HUD** on every frame: the network's +predicted coefficients (α, β, γ), wall clearance, speed, escape mode, and building +penetration. `--source cvc` on `build-sdf` uses CVC's mesh-exact 3-D SDF instead of the +footprint EDT. See [docs/training-navigation-on-geometry.md](docs/training-navigation-on-geometry.md). -The navigator can be **trained on a real scene** (terrain + city mesh) and driven -live in volrover3. Two prep steps (once), then load a demo as a volrover3 job. See -[docs/training-navigation-on-geometry.md](docs/training-navigation-on-geometry.md). +### Live inside VolRover3 -**Prep — build the SDF + train the policy** (CPU, a few minutes): ```bash -python scripts/build_sdf.py --source edt # footprint SDF (or --source cvc for cvc::sdf) -python scripts/train_sdf.py /nav_sdf.npz -o checkpoints/coef_sdf.pt --steps 1500 +grl-snam demo --list # the registered demos +grl-snam demo austin-freedrive --bundle path/to/austin_south --checkpoint coef_sdf.pt ``` -**Run inside volrover3** — from the app's **Jobs tab → Load Script… → Run as Job**, -or headless via the CLI: -```bash -export GRL_SNAM_CHECKPOINT=$PWD/checkpoints/coef_sdf.pt -export GRL_SNAM_SCENE_BUNDLE=/path/to/austin_south # dir with terrain.json + buildings.glb -volrover3 --run-job examples/volrover_grl_snam_austin_freedrive.py -``` -`--run-job` loads a file that defines `step(dt)` and ticks it cooperatively (it -appears in the Python Console Jobs tab). Env vars: `GRL_SNAM_CHECKPOINT` (trained -`.pt`), `GRL_SNAM_SCENE_BUNDLE` (scene dir), `GRL_SNAM_SDF` (prebuilt `nav_sdf.npz`; -else built at load), `GRL_SNAM_ROOT` (repo path if not on the default). - -Drive demos in `examples/`: -- **`volrover_grl_snam_austin_freedrive.py`** — end-to-end: the learned SDF policy - finds its own way START→GOAL with **no route**, plus a wall-follow escape for - local minima. Retarget the goal live to watch it redirect. -- **`volrover_grl_snam_austin_learned.py`** — stagewise: an A* occupancy route is a - collision-free spine; the learned policy drives locally between sub-goals (robust - for adversarial start/goal pairs). -- **`volrover_grl_snam_planner.py`** — the surrogate's native sparse-obstacle regime. - -**Capture a drive to video — offscreen, no window** (the same way the demo films are -made: a scripted chase camera through VTK — the VolRover3 engine — not a screen grab): +`grl-snam demo NAME` shells out to `volrover3 --run-job` with the demo module (set +`VOLROVER3_BIN` if `volrover3` isn't on `PATH`); you can also load the module from the +app's **Jobs tab → Load Script → Run as Job**. Each demo publishes the same live metrics +into the state tree under `grl_snam.metrics.*` — the hook for an in-app HUD panel — and +prints them to the console. Demos: + +- **`austin-freedrive`** — end-to-end learned SDF free-drive, **no route**, live goal retarget. +- **`austin-learned`** — stagewise: an A* route spine + learned SDF local control (robust). +- **`austin-planner`** — the surrogate in its native sparse-obstacle regime. +- **`austin-patrol`** — a grounded A* street patrol (non-learned baseline). +- **`lab`** — the analytic terrain + an agent walking a draped loop. + +Env for the Austin demos: `GRL_SNAM_SCENE_BUNDLE`, `GRL_SNAM_CHECKPOINT`, `GRL_SNAM_SDF` +(prebuilt `nav_sdf.npz`; else built at load). + +### selftest & lab-demo + ```bash -# one fixed A→B run -python scripts/capture_drive_video.py checkpoints/coef_sdf.pt \ - --start -361 114 --goal 185 50 -o drive.mp4 -# the DYNAMIC multi-goal free-drive: goals re-targeted live, drone chase cam, ~3 min -python scripts/capture_multigoal_video.py checkpoints/coef_sdf.pt \ - --sdf /nav_sdf.npz --minutes 3 -o multigoal.mp4 +grl-snam selftest # four PASS checks: goal-seeking, IPC-finite, deterministic +grl-snam lab-demo out.png # offscreen snapshot of the analytic lab (omit PNG for a window) ``` -Both run headless (offscreen GL is fine) and only need the volrover env + a trained -checkpoint + a scene bundle + ffmpeg — no display, no live app window. -### 5. Dataset-specific extensions +### Dataset-specific extensions Applied, dataset-specific work lives in **separate downstream projects** that depend on GRL-SNAM (not the other way round) and keep their dataset-specific code and models out of this general library. Such an extension reuses exactly the pieces shown above — the diff --git a/examples/NOTICE-austin.md b/examples/NOTICE-austin.md deleted file mode 100644 index 5ca87e4..0000000 --- a/examples/NOTICE-austin.md +++ /dev/null @@ -1,19 +0,0 @@ -# Attribution — Austin scene geometry (`volrover_austin_demo.py`) - -The `volrover_austin_demo.py` example renders real-world Austin, TX geometry from -a `geometry_bundle` (e.g. an `austin_south` OSM/SRTM export). That geometry is -derived from public sources and must be credited when you publish renders: - -- **Buildings, roads, water** — © **OpenStreetMap** contributors, licensed under - the **Open Database License (ODbL)**, https://openstreetmap.org/copyright. - Attribution to OpenStreetMap is **required** for any published render or - derived work. -- **Terrain elevation** — **SRTM** (NASA/USGS), **U.S. public domain**. - -Do **not** ship the Esri `satellite.png` overlay that some bundles carry — it is -proprietary Esri World Imagery and is not redistributable. The loaders in -`grl_snam_lab.scenes` never load it; the demo colors the ground itself. - -The bundle geometry itself is not formally licensed in-repo; the above is read -from its generation pipeline (OSM via osmnx/Overpass, SRTM via Open Topo Data). -When in doubt, treat it as ODbL + attribution-required. diff --git a/examples/lab_demo.py b/examples/lab_demo.py deleted file mode 100644 index 90fc193..0000000 --- a/examples/lab_demo.py +++ /dev/null @@ -1,76 +0,0 @@ -"""Demo: build a synthetic GRL-SNAM scene in the lab and show / snapshot it. - -General GRL-SNAM only — a wavy terrain, a couple of obstacle meshes, two -agent trajectories, agent markers, and a risk field. A downstream adapter -would instead feed its own real track data through the same Lab API. - -Run (with pycvc / pycvc_gl on PYTHONPATH and a display): - python examples/lab_demo.py # interactive window - python examples/lab_demo.py out.png # offscreen snapshot -""" - -import math -import sys - -from grl_snam_lab import Lab - - -def _wavy_terrain(n=40, amp=3.0): - return [ - [amp * math.sin(0.3 * r) * math.cos(0.3 * c) for c in range(n)] - for r in range(n) - ] - - -def build_demo() -> Lab: - lab = Lab() - lab.add_terrain(_wavy_terrain(), bounds=(-100, -100, 100, 100), color=(0.3, 0.5, 0.3)) - - # two boxy "obstacles" as simple prisms (just top faces here for brevity) - lab.add_mesh( - "obstacle_a", - [-40, -40, 5, -20, -40, 5, -20, -20, 5, -40, -20, 5], - [0, 1, 2, 0, 2, 3], - color=(0.6, 0.6, 0.65), - ) - lab.add_mesh( - "obstacle_b", - [20, 10, 8, 45, 10, 8, 45, 35, 8, 20, 35, 8], - [0, 1, 2, 0, 2, 3], - color=(0.6, 0.6, 0.65), - ) - - # two agent trajectories crossing the field - lab.add_path( - "agent0", - [(-90 + i * 4, -30 + 20 * math.sin(0.15 * i), 6) for i in range(45)], - color=(1.0, 0.2, 0.2), - ) - lab.add_path( - "agent1", - [(-90 + i * 4, 40 - 15 * math.sin(0.2 * i), 6) for i in range(45)], - color=(0.2, 0.4, 1.0), - ) - lab.add_markers("agents", [(-90, -30, 6), (-90, 40, 6)]) - - # a risk field (higher near the obstacles); flat 8x8x4 grid - nx, ny, nz = 8, 8, 4 - field = [] - for k in range(nz): - for j in range(ny): - for i in range(nx): - field.append(float((i - 4) ** 2 + (j - 4) ** 2 + k)) - lab.add_field("risk", field, dims=(nx, ny, nz), bounds=(-100, -100, 0, 100, 100, 30)) - - lab.pump() - return lab - - -if __name__ == "__main__": - lab = build_demo() - print(f"GRL-SNAM lab demo: {lab.num_nodes()} scene nodes") - if len(sys.argv) > 1: - lab.render_png(sys.argv[1]) - print(f"wrote {sys.argv[1]}") - else: - lab.show() diff --git a/examples/volrover_austin_demo.py b/examples/volrover_austin_demo.py deleted file mode 100644 index 73a6fcb..0000000 --- a/examples/volrover_austin_demo.py +++ /dev/null @@ -1,172 +0,0 @@ -# examples/volrover_austin_demo.py — a robot patrols REAL Austin, TX geometry in -# a LIVE volrover3 window, GROUNDED: it drives the STREETS (routed around the -# buildings, draped on the terrain), filmed by a low third-person chase camera. -# -# Loads a geometry_bundle (terrain heightfield + glTF city mesh) — by default an -# "austin_south" OSM/SRTM bundle (3 km x 3 km around the UT campus) — rasterizes the -# buildings into an occupancy grid, A*-plans a grounded patrol loop through the -# free space (streets), drapes it on the real terrain, and follows the vehicle -# with grl_snam_lab.camera.ChaseCamera (position-driven, ready for live GRL-SNAM -# paths). -# -# Load it from volrover3's Python Console -> Jobs tab -> "Load Script..." -> -# "Run as Job" (the glTF is ~978k triangles, so the first tick takes several -# seconds to load + rasterize). Point it at a different bundle with the -# GRL_SNAM_SCENE_BUNDLE env var. -# -# ATTRIBUTION: the geometry is derived from OpenStreetMap (c) OpenStreetMap -# contributors, ODbL (https://openstreetmap.org/copyright); SRTM terrain is US -# public domain. Credit OpenStreetMap in any published render. -# -# LAYERING: volrover3 does NOT depend on GRL-SNAM (or on the scene assets). This -# runs under volrover3's generic job runner when GRL-SNAM is installed and the -# bundle is present. - -import math -import os - -from grl_snam_lab.camera import ChaseCamera -from grl_snam_lab.lab import Lab -from grl_snam_lab.scenes import ( - building_occupancy, - load_geometry_bundle, - plan_ground_route, - resample_polyline, - terrain_grid, -) -from grl_snam_lab.vehicle import VehiclePose - -try: - import pycvc - import vrhost -except ImportError as exc: # pragma: no cover - only meaningful inside volrover3 - raise RuntimeError( - "volrover_austin_demo: `vrhost`/`pycvc` not found — load this INSIDE " - "volrover3's embedded Python (Jobs tab -> Load Script -> Run as Job)." - ) from exc - -_BUNDLE = os.environ.get( - "GRL_SNAM_SCENE_BUNDLE", - os.path.expanduser("~/scenes/austin_south"), # set GRL_SNAM_SCENE_BUNDLE to your bundle dir -) - -_app = vrhost.app() -_lab = Lab(app=_app, scene=vrhost.scene()) -_lab.set_axis_visible(False) # hide the XYZ gnomon — a distraction for this demo -# Real Austin terrain + buildings; `_sample(x, y)` is the terrain height (drape). -_sample = load_geometry_bundle(_lab, _BUNDLE) - -# ── plan a GROUNDED patrol route through the streets (around the buildings) ─── -# Rasterize the city into an occupancy grid, then A*-route a closed patrol loop -# through the free space and drape it on the terrain. The vehicle stays on the -# ground and never drives through a building. -_bounds = terrain_grid(os.path.join(_BUNDLE, "terrain.json"))[1] -print("volrover_austin_demo: building/loading the building occupancy grid...") -# SOLID footprint mask (filled), grown by ~12 m so a vehicle keeps clear of walls. -# Cached next to the .glb, so this is instant after the first run. -_occ = building_occupancy(os.path.join(_BUNDLE, "buildings.glb"), _bounds, nx=512, ny=512, inflate_m=12.0) - -_R = 430.0 # patrol radius; waypoints are snapped to the nearest open street -_WAYPTS = [(_R * math.cos(a), _R * math.sin(a)) for a in - [0.0, math.pi / 3, 2 * math.pi / 3, math.pi, 4 * math.pi / 3, 5 * math.pi / 3]] -_route2d = resample_polyline(plan_ground_route(_occ, _bounds, _WAYPTS, close_loop=True), spacing=4.0) -if len(_route2d) < 2: # pathing failed (bad bundle) — fall back to a draped ring - _route2d = [( _R * math.cos(2 * math.pi * k / 240), _R * math.sin(2 * math.pi * k / 240)) - for k in range(241)] -print("volrover_austin_demo: grounded route = %d pts along the streets." % len(_route2d)) - -_PPS = 3.0 # route points/sec -> ~12 m/s at 4 m spacing - - -def _pose(t): - # (x, y) along the uniformly-resampled loop (seamless wrap). Draping (z) and - # heading are handled downstream — VehiclePose derives heading + slope from the - # position stream, and the chase camera drapes its own target. - n = len(_route2d) - f = (t * _PPS) % n - i = int(f) - j = (i + 1) % n - a = f - i - x = _route2d[i][0] * (1 - a) + _route2d[j][0] * a - y = _route2d[i][1] * (1 - a) + _route2d[j][1] * a - return x, y - - -def _vehicle_mesh(): - # A car-sized box in LOCAL coords: forward = +X, centered in X/Y, base at z=0. - L, W, H = 4.6, 2.0, 1.6 - hx, hy = L / 2.0, W / 2.0 - v = [-hx, -hy, 0, hx, -hy, 0, hx, hy, 0, -hx, hy, 0, - -hx, -hy, H, hx, -hy, H, hx, hy, H, -hx, hy, H] - t = [0, 1, 2, 0, 2, 3, 4, 6, 5, 4, 7, 6, 1, 2, 6, 1, 6, 5, - 0, 7, 4, 0, 3, 7, 3, 2, 6, 3, 6, 7, 0, 5, 1, 0, 4, 5] - return v, t - - -_vv, _vt = _vehicle_mesh() -_lab.add_mesh("agent0", _vv, _vt, color=(0.90, 0.12, 0.12)) -_lab.add_path("agent0_track", [(x, y, _sample(x, y) + 0.6) for x, y in _route2d], color=(0.95, 0.75, 0.10)) - -# Smoothly yaw the vehicle to its heading AND pitch/roll it to the terrain normal -# (banks over crests, leans on side-slopes) with camera-style damping. -_vpose = VehiclePose(_sample, lift=0.25) - - -def _place_vehicle(x, y, dt): - _lab.node("agent0").setTransform(_vpose.update(x, y, dt)) - - -# ── low third-person chase camera (grounded, street-level) ─────────────────── -_CAM = "volrover3.camera" -_chase = ChaseCamera(back=34.0, height=13.0, look_up=2.5, up=(0.0, 0.0, 1.0)) - - -def _cset(k, v): - pycvc.state_set(_app, _CAM + "." + k, "%.6f" % float(v)) - - -def _drive(eye, tgt, up): - vx, vy, vz = (tgt[i] - eye[i] for i in range(3)) - m = math.sqrt(vx * vx + vy * vy + vz * vz) or 1.0 - _cset("position.x", eye[0]) - _cset("position.y", eye[1]) - _cset("position.z", eye[2]) - _cset("view_direction.x", vx / m) - _cset("view_direction.y", vy / m) - _cset("view_direction.z", vz / m) - _cset("up_vector.x", up[0]) - _cset("up_vector.y", up[1]) - _cset("up_vector.z", up[2]) - _cset("fov", 60.0) - - -def _cam_pos(x, y): - # feed the chase camera the vehicle's draped ground position (x, y, z) - return (x, y, _sample(x, y)) - - -_DT = 1.0 / 30.0 -_t = 0.0 -for _i in range(30): # prime the camera + vehicle smoothing to a settled pose - _x, _y = _pose(_i * _DT) - _place_vehicle(_x, _y, _DT) - _e, _g, _u = _chase.update(_cam_pos(_x, _y), _DT) -_drive(_e, _g, _u) -_t = 30 * _DT -_lab.pump() - -print( - "grl_snam_lab: Austin GROUNDED chase-cam demo loaded (OSM/ODbL + SRTM terrain). " - "The red agent drives the streets, routed around the buildings; camera follows. " - "Pause/stop from the Jobs tab." -) - - -def step(dt): - global _t - _t += dt - x, y = _pose(_t) - _place_vehicle(x, y, dt) # yaw to heading + pitch/roll to the terrain, smoothed - eye, tgt, up = _chase.update(_cam_pos(x, y), dt) - _drive(eye, tgt, up) - _lab.pump() diff --git a/examples/volrover_grl_snam_austin_freedrive.py b/examples/volrover_grl_snam_austin_freedrive.py deleted file mode 100644 index e3932bf..0000000 --- a/examples/volrover_grl_snam_austin_freedrive.py +++ /dev/null @@ -1,172 +0,0 @@ -# examples/volrover_grl_snam_austin_freedrive.py — the learned GRL-SNAM navigator -# finding its OWN way across REAL Austin, TX in a LIVE volrover3 window, with NO A* -# route. Pick a START and a GOAL; the trained SDF policy drives there end-to-end, -# purely by reacting to the buildings: a moving carrot toward the goal + the SDF -# barrier deflecting around footprints. This is "the vehicle finds its own path -# from A to B based on environmental factors" — no precomputed path at all. -# -# vs volrover_grl_snam_austin_learned.py (which uses an A* route as a stage spine): -# this one is pure end-to-end learned navigation. A pure potential field can stall -# in a local minimum (a dead-end / U-shaped block) for adversarial A->B pairs, so -# the START/GOAL below are chosen to be navigable; retarget as you like. -# -# PREP (once): build the SDF + train the coefficients (see docs) -> -# python scripts/build_sdf.py --source edt -# python scripts/train_sdf.py /nav_sdf.npz -o checkpoints/coef_sdf_austin.pt -# ENV: GRL_SNAM_SCENE_BUNDLE, GRL_SNAM_CHECKPOINT (coef_sdf .pt), GRL_SNAM_SDF -# (prebuilt nav_sdf.npz; else built from occupancy at load), GRL_SNAM_ROOT. -# RUN (inside volrover3): Python Console -> Jobs tab -> "Load Script..." -> Run as Job. -# ATTRIBUTION: OpenStreetMap (c) contributors, ODbL; SRTM terrain US public domain. - -import math -import os -import sys - -import numpy as np - -_ROOT = os.environ.get("GRL_SNAM_ROOT", "/home/joe/src/cvc/GRL-SNAM") -if _ROOT not in sys.path: - sys.path.insert(0, _ROOT) - -import torch # noqa: E402 - -import sdf_nav # noqa: E402 -from pycvc_gl.camera import ChaseCamera # noqa: E402 -from pycvc_gl.lab import Lab # noqa: E402 -from pycvc_gl.scenes import building_occupancy, load_geometry_bundle, terrain_grid # noqa: E402 -from pycvc_gl.vehicle import VehiclePose # noqa: E402 - -try: - import pycvc - import vrhost -except ImportError as exc: # pragma: no cover - raise RuntimeError("run INSIDE volrover3 (Jobs tab -> Load Script)") from exc - -_BUNDLE = os.environ.get("GRL_SNAM_SCENE_BUNDLE", - os.path.expanduser("~/scenes/austin_south")) # set to your bundle dir -_CKPT = os.environ.get("GRL_SNAM_CHECKPOINT", os.path.join(_ROOT, "checkpoints", "coef_sdf_austin.pt")) -torch.set_num_threads(2) - -_ck = torch.load(_CKPT, map_location="cpu"); _m = _ck["meta"] -_S = float(_m["scale"]); _CTR = np.asarray(_m["center"], np.float32) -_RR = float(_m["rr"]); _DHAT = float(_m["d_hat"]); _DT = float(_m["dt"]); _NSUB = int(_m["nsub"]); _VMAX = float(_m["vmax"]) -_model = sdf_nav.CoefMLP(); _model.load_state_dict(_ck["model_state_dict"]); _model.eval() - -# START -> GOAL (world x,y). Validated navigable on austin_south; retarget freely. -_START = (-361.0, 114.0) -_GOAL = (185.0, 50.0) - -_app = vrhost.app() -_lab = Lab(app=_app, scene=vrhost.scene()); _lab.set_axis_visible(False) -_sample = load_geometry_bundle(_lab, _BUNDLE) -_bounds = terrain_grid(os.path.join(_BUNDLE, "terrain.json"))[1]; _mnx, _mny, _mxx, _mxy = _bounds -print("grl_snam_austin_freedrive: loading occupancy + SDF...") -_occ0 = building_occupancy(os.path.join(_BUNDLE, "buildings.glb"), _bounds, nx=512, ny=512, inflate_m=0.0) -_NY, _NX = _occ0.shape -_sdf_npz = os.environ.get("GRL_SNAM_SDF", "") -if _sdf_npz and os.path.exists(_sdf_npz): - _d = np.load(_sdf_npz); _phi, _nxg, _nyg = _d["phi"], _d["normal_x"], _d["normal_y"] -else: - _phi, _nxg, _nyg = sdf_nav.build_sdf(_occ0, _bounds, _S) -_field = sdf_nav.SDFField(_phi, _nxg, _nyg, _bounds, _CTR, _S, device="cpu") -_kw = dict(rr=_RR, d_hat=_DHAT, dt=_DT, vmax=_VMAX) - - -def _w2n(p): - return np.array([(p[0] - _CTR[0]) * _S, (p[1] - _CTR[1]) * _S], np.float32) - - -def _n2w(on): - return np.array([on[0] / _S + _CTR[0], on[1] / _S + _CTR[1]], np.float32) - - -_lab.add_markers("start", [(_START[0], _START[1], _sample(*_START) + 1.0)], color=(0.15, 0.85, 0.25)) -_lab.add_markers("goal", [(_GOAL[0], _GOAL[1], _sample(*_GOAL) + 1.0)], color=(0.95, 0.80, 0.10)) - - -def _vehicle_mesh(): - L, W, Hh = 4.6, 2.0, 1.6; hx, hy = L / 2, W / 2 - v = [-hx, -hy, 0, hx, -hy, 0, hx, hy, 0, -hx, hy, 0, -hx, -hy, Hh, hx, -hy, Hh, hx, hy, Hh, -hx, hy, Hh] - t = [0, 1, 2, 0, 2, 3, 4, 6, 5, 4, 7, 6, 1, 2, 6, 1, 6, 5, 0, 7, 4, 0, 3, 7, 3, 2, 6, 3, 6, 7, 0, 5, 1, 0, 4, 5] - return v, t - - -_vv, _vt = _vehicle_mesh(); _lab.add_mesh("agent0", _vv, _vt, color=(0.90, 0.12, 0.12)) -_vpose = VehiclePose(_sample, lift=0.25) - -# ── end-to-end learned navigation state (no route — carrot toward the true goal) ── -_o = torch.from_numpy(_w2n(_START)).unsqueeze(0).float(); _v = torch.zeros(1, 2) -_GN = _w2n(_GOAL); _done = False -_best = 1e9; _stall = 0; _mode = "seek"; _turn = 1.0; _dhit = 0.0 # wall-follow escape state - - -@torch.no_grad() -def _sdf_normal(on): - _, nrm = _field.sample(torch.from_numpy(on).unsqueeze(0).float()) - return nrm[0].numpy() - - -@torch.no_grad() -def _nav_step(): - """One end-to-end step. Normally aim a carrot straight at the goal; if progress - STALLS (a potential-field local minimum), follow the wall tangentially — a - 'bug'-style escape — until the agent rounds the obstacle and progress resumes.""" - global _o, _v, _done, _best, _stall, _mode, _turn, _dhit - if _done: - w = _n2w(_o[0].numpy()); return float(w[0]), float(w[1]) - p = _o[0].numpy(); dg = float(np.linalg.norm(_GN - p)); gdir = (_GN - p) / (dg + 1e-6) - if dg < _best - 1e-3: - _best = dg; _stall = 0 - else: - _stall += 1 - if _mode == "seek" and _stall > 70: # stuck -> start following the wall - nrm = _sdf_normal(p); t = np.array([-nrm[1], nrm[0]], np.float32) - _turn = 1.0 if np.dot(t, gdir) >= 0 else -1.0; _dhit = dg; _mode = "wall"; _stall = 0 - if _mode == "wall": - nrm = _sdf_normal(p); t = _turn * np.array([-nrm[1], nrm[0]], np.float32) - carrot = (p + (0.6 * t + 0.4 * nrm) * 1.6).astype(np.float32) # slide along the wall - if dg < _dhit - 1.2 or _stall > 240: # rounded it (or give up escaping) - _mode = "seek"; _best = dg; _stall = 0 - else: - carrot = (p + gdir * min(1.8, dg)).astype(np.float32) # local goal toward the true goal - gt = torch.from_numpy(carrot).unsqueeze(0) - al, be, ga = _model(sdf_nav.coef_feats(_field, _o, gt)) - _o, _v, _ = sdf_nav.sdf_rollout(_field, _o, _v, gt, al, be, ga, 1, nsub=_NSUB, **_kw) - if float(np.linalg.norm(_o[0].numpy() - _GN)) < 0.4: - _done = True - w = _n2w(_o[0].numpy()); return float(w[0]), float(w[1]) - - -_CAM = "volrover3.camera"; _chase = ChaseCamera(back=34.0, height=13.0, look_up=2.5, up=(0.0, 0.0, 1.0)) - - -def _cset(k, val): - pycvc.state_set(_app, _CAM + "." + k, "%.6f" % float(val)) - - -def _drive_cam(eye, tgt, up): - vx, vy, vz = (tgt[i] - eye[i] for i in range(3)); mm = math.sqrt(vx * vx + vy * vy + vz * vz) or 1.0 - _cset("position.x", eye[0]); _cset("position.y", eye[1]); _cset("position.z", eye[2]) - _cset("view_direction.x", vx / mm); _cset("view_direction.y", vy / mm); _cset("view_direction.z", vz / mm) - _cset("up_vector.x", up[0]); _cset("up_vector.y", up[1]); _cset("up_vector.z", up[2]); _cset("fov", 60.0) - - -def _place(x, y, dt): - _lab.node("agent0").setTransform(_vpose.update(x, y, dt)) - - -_STEPS_PER_FRAME = 3; _FDT = 1.0 / 30.0 -for _i in range(30): - _place(_START[0], _START[1], _FDT) - _e, _g, _u = _chase.update((_START[0], _START[1], _sample(*_START)), _FDT) -_drive_cam(_e, _g, _u); _lab.pump() -print("grl_snam_austin_freedrive: the learned SDF policy is finding its own way " - "%s -> %s across Austin (no route). Pause/stop from Jobs." % (_START, _GOAL)) - - -def step(dt): - for _ in range(_STEPS_PER_FRAME): - x, y = _nav_step() - _place(x, y, dt) - eye, tgt, up = _chase.update((x, y, _sample(x, y)), dt) - _drive_cam(eye, tgt, up); _lab.pump() diff --git a/examples/volrover_grl_snam_austin_learned.py b/examples/volrover_grl_snam_austin_learned.py deleted file mode 100644 index e50d868..0000000 --- a/examples/volrover_grl_snam_austin_learned.py +++ /dev/null @@ -1,183 +0,0 @@ -# examples/volrover_grl_snam_austin_learned.py — the LEARNED GRL-SNAM navigator -# driving REAL Austin, TX in a LIVE volrover3 window, using the SDF obstacle model. -# -# A vehicle drives START -> GOAL across the city. At every step the trained SDF -# coefficient net (sdf_nav.CoefMLP) predicts the navigation coefficients and the -# DIFFERENTIABLE SDF surrogate (sdf_nav.sdf_rollout) integrates the motion — the -# learned policy genuinely NAVIGATES the streets (a signed distance field repels -# along the true wall normal, so it doesn't clip corners the way the circular- -# obstacle surrogate did). An A*/occupancy route supplies the global topology -# (GRL-SNAM's stagewise decomposition — pure potential fields have local minima); -# the learned surrogate drives within the street corridor, and a footprint check + -# the route spine guarantee the vehicle never enters a building. -# -# PREP (once, outside volrover3): build the SDF and train the coefficients -> -# python scripts/build_sdf.py --source edt # or --source cvc (mesh-exact 3-D) -# python scripts/train_sdf.py /nav_sdf.npz -o checkpoints/coef_sdf_austin.pt -# See docs/training-navigation-on-geometry.md. -# -# ENV: GRL_SNAM_SCENE_BUNDLE (scene dir), GRL_SNAM_CHECKPOINT (coef_sdf .pt), -# GRL_SNAM_SDF (prebuilt nav_sdf.npz; else built from the occupancy at load), -# GRL_SNAM_ROOT (repo path if not default). -# -# RUN (inside volrover3): Python Console -> Jobs tab -> "Load Script..." -> Run as Job. -# ATTRIBUTION: OpenStreetMap (c) contributors, ODbL; SRTM terrain US public domain. - -import math -import os -import sys - -import numpy as np - -_ROOT = os.environ.get("GRL_SNAM_ROOT", "/home/joe/src/cvc/GRL-SNAM") -if _ROOT not in sys.path: - sys.path.insert(0, _ROOT) - -import torch # noqa: E402 - -import sdf_nav # noqa: E402 -from pycvc_gl.camera import ChaseCamera # noqa: E402 -from pycvc_gl.lab import Lab # noqa: E402 -from pycvc_gl.scenes import ( # noqa: E402 - building_occupancy, load_geometry_bundle, plan_ground_route, resample_polyline, terrain_grid, -) -from pycvc_gl.vehicle import VehiclePose # noqa: E402 - -try: - import pycvc - import vrhost -except ImportError as exc: # pragma: no cover - raise RuntimeError("run INSIDE volrover3 (Jobs tab -> Load Script)") from exc - -_BUNDLE = os.environ.get("GRL_SNAM_SCENE_BUNDLE", - os.path.expanduser("~/scenes/austin_south")) # set to your bundle dir -_CKPT = os.environ.get("GRL_SNAM_CHECKPOINT", os.path.join(_ROOT, "checkpoints", "coef_sdf_austin.pt")) -torch.set_num_threads(2) - -# ── trained SDF policy + physics metadata ──────────────────────────────────── -_ck = torch.load(_CKPT, map_location="cpu"); _m = _ck["meta"] -_S = float(_m["scale"]); _CTR = np.asarray(_m["center"], np.float32) -_RR = float(_m["rr"]); _DHAT = float(_m["d_hat"]); _DT = float(_m["dt"]); _NSUB = int(_m["nsub"]); _VMAX = float(_m["vmax"]) -_REGION = float(_m["region"]) -_model = sdf_nav.CoefMLP(); _model.load_state_dict(_ck["model_state_dict"]); _model.eval() - -# ── live scene: real Austin terrain + buildings ────────────────────────────── -_app = vrhost.app() -_lab = Lab(app=_app, scene=vrhost.scene()); _lab.set_axis_visible(False) -_sample = load_geometry_bundle(_lab, _BUNDLE) -_bounds = terrain_grid(os.path.join(_BUNDLE, "terrain.json"))[1]; _mnx, _mny, _mxx, _mxy = _bounds -print("grl_snam_austin_learned(SDF): loading occupancy + SDF...") -_occR = building_occupancy(os.path.join(_BUNDLE, "buildings.glb"), _bounds, nx=512, ny=512, inflate_m=12.0) -_occ0 = building_occupancy(os.path.join(_BUNDLE, "buildings.glb"), _bounds, nx=512, ny=512, inflate_m=0.0) -_NY, _NX = _occ0.shape - -_sdf_npz = os.environ.get("GRL_SNAM_SDF", "") -if _sdf_npz and os.path.exists(_sdf_npz): - _d = np.load(_sdf_npz); _phi, _nxg, _nyg = _d["phi"], _d["normal_x"], _d["normal_y"] -else: # build the footprint-EDT SDF from the occupancy (no extra deps) - _phi, _nxg, _nyg = sdf_nav.build_sdf(_occ0, _bounds, _S) -_field = sdf_nav.SDFField(_phi, _nxg, _nyg, _bounds, _CTR, _S, device="cpu") - - -def _w2n(p): - return np.array([(p[0] - _CTR[0]) * _S, (p[1] - _CTR[1]) * _S], np.float32) - - -def _n2w(on): - return np.array([on[0] / _S + _CTR[0], on[1] / _S + _CTR[1]], np.float32) - - -def _in_building(xw, yw): - c = int((xw - _mnx) / (_mxx - _mnx) * (_NX - 1)); r = int((yw - _mny) / (_mxy - _mny) * (_NY - 1)) - return 0 <= r < _NY and 0 <= c < _NX and bool(_occ0[r, c]) - - -# ── global route (the stage planner): collision-free spine START -> GOAL ────── -_START = (-_REGION * 0.9, -_REGION * 0.75); _GOAL = (_REGION * 0.9, _REGION * 0.8) -print("grl_snam_austin_learned(SDF): planning route...") -_route = plan_ground_route(_occR, _bounds, [_START, _GOAL], close_loop=False) -if not _route or len(_route) < 2: - raise RuntimeError("route planning failed for this bundle/region") -_ROUTE = np.asarray(resample_polyline(_route, spacing=0.12 / _S), np.float32) -_lab.add_markers("start", [(_START[0], _START[1], _sample(*_START) + 1.0)], color=(0.15, 0.85, 0.25)) -_lab.add_markers("goal", [(_GOAL[0], _GOAL[1], _sample(*_GOAL) + 1.0)], color=(0.95, 0.80, 0.10)) -_lab.add_path("spine", [(w[0], w[1], _sample(w[0], w[1]) + 0.5) for w in _ROUTE], color=(0.30, 0.55, 0.95)) - - -def _vehicle_mesh(): - L, W, Hh = 4.6, 2.0, 1.6; hx, hy = L / 2, W / 2 - v = [-hx, -hy, 0, hx, -hy, 0, hx, hy, 0, -hx, hy, 0, -hx, -hy, Hh, hx, -hy, Hh, hx, hy, Hh, -hx, hy, Hh] - t = [0, 1, 2, 0, 2, 3, 4, 6, 5, 4, 7, 6, 1, 2, 6, 1, 6, 5, 0, 7, 4, 0, 3, 7, 3, 2, 6, 3, 6, 7, 0, 5, 1, 0, 4, 5] - return v, t - - -_vv, _vt = _vehicle_mesh(); _lab.add_mesh("agent0", _vv, _vt, color=(0.90, 0.12, 0.12)) -_vpose = VehiclePose(_sample, lift=0.25) - -# ── learned SDF navigation state ───────────────────────────────────────────── -_CORR = 0.35 / _S; _LOOK = 8 -_o = torch.from_numpy(_w2n(_ROUTE[0])).unsqueeze(0).float(); _v = torch.zeros(1, 2) -_ri = 0; _stall = 0; _done = False -_kw = dict(rr=_RR, d_hat=_DHAT, dt=_DT, vmax=_VMAX) - - -@torch.no_grad() -def _nav_step(): - """One learned-SDF step: the surrogate drives toward the look-ahead sub-goal and - is accepted while it stays in the street corridor and out of footprints; else the - drive advances along the clean spine. Returns the vehicle world (x, y).""" - global _o, _v, _ri, _stall, _done - if _done: - w = _ROUTE[_ri]; return float(w[0]), float(w[1]) - sub = _ROUTE[min(_ri + _LOOK, len(_ROUTE) - 1)]; goal = torch.from_numpy(_w2n(sub)).unsqueeze(0) - al, be, ga = _model(sdf_nav.coef_feats(_field, _o, goal)) - o2, v2, _ = sdf_nav.sdf_rollout(_field, _o, _v, goal, al, be, ga, 1, nsub=_NSUB, **_kw) - w = _n2w(o2[0].numpy()) - lo = max(0, _ri - 2); seg = _ROUTE[lo:min(_ri + 3 * _LOOK, len(_ROUTE))] - j = lo + int(np.argmin(np.linalg.norm(seg - w, axis=1))) - in_corr = (not _in_building(float(w[0]), float(w[1]))) and np.linalg.norm(w - _ROUTE[j]) < _CORR - if in_corr: - _o, _v = o2, v2; _stall = _stall + 1 if j <= _ri else 0; _ri = max(_ri, j) - else: - _stall += 1 - if (not in_corr) or _stall > 25: # spine nudge if strayed or stuck - _ri = min(_ri + 1, len(_ROUTE) - 1); w = _ROUTE[_ri].copy() - _o = torch.from_numpy(_w2n(w)).unsqueeze(0); _v = torch.zeros(1, 2); _stall = 0 - if np.linalg.norm(w - np.asarray(_GOAL, np.float32)) < 40.0 or _ri >= len(_ROUTE) - 1: - _done = True - return float(w[0]), float(w[1]) - - -# ── chase camera ───────────────────────────────────────────────────────────── -_CAM = "volrover3.camera"; _chase = ChaseCamera(back=34.0, height=13.0, look_up=2.5, up=(0.0, 0.0, 1.0)) - - -def _cset(k, val): - pycvc.state_set(_app, _CAM + "." + k, "%.6f" % float(val)) - - -def _drive_cam(eye, tgt, up): - vx, vy, vz = (tgt[i] - eye[i] for i in range(3)); mm = math.sqrt(vx * vx + vy * vy + vz * vz) or 1.0 - _cset("position.x", eye[0]); _cset("position.y", eye[1]); _cset("position.z", eye[2]) - _cset("view_direction.x", vx / mm); _cset("view_direction.y", vy / mm); _cset("view_direction.z", vz / mm) - _cset("up_vector.x", up[0]); _cset("up_vector.y", up[1]); _cset("up_vector.z", up[2]); _cset("fov", 60.0) - - -def _place(x, y, dt): - _lab.node("agent0").setTransform(_vpose.update(x, y, dt)) - - -_STEPS_PER_FRAME = 3; _FDT = 1.0 / 30.0 -for _i in range(30): - _place(_START[0], _START[1], _FDT) - _e, _g, _u = _chase.update((_START[0], _START[1], _sample(*_START)), _FDT) -_drive_cam(_e, _g, _u); _lab.pump() -print("grl_snam_austin_learned(SDF): the learned SDF policy is driving Austin. Pause/stop from Jobs.") - - -def step(dt): - for _ in range(_STEPS_PER_FRAME): - x, y = _nav_step() - _place(x, y, dt) - eye, tgt, up = _chase.update((x, y, _sample(x, y)), dt) - _drive_cam(eye, tgt, up); _lab.pump() diff --git a/examples/volrover_grl_snam_planner.py b/examples/volrover_grl_snam_planner.py deleted file mode 100644 index 07028b2..0000000 --- a/examples/volrover_grl_snam_planner.py +++ /dev/null @@ -1,229 +0,0 @@ -# examples/volrover_grl_snam_planner.py — a REAL GRL-SNAM navigation demo in a -# live volrover3 window: the GRL-SNAM surrogate planner determines a path from a -# START to a GOAL that navigates AROUND obstacles, and an agent drives it, -# filmed by the third-person chase camera. -# -# The path is produced by iterating grl_snam's differentiable surrogate rollout -# (dynamics.integrate_surrogate_v2) one step at a time — the same physics-informed -# integrator downstream training runs through (semi-implicit Euler + IPC obstacle -# barrier). Control coefficients here are hand-set to a known-good regime; with a -# trained CoefEnergyNet checkpoint you'd predict them per step and refine online -# with grl_snam.adaptation.HistSecantController (the "continuous learning" loop) — -# see the CONTINUOUS-LEARNING note below. -# -# HOW TO RUN (inside volrover3): Jobs tab -> "Load Script..." -> pick this file. -# Needs torch (a GRL-SNAM dependency) in the interpreter. Set GRL_SNAM_ROOT if the -# repo isn't at the default path. -# -# LAYERING: volrover3 has no GRL-SNAM dependency; this runs under its generic job -# runner when GRL-SNAM + torch are installed. - -import math -import os -import sys - -# grl_snam imports surrogate_robust from the repo root, so ensure it's importable. -_ROOT = os.environ.get("GRL_SNAM_ROOT", "/home/joe/src/cvc/GRL-SNAM") -if _ROOT not in sys.path: - sys.path.insert(0, _ROOT) - -import torch # noqa: E402 -from grl_snam.dynamics import integrate_surrogate_v2 # noqa: E402 - -from grl_snam_lab.camera import ChaseCamera # noqa: E402 -from grl_snam_lab.demo import terrain_height, _terrain_heights, TERRAIN_BOUNDS # noqa: E402 -from grl_snam_lab.lab import Lab # noqa: E402 - -try: - import pycvc - import vrhost -except ImportError as exc: # pragma: no cover - raise RuntimeError( - "volrover_grl_snam_planner: run INSIDE volrover3 (Jobs tab -> Load Script)." - ) from exc - -# ── scene: start, goal, obstacles (all in the terrain's x,y plane) ─────────── -START = (-70.0, -60.0) -GOAL = (70.0, 60.0) -OBSTACLES = [ # (x, y, radius) - (10.0, 20.0, 14.0), - (-20.0, 10.0, 12.0), - (25.0, -15.0, 12.0), - (-10.0, -25.0, 10.0), -] -ROBOT_RADIUS = 3.0 -LIFT = 1.0 - - -def _drape(x, y): - return terrain_height(x, y) + LIFT - - -# ── GRL-SNAM plan: iterate the surrogate rollout to a path of (x,y) ────────── -def plan_path(): - f64 = torch.float64 - o = torch.tensor([[START[0], START[1]]], dtype=f64) - v = torch.zeros(1, 2, dtype=f64) - goal = torch.tensor([[GOAL[0], GOAL[1]]], dtype=f64) - C = torch.tensor([[[ox, oy] for ox, oy, _ in OBSTACLES]], dtype=f64) - R = torch.tensor([[r for _, _, r in OBSTACLES]], dtype=f64) - mask = torch.ones(1, len(OBSTACLES), dtype=torch.bool) - rr = torch.full((1,), ROBOT_RADIUS, dtype=f64) # (B,) tensor, not a scalar - # known-good regime (goal spring / damping / barrier activation) for this scale - alphas = torch.full((1, len(OBSTACLES)), 3.0, dtype=f64) - beta = torch.full((1,), 2.0, dtype=f64) - gamma = torch.full((1,), 4.0, dtype=f64) - d_hat = torch.full((1,), 16.0, dtype=f64) - dt = torch.full((1,), 0.06, dtype=f64) - H1 = torch.tensor([1], dtype=torch.long) - path = [START] - min_clear = float("inf") - for _ in range(9000): - o, v, clr = integrate_surrogate_v2( - o, - v, - goal, - C, - R, - mask, - alphas, - beta, - gamma, - d_hat, - dt, - H1, - robot_radius=rr, - margin_factor=0.5, - ) - min_clear = min(min_clear, float(clr[0])) - path.append((float(o[0, 0]), float(o[0, 1]))) - if torch.linalg.norm(o - goal).item() < 1.5: - break - return path, min_clear - - -_path2d, _clear = plan_path() -print( - "grl_snam planner: %d-waypoint path, min obstacle clearance %.1f — building scene..." - % (len(_path2d), _clear) -) - -# ── build the live scene ───────────────────────────────────────────────────── -_app = vrhost.app() -_lab = Lab(app=_app, scene=vrhost.scene()) -_lab.add_terrain(_terrain_heights(), bounds=TERRAIN_BOUNDS, color=(0.32, 0.40, 0.27)) - -# obstacles as translucent red pillars (children of the terrain, so they align) -from vtkmodules.vtkFiltersSources import vtkCylinderSource # noqa: E402 -from vtkmodules.vtkFiltersGeneral import vtkTransformPolyDataFilter # noqa: E402 -from vtkmodules.vtkCommonTransforms import vtkTransform # noqa: E402 -from vtkmodules.vtkRenderingCore import vtkActor, vtkPolyDataMapper # noqa: E402 - - -def _add_obstacle(name, x, y, r, height=40.0): - cyl = vtkCylinderSource() - cyl.SetRadius(r) - cyl.SetHeight(height) - cyl.SetResolution(28) - tr = vtkTransform() # cylinder is Y-up; stand it up along Z at (x,y) - tr.Translate(x, y, terrain_height(x, y) + height * 0.5) - tr.RotateX(90.0) - tf = vtkTransformPolyDataFilter() - tf.SetTransform(tr) - tf.SetInputConnection(cyl.GetOutputPort()) - tf.Update() - m = vtkPolyDataMapper() - m.SetInputData(tf.GetOutput()) - m.SetStatic(1) - a = vtkActor() - a.SetMapper(m) - a.GetProperty().SetColor(0.85, 0.20, 0.18) - a.GetProperty().SetOpacity(0.45) - b = tf.GetOutput().GetBounds() - _lab.add_prop(name, a, (b[0], b[2], b[4], b[1], b[3], b[5]), parent="terrain") - - -for _i, (_ox, _oy, _r) in enumerate(OBSTACLES): - _add_obstacle("obstacle%d" % _i, _ox, _oy, _r) - -# start (green) + goal (gold) markers, the planned path (yellow), and the agent -_lab.add_markers("start", [(*START, _drape(*START))], color=(0.15, 0.85, 0.25)) -_lab.add_markers("goal", [(*GOAL, _drape(*GOAL))], color=(0.95, 0.80, 0.10)) -_lab.add_path("plan", [(x, y, _drape(x, y)) for x, y in _path2d], color=(0.95, 0.85, 0.20)) -_S = 5.0 -_lab.add_mesh( - "agent0", - [0, 0, _S, -_S, -_S, 0, _S, -_S, 0, 0, _S, 0], - [0, 1, 2, 0, 2, 3, 0, 3, 1, 1, 3, 2], - color=(0.95, 0.15, 0.15), -) - -# ── chase camera along the planned path ────────────────────────────────────── -_CAM = "volrover3.camera" -_chase = ChaseCamera(back=55.0, height=42.0, look_up=3.0, up=(0.0, 0.0, 1.0)) -_SPEED = 28.0 # waypoints per second along the plan - - -def _agent_at(t): - f = t * _SPEED - i = min(int(f), len(_path2d) - 1) - j = min(i + 1, len(_path2d) - 1) - a = f - i - x = _path2d[i][0] * (1 - a) + _path2d[j][0] * a - y = _path2d[i][1] * (1 - a) + _path2d[j][1] * a - return (x, y, _drape(x, y)) - - -def _cset(k, val): - pycvc.state_set(_app, _CAM + "." + k, "%.6f" % float(val)) - - -def _drive(eye, tgt, up): - vx, vy, vz = (tgt[i] - eye[i] for i in range(3)) - m = math.sqrt(vx * vx + vy * vy + vz * vz) or 1.0 - _cset("position.x", eye[0]) - _cset("position.y", eye[1]) - _cset("position.z", eye[2]) - _cset("view_direction.x", vx / m) - _cset("view_direction.y", vy / m) - _cset("view_direction.z", vz / m) - _cset("up_vector.x", up[0]) - _cset("up_vector.y", up[1]) - _cset("up_vector.z", up[2]) - _cset("fov", 55.0) - - -_DT = 1.0 / 30.0 -_t = 0.0 -for _k in range(30): # prime the camera - _e, _g, _u = _chase.update(_agent_at(_k * _DT), _DT) -_drive(_e, _g, _u) -_lab.move("agent0", *_agent_at(0.0)) -_lab.pump() -print( - "grl_snam planner demo: agent drives the planned path A->B around obstacles. Jobs tab -> Stop." -) - -_T_END = (len(_path2d) - 1) / _SPEED - - -def step(dt): - global _t - _t += dt - if _t > _T_END + 1.5: # loop - _t = 0.0 - pos = _agent_at(_t) - _lab.move("agent0", *pos) - _drive(*_chase.update(pos, dt)) - _lab.pump() - - -# ── CONTINUOUS-LEARNING note ───────────────────────────────────────────────── -# With a trained policy checkpoint you replace the hand-set coefficients with -# per-step predictions and adapt them ONLINE (no retraining): -# from grl_snam import CoefEnergyNet -# from grl_snam.adaptation import HistSecantController # rank-1 secant Jacobian -# Each step: predict (alphas,beta,gamma) from the local obstacle features, form an -# observable [clearance, goal-distance, speed], and let HistSecantController nudge -# the coefficients toward (clear, closing, moving) — the path then bends in real -# time. That online correction IS the "continuous learning after training." diff --git a/examples/volrover_lab_animated.py b/examples/volrover_lab_animated.py deleted file mode 100644 index 2e81d0f..0000000 --- a/examples/volrover_lab_animated.py +++ /dev/null @@ -1,87 +0,0 @@ -# examples/volrover_lab_animated.py — the GRL-SNAM agent WALKS its track in a -# LIVE volrover3 window, filmed by a third-person CHASE CAMERA. -# -# Load it from volrover3's Python Console -> Jobs tab -> "Load Script..." -> -# "Run as Job". The scheduler calls step(dt) each tick; the agent advances by -# MUTATING its node transform (no destroy/recreate), and a position-driven chase -# camera follows it. Pause / stop from the Jobs tab. -# -# The camera is driven ONLY by the stream of agent positions (grl_snam_lab.camera -# .ChaseCamera) — it estimates the heading from a smoothed velocity and damps the -# pose, so it works unchanged when the LIVE GRL-SNAM planner supplies real (noisy) -# positions instead of this demo path. -# -# LAYERING: volrover3 does NOT depend on GRL-SNAM. This runs under volrover3's -# GENERIC job runner only when GRL-SNAM is installed. `vrhost` + `pycvc` are -# provided by the volrover3 host at runtime. - -import math - -from grl_snam_lab.camera import ChaseCamera -from grl_snam_lab.demo import agent_position, demo_scene -from grl_snam_lab.lab import Lab - -try: - import pycvc - import vrhost # host-injected inside volrover3; a soft dependency -except ImportError as exc: # pragma: no cover - only meaningful inside volrover3 - raise RuntimeError( - "volrover_lab_animated: `vrhost`/`pycvc` not found — load this INSIDE " - "volrover3's embedded Python (Jobs tab -> Load Script -> Run as Job)." - ) from exc - -_app = vrhost.app() -_lab = demo_scene(Lab(app=_app, scene=vrhost.scene())) -_lab.pump() - -# ── third-person chase camera → volrover3's live camera (its state tree) ────── -_CAM = "volrover3.camera" -_chase = ChaseCamera(back=55.0, height=40.0, look_up=3.0, up=(0.0, 0.0, 1.0)) - - -def _cset(k, v): - pycvc.state_set(_app, _CAM + "." + k, "%.6f" % float(v)) - - -def _drive_camera(eye, tgt, up): - vx, vy, vz = (tgt[i] - eye[i] for i in range(3)) - m = math.sqrt(vx * vx + vy * vy + vz * vz) or 1.0 - _cset("position.x", eye[0]) - _cset("position.y", eye[1]) - _cset("position.z", eye[2]) - _cset("view_direction.x", vx / m) - _cset("view_direction.y", vy / m) - _cset("view_direction.z", vz / m) - _cset("up_vector.x", up[0]) - _cset("up_vector.y", up[1]) - _cset("up_vector.z", up[2]) - _cset("fov", 55.0) - - -# Prime the camera over the first ~1s of motion so it starts settled (no swing). -_DT = 1.0 / 30.0 -_t = 0.0 -for _i in range(30): - _e, _g, _u = _chase.update(agent_position(_i * _DT), _DT) -_drive_camera(_e, _g, _u) -_t = 30 * _DT -_lab.move("agent0", *agent_position(_t)) -_lab.pump() - -print( - "grl_snam_lab: animated chase-cam job loaded — the red agent walks the track, " - "the camera follows. Pause/stop from the Jobs tab." -) - - -def step(dt): - """Scheduler tick (dt = seconds since last tick): advance the agent and let the - position-driven chase camera follow. Only the agent0 node's transform and the - camera state change — nothing is added or removed.""" - global _t - _t += dt - pos = agent_position(_t) - _lab.move("agent0", *pos) - eye, tgt, up = _chase.update(pos, dt) # camera reads ONLY the position stream - _drive_camera(eye, tgt, up) - _lab.pump() diff --git a/examples/volrover_lab_demo.py b/examples/volrover_lab_demo.py deleted file mode 100644 index d839b8f..0000000 --- a/examples/volrover_lab_demo.py +++ /dev/null @@ -1,17 +0,0 @@ -# examples/volrover_lab_demo.py — run the GRL-SNAM lab in a LIVE volrover3. -# -# A thin wrapper over grl_snam_lab.run_in_volrover() (the real entry point). Run -# it INSIDE volrover3's embedded Python console: -# -# * REPL tab: exec(open(".../examples/volrover_lab_demo.py").read()) -# * or simply, in the REPL: import grl_snam_lab; grl_snam_lab.run_in_volrover() -# * or Jobs tab -> "Load Script..." and pick this file. -# -# LAYERING: volrover3 does NOT depend on GRL-SNAM. This lives here and is run by -# volrover3's GENERIC script runner only when GRL-SNAM is installed in the -# interpreter's environment. `vrhost` is injected by the volrover3 host at runtime -# (a soft dependency); outside volrover3 run_in_volrover() explains that. - -from grl_snam_lab import run_in_volrover - -run_in_volrover() diff --git a/grl_snam/cli.py b/grl_snam/cli.py new file mode 100644 index 0000000..aadea64 --- /dev/null +++ b/grl_snam/cli.py @@ -0,0 +1,234 @@ +"""``grl-snam`` — the unified command-line interface. + +One entry point for every workflow, from the world model to a running demo: + +\b + grl-snam selftest numerical correctness check (no data) + grl-snam obstacles BUNDLE world model -> circular obstacles (.npz) + grl-snam build-sdf BUNDLE world model -> navigation SDF (.npz) + grl-snam train SDF.npz self-supervised SDF-coefficient training + grl-snam capture drive|multigoal ... drive the policy -> mp4 (offscreen, HUD) + grl-snam pipeline BUNDLE world model -> SDF -> train -> video (all) + grl-snam demo NAME run a demo live inside VolRover3 + grl-snam lab-demo [PNG] standalone lab visualization + grl-snam eval [args...] CoefEnergyNet visual eval (legacy trainer) + grl-snam train-coef [args...] CoefEnergyNet dataset training (legacy) + +The heavy work lives in :mod:`grl_snam.tools`; the volrover demos in +:mod:`grl_snam.demos`. Environment-specific deps (``pycvc_gl``, VTK, ``vrhost``) load +lazily, so ``grl-snam --help`` works anywhere. +""" + +from __future__ import annotations + +import importlib +import os +import sys + +import click + +from . import __version__ + + +@click.group(help=__doc__, context_settings=dict(help_option_names=["-h", "--help"])) +@click.version_option(version=__version__, prog_name="grl-snam") +def main() -> None: + pass + + +# ── correctness + standalone viz ───────────────────────────────────────────── +@main.command() +def selftest() -> None: + """Exercise the surrogate navigation dynamics (goal-seeking / finiteness / determinism).""" + from .selftest import main as _selftest + + raise SystemExit(_selftest()) + + +@main.command("lab-demo") +@click.argument("png", required=False) +def lab_demo(png: str | None) -> None: + """Standalone lab demo (terrain + agent track + marker). PNG => offscreen snapshot.""" + from .demos import lab + + raise SystemExit(lab.main([png] if png else [])) + + +# ── world model -> obstacles / SDF ─────────────────────────────────────────── +@main.command() +@click.argument("bundle", type=click.Path(exists=True, file_okay=False)) +@click.option("-o", "--out", default=None, help="output .npz (default /obstacles.npz)") +@click.option("--grid", default=512, show_default=True, help="occupancy raster resolution") +@click.option( + "--block", default=2, show_default=True, help="obstacle coarsening (finer=smaller circles)" +) +def obstacles(bundle: str, out: str | None, grid: int, block: int) -> None: + """World model (terrain + buildings) -> circular obstacle set (.npz).""" + from .tools import obstacles as _obs + + _obs.extract_to_npz(bundle, out, grid=grid, block=block) + + +@main.command("build-sdf") +@click.argument("bundle", type=click.Path(exists=True, file_okay=False)) +@click.option("--source", type=click.Choice(["edt", "cvc"]), default="edt", show_default=True) +@click.option( + "--region", default=430.0, show_default=True, help="working-region half-extent (world)" +) +@click.option("--grid", default=512, show_default=True, help="2-D field resolution") +@click.option("-o", "--out", default=None, help="output .npz (default /nav_sdf.npz)") +def build_sdf(bundle: str, source: str, region: float, grid: int, out: str | None) -> None: + """World model -> navigation SDF (.npz). --source cvc uses CVC's mesh-exact 3-D SDF.""" + from .tools import sdf as _sdf + + _sdf.build(bundle, source=source, region=region, grid=grid, out=out) + + +# ── training ───────────────────────────────────────────────────────────────── +@main.command() +@click.argument("sdf_npz", type=click.Path(exists=True, dir_okay=False)) +@click.option("-o", "--out", default="coef_sdf.pt", show_default=True) +@click.option("--steps", default=1500, show_default=True) +@click.option("--batch", default=128, show_default=True) +@click.option("--lr", default=3e-4, show_default=True) +@click.option("--threads", default=6, show_default=True) +@click.option("--seed", default=0, show_default=True) +def train( + sdf_npz: str, out: str, steps: int, batch: int, lr: float, threads: int, seed: int +) -> None: + """Self-supervised training of the SDF navigation coefficients from a nav_sdf.npz.""" + from .tools import train as _train + + _train.train_sdf(sdf_npz, out, steps=steps, batch=batch, lr=lr, threads=threads, seed=seed) + + +# ── drive + render ─────────────────────────────────────────────────────────── +@main.group() +def capture() -> None: + """Drive the learned policy and render an mp4 OFFSCREEN (no window), with a live HUD.""" + + +@capture.command("drive") +@click.argument("bundle", type=click.Path(exists=True, file_okay=False)) +@click.argument("checkpoint", type=click.Path(exists=True, dir_okay=False)) +@click.option("--start", nargs=2, type=float, required=True) +@click.option("--goal", nargs=2, type=float, required=True) +@click.option( + "--sdf", "sdf_npz", default=None, help="prebuilt nav_sdf.npz (else built from occupancy)" +) +@click.option("-o", "--out", default="drive.mp4", show_default=True) +@click.option("--minutes", default=1.0, show_default=True) +@click.option("--no-hud", is_flag=True, help="disable the metrics HUD overlay") +def capture_drive_cmd(bundle, checkpoint, start, goal, sdf_npz, out, minutes, no_hud) -> None: + """A single fixed A->B run.""" + from .tools import capture as _cap + + _cap.capture_drive( + bundle, checkpoint, start, goal, out, sdf_npz=sdf_npz, minutes=minutes, hud=not no_hud + ) + + +@capture.command("multigoal") +@click.argument("bundle", type=click.Path(exists=True, file_okay=False)) +@click.argument("checkpoint", type=click.Path(exists=True, dir_okay=False)) +@click.option("--sdf", "sdf_npz", default=None, help="prebuilt nav_sdf.npz") +@click.option("-o", "--out", default="multigoal.mp4", show_default=True) +@click.option("--minutes", default=3.0, show_default=True) +@click.option("--no-hud", is_flag=True, help="disable the metrics HUD overlay") +def capture_multigoal_cmd(bundle, checkpoint, sdf_npz, out, minutes, no_hud) -> None: + """The dynamic multi-goal free-drive (goals re-targeted live, drone chase cam).""" + from .tools import capture as _cap + + _cap.capture_multigoal( + bundle, checkpoint, out, sdf_npz=sdf_npz, minutes=minutes, hud=not no_hud + ) + + +# ── full pipeline ──────────────────────────────────────────────────────────── +@main.command() +@click.argument("bundle", type=click.Path(exists=True, file_okay=False)) +@click.option("--source", type=click.Choice(["edt", "cvc"]), default="edt", show_default=True) +@click.option("--steps", default=1500, show_default=True, help="training steps") +@click.option("--minutes", default=3.0, show_default=True, help="demo video length") +@click.option( + "--drive", type=click.Choice(["multigoal", "single"]), default="multigoal", show_default=True +) +@click.option("--out-dir", default=None, help="where to write nav_sdf.npz / checkpoint / video") +def pipeline(bundle, source, steps, minutes, drive, out_dir) -> None: + """Run the WHOLE pipeline: world model -> SDF -> train -> demo video.""" + from .tools import pipeline as _pipe + + res = _pipe.run( + bundle, source=source, steps=steps, minutes=minutes, drive=drive, out_dir=out_dir + ) + click.echo("artifacts: " + ", ".join(f"{k}={v}" for k, v in res.items())) + + +# ── live-in-VolRover demos ─────────────────────────────────────────────────── +@main.command() +@click.argument("name", required=False) +@click.option("--bundle", default=None, help="scene bundle (sets GRL_SNAM_SCENE_BUNDLE)") +@click.option("--checkpoint", default=None, help="trained .pt (sets GRL_SNAM_CHECKPOINT)") +@click.option("--list", "list_", is_flag=True, help="list available demos") +def demo(name: str | None, bundle: str | None, checkpoint: str | None, list_: bool) -> None: + """Run a demo live inside VolRover3 (shells out to `volrover3 --run-job`). + + Set VOLROVER3_BIN if `volrover3` is not on PATH. NAME is one of the registered + demos (see `--list`).""" + from . import demos + + if list_ or not name: + click.echo("available demos (grl-snam demo NAME):") + for key, desc in demos.registry().items(): + click.echo(f" {key:18s} {desc}") + return + path = demos.demo_path(name) + if path is None: + raise click.ClickException(f"unknown demo {name!r}; try `grl-snam demo --list`") + env = os.environ.copy() + if bundle: + env["GRL_SNAM_SCENE_BUNDLE"] = bundle + if checkpoint: + env["GRL_SNAM_CHECKPOINT"] = checkpoint + binary = os.environ.get("VOLROVER3_BIN", "volrover3") + click.echo(f"launching: {binary} --run-job {path}") + import subprocess + + raise SystemExit(subprocess.call([binary, "--run-job", path], env=env)) + + +# ── legacy CoefEnergyNet trainer/evaluator (kept; future real-time HUD source) ─ +def _forward(module_name: str, args) -> None: + mod = importlib.import_module(module_name) + old = sys.argv + sys.argv = [module_name, *args] + try: + mod.main() + finally: + sys.argv = old + + +@main.command("eval", context_settings=dict(ignore_unknown_options=True)) +@click.argument("args", nargs=-1, type=click.UNPROCESSED) +def eval_cmd(args) -> None: + """CoefEnergyNet visual eval (online correction/finetune, GIF/MP4). Forwards args to eval_coef_energy.""" + _forward("eval_coef_energy", args) + + +@main.command("train-coef", context_settings=dict(ignore_unknown_options=True)) +@click.argument("args", nargs=-1, type=click.UNPROCESSED) +def train_coef_cmd(args) -> None: + """CoefEnergyNet dataset training. Forwards args to train_coef_energy.""" + _forward("train_coef_energy", args) + + +@main.command("train-geo", context_settings=dict(ignore_unknown_options=True)) +@click.argument("args", nargs=-1, type=click.UNPROCESSED) +def train_geo_cmd(args) -> None: + """Train a CoefEnergyNet directly on a scene's obstacles.npz (circle-obstacle track). + Forwards args to scripts.train_on_geometry.""" + _forward("scripts.train_on_geometry", args) + + +if __name__ == "__main__": + main() diff --git a/grl_snam/demos/__init__.py b/grl_snam/demos/__init__.py new file mode 100644 index 0000000..d95ec98 --- /dev/null +++ b/grl_snam/demos/__init__.py @@ -0,0 +1,46 @@ +"""grl_snam.demos — the live VolRover3 demos, in the package (previously ``examples/``). + +Each demo is a module defining ``step(dt)`` (and usually ``setup()``); they run inside +VolRover3 via ``grl-snam demo `` (which shells out to ``volrover3 --run-job``) or +the app's Jobs tab -> Load Script. Host-specific deps (``pycvc``, ``vrhost``, +``pycvc_gl``) import lazily, so importing this package never needs the running host. +""" + +from __future__ import annotations + +import importlib.util + +_DEMOS = { + "austin-freedrive": ( + "grl_snam.demos.austin_freedrive", + "end-to-end learned SDF free-drive on real Austin (no route)", + ), + "austin-learned": ( + "grl_snam.demos.austin_learned", + "stagewise: A* route spine + learned SDF local control", + ), + "austin-planner": ( + "grl_snam.demos.austin_planner", + "surrogate planner in its native sparse-obstacle regime", + ), + "austin-patrol": ( + "grl_snam.demos.austin_patrol", + "grounded A* street patrol (non-learned baseline)", + ), + "lab": ("grl_snam.demos.lab", "analytic terrain + an agent walking a draped loop"), +} + + +def registry() -> dict[str, str]: + """``{demo_name: description}`` for every registered demo.""" + return {k: v[1] for k, v in _DEMOS.items()} + + +def demo_path(name: str) -> str | None: + """Absolute file path of a demo module (for ``volrover3 --run-job``), or None if + unknown. Resolves without importing the module (so no host deps are needed).""" + entry = _DEMOS.get(name) + if entry is None: + return None + spec = importlib.util.find_spec(entry[0]) + return spec.origin if spec and spec.origin else None diff --git a/grl_snam/demos/_common.py b/grl_snam/demos/_common.py new file mode 100644 index 0000000..186253b --- /dev/null +++ b/grl_snam/demos/_common.py @@ -0,0 +1,157 @@ +"""Shared scaffolding for the live VolRover3 demos. + +The demo modules in this package are loaded two ways: imported as part of the +``grl_snam`` package (so the CLI can find them), and exec'd as standalone files by +``volrover3 --run-job``. They therefore use absolute imports and defer every +host-specific import (``pycvc``, ``vrhost``, ``pycvc_gl``) to run time, so importing +the package never needs the compiled bindings or the running host. + +Common pieces: a host guard, the vehicle box mesh, a camera driver that writes the +VolRover3 camera through the state tree, and a metrics publisher that pushes the +navigator's live :class:`~grl_snam.metrics.NavMetrics` into the state tree (the hook +a HUD panel reads) and prints a compact line. +""" + +from __future__ import annotations + +import math + +from grl_snam.metrics import NavMetrics, hud_lines + + +def require_host(): + """Return ``(pycvc, vrhost)`` or raise a clear error outside VolRover3.""" + try: + import pycvc + import vrhost + except ImportError as exc: # pragma: no cover - only meaningful inside volrover3 + raise RuntimeError( + "this demo must run INSIDE volrover3's embedded Python " + "(Python Console -> Jobs tab -> Load Script -> Run as Job, or " + "`volrover3 --run-job `)." + ) from exc + return pycvc, vrhost + + +def vehicle_box_mesh(length=4.6, width=2.0, height=1.6): + """Flat ``(verts, tris)`` for a simple vehicle box centred at the origin, base at z=0.""" + hx, hy = length / 2, width / 2 + v = [ + -hx, + -hy, + 0, + hx, + -hy, + 0, + hx, + hy, + 0, + -hx, + hy, + 0, + -hx, + -hy, + height, + hx, + -hy, + height, + hx, + hy, + height, + -hx, + hy, + height, + ] + t = [ + 0, + 1, + 2, + 0, + 2, + 3, + 4, + 6, + 5, + 4, + 7, + 6, + 1, + 2, + 6, + 1, + 6, + 5, + 0, + 7, + 4, + 0, + 3, + 7, + 3, + 2, + 6, + 3, + 6, + 7, + 0, + 5, + 1, + 0, + 4, + 5, + ] + return v, t + + +class CameraDriver: + """Write the VolRover3 camera each frame through the state tree (`volrover3.camera.*`).""" + + def __init__(self, app, pycvc, path="volrover3.camera", fov=60.0): + self._app = app + self._pycvc = pycvc + self._path = path + self._fov = fov + + def _set(self, key, val): + self._pycvc.state_set(self._app, f"{self._path}.{key}", "%.6f" % float(val)) + + def look(self, eye, target, up=(0.0, 0.0, 1.0)): + vx, vy, vz = (target[i] - eye[i] for i in range(3)) + mm = math.sqrt(vx * vx + vy * vy + vz * vz) or 1.0 + self._set("position.x", eye[0]) + self._set("position.y", eye[1]) + self._set("position.z", eye[2]) + self._set("view_direction.x", vx / mm) + self._set("view_direction.y", vy / mm) + self._set("view_direction.z", vz / mm) + self._set("up_vector.x", up[0]) + self._set("up_vector.y", up[1]) + self._set("up_vector.z", up[2]) + self._set("fov", self._fov) + + +class MetricsPublisher: + """Publish live :class:`NavMetrics` into the state tree (for a HUD panel) + console. + + Writes each numeric field to ``grl_snam.metrics.`` so a VolRover3 HUD/overlay + can subscribe, and prints a compact multi-line read-out every ``print_every`` steps. + """ + + def __init__(self, app, pycvc, base="grl_snam.metrics", print_every=45): + self._app = app + self._pycvc = pycvc + self._base = base + self._print_every = print_every + self._n = 0 + + def publish(self, m: NavMetrics) -> None: + self._n += 1 + for k, v in m.as_dict().items(): + if isinstance(v, bool): + v = int(v) + if isinstance(v, int | float): + self._pycvc.state_set(self._app, f"{self._base}.{k}", "%.4f" % float(v)) + else: + self._pycvc.state_set(self._app, f"{self._base}.{k}", str(v)) + if self._print_every and self._n % self._print_every == 1: + print(" | ".join(hud_lines(m)), flush=True) diff --git a/grl_snam/demos/austin_freedrive.py b/grl_snam/demos/austin_freedrive.py new file mode 100644 index 0000000..9b8a508 --- /dev/null +++ b/grl_snam/demos/austin_freedrive.py @@ -0,0 +1,118 @@ +"""Live demo: the learned SDF navigator finding its OWN way A->B across real Austin, +TX in a running VolRover3 window — no precomputed route. The trained policy reacts to +the buildings through the SDF (a moving carrot toward the goal + the wall barrier), +with a bug-style wall-follow escape for potential-field dead-ends. Live metrics (the +network's coefficients, clearance, mode) publish to the state tree + console every step. + +Prep once, then run inside VolRover3 (Jobs tab -> Load Script, or `grl-snam demo +austin-freedrive`): + grl-snam build-sdf + grl-snam train /nav_sdf.npz -o checkpoints/coef_sdf.pt + +Env: GRL_SNAM_SCENE_BUNDLE (bundle dir), GRL_SNAM_CHECKPOINT (.pt), GRL_SNAM_SDF +(prebuilt nav_sdf.npz; else built from occupancy), GRL_SNAM_START / GRL_SNAM_GOAL +("x,y"; defaults are a navigable pair on austin_south). +""" + +from __future__ import annotations + +import os + +import numpy as np +import torch + +import sdf_nav +from grl_snam.demos._common import CameraDriver, MetricsPublisher, require_host, vehicle_box_mesh +from grl_snam.nav import SdfNavigator + +_S: dict = {} + + +def _xy(env, default): + v = os.environ.get(env) + if not v: + return default + a, b = v.split(",") + return (float(a), float(b)) + + +def setup() -> None: + pycvc, vrhost = require_host() + from pycvc_gl.camera import ChaseCamera + from pycvc_gl.lab import Lab + from pycvc_gl.scenes import building_occupancy, load_geometry_bundle, terrain_grid + from pycvc_gl.vehicle import VehiclePose + + bundle = os.environ.get("GRL_SNAM_SCENE_BUNDLE", os.path.expanduser("~/scenes/austin_south")) + ckpt = os.environ.get("GRL_SNAM_CHECKPOINT", "checkpoints/coef_sdf.pt") + torch.set_num_threads(2) + ck = torch.load(ckpt, map_location="cpu") + meta = ck["meta"] + model = sdf_nav.CoefMLP() + model.load_state_dict(ck["model_state_dict"]) + model.eval() + + app = vrhost.app() + lab = Lab(app=app, scene=vrhost.scene()) + lab.set_axis_visible(False) + sample = load_geometry_bundle(lab, bundle) + bounds = terrain_grid(os.path.join(bundle, "terrain.json"))[1] + sdf_npz = os.environ.get("GRL_SNAM_SDF", "") + if sdf_npz and os.path.exists(sdf_npz): + d = np.load(sdf_npz) + phi, nxg, nyg = d["phi"], d["normal_x"], d["normal_y"] + else: + print("grl_snam austin_freedrive: rasterizing occupancy + building SDF...", flush=True) + occ = building_occupancy( + os.path.join(bundle, "buildings.glb"), bounds, 512, 512, inflate_m=0.0 + ) + phi, nxg, nyg = sdf_nav.build_sdf(occ, bounds, float(meta["scale"])) + field = sdf_nav.SDFField( + phi, nxg, nyg, bounds, meta["center"], float(meta["scale"]), device="cpu" + ) + + start = _xy("GRL_SNAM_START", (-361.0, 114.0)) + goal = _xy("GRL_SNAM_GOAL", (185.0, 50.0)) + nav = SdfNavigator(field, model, meta, reach_tol=0.5).start(start, goal) + + lab.add_markers("start", [(start[0], start[1], sample(*start) + 1.0)], color=(0.15, 0.85, 0.25)) + lab.add_markers("goal", [(goal[0], goal[1], sample(*goal) + 1.0)], color=(0.95, 0.80, 0.10)) + vv, vt = vehicle_box_mesh() + lab.add_mesh("agent0", vv, vt, color=(0.90, 0.12, 0.12)) + vpose = VehiclePose(sample, lift=0.25) + chase = ChaseCamera(back=34.0, height=13.0, look_up=2.5, up=(0.0, 0.0, 1.0)) + cam = CameraDriver(app, pycvc) + metrics = MetricsPublisher(app, pycvc) + + lab.node("agent0").setTransform(vpose.update(start[0], start[1], 1.0 / 30.0)) + lab.pump() + print( + "grl_snam austin_freedrive: the learned SDF policy is finding its own way " + "%s -> %s across Austin (no route)." % (start, goal), + flush=True, + ) + _S.update( + lab=lab, + nav=nav, + sample=sample, + vpose=vpose, + chase=chase, + cam=cam, + metrics=metrics, + steps_per_frame=3, + ) + + +def step(dt: float) -> None: + if not _S: + setup() + nav = _S["nav"] + m = None + for _ in range(_S["steps_per_frame"]): + m = nav.step() + x, y = m.x, m.y + _S["lab"].node("agent0").setTransform(_S["vpose"].update(x, y, dt)) + eye, tgt, up = _S["chase"].update((x, y, _S["sample"](x, y)), dt) + _S["cam"].look(eye, tgt, up) + _S["metrics"].publish(m) + _S["lab"].pump() diff --git a/grl_snam/demos/austin_learned.py b/grl_snam/demos/austin_learned.py new file mode 100644 index 0000000..d4a2771 --- /dev/null +++ b/grl_snam/demos/austin_learned.py @@ -0,0 +1,212 @@ +"""Live demo: the learned SDF navigator driving real Austin with an A* route spine +(GRL-SNAM's stagewise decomposition). An A*/occupancy route supplies global topology; +the trained SDF surrogate drives locally within the street corridor, accepted while it +stays out of footprints, else nudged along the clean spine — robust for adversarial +start/goal pairs where a pure potential field would stall. Live metrics publish each +step. Prep + env are the same as ``austin_freedrive`` (build-sdf + train). +""" + +from __future__ import annotations + +import os + +import numpy as np +import torch + +import sdf_nav +from grl_snam.demos._common import CameraDriver, MetricsPublisher, require_host, vehicle_box_mesh +from grl_snam.metrics import NavMetrics + +_S: dict = {} + + +def setup() -> None: + pycvc, vrhost = require_host() + from pycvc_gl.camera import ChaseCamera + from pycvc_gl.lab import Lab + from pycvc_gl.scenes import ( + building_occupancy, + load_geometry_bundle, + plan_ground_route, + resample_polyline, + terrain_grid, + ) + from pycvc_gl.vehicle import VehiclePose + + bundle = os.environ.get("GRL_SNAM_SCENE_BUNDLE", os.path.expanduser("~/scenes/austin_south")) + ckpt = os.environ.get("GRL_SNAM_CHECKPOINT", "checkpoints/coef_sdf.pt") + torch.set_num_threads(2) + ck = torch.load(ckpt, map_location="cpu") + meta = ck["meta"] + model = sdf_nav.CoefMLP() + model.load_state_dict(ck["model_state_dict"]) + model.eval() + scale = float(meta["scale"]) + center = np.asarray(meta["center"], np.float32) + region = float(meta["region"]) + kw = dict( + rr=float(meta["rr"]), + d_hat=float(meta["d_hat"]), + dt=float(meta["dt"]), + vmax=float(meta["vmax"]), + ) + nsub = int(meta["nsub"]) + + app = vrhost.app() + lab = Lab(app=app, scene=vrhost.scene()) + lab.set_axis_visible(False) + sample = load_geometry_bundle(lab, bundle) + bounds = terrain_grid(os.path.join(bundle, "terrain.json"))[1] + mnx, mny, mxx, mxy = bounds + print("grl_snam austin_learned: occupancy + SDF + route...", flush=True) + occ_r = building_occupancy( + os.path.join(bundle, "buildings.glb"), bounds, 512, 512, inflate_m=12.0 + ) + occ0 = building_occupancy( + os.path.join(bundle, "buildings.glb"), bounds, 512, 512, inflate_m=0.0 + ) + ny, nx = occ0.shape + sdf_npz = os.environ.get("GRL_SNAM_SDF", "") + if sdf_npz and os.path.exists(sdf_npz): + d = np.load(sdf_npz) + phi, nxg, nyg = d["phi"], d["normal_x"], d["normal_y"] + else: + phi, nxg, nyg = sdf_nav.build_sdf(occ0, bounds, scale) + field = sdf_nav.SDFField(phi, nxg, nyg, bounds, center, scale, device="cpu") + + start = (-region * 0.9, -region * 0.75) + goal = (region * 0.9, region * 0.8) + route = plan_ground_route(occ_r, bounds, [start, goal], close_loop=False) + if not route or len(route) < 2: + raise RuntimeError("route planning failed for this bundle/region") + route = np.asarray(resample_polyline(route, spacing=0.12 / scale), np.float32) + lab.add_markers("start", [(start[0], start[1], sample(*start) + 1.0)], color=(0.15, 0.85, 0.25)) + lab.add_markers("goal", [(goal[0], goal[1], sample(*goal) + 1.0)], color=(0.95, 0.80, 0.10)) + lab.add_path( + "spine", [(w[0], w[1], sample(w[0], w[1]) + 0.5) for w in route], color=(0.30, 0.55, 0.95) + ) + vv, vt = vehicle_box_mesh() + lab.add_mesh("agent0", vv, vt, color=(0.90, 0.12, 0.12)) + vpose = VehiclePose(sample, lift=0.25) + chase = ChaseCamera(back=34.0, height=13.0, look_up=2.5, up=(0.0, 0.0, 1.0)) + + o = torch.from_numpy( + np.array([(route[0][0] - center[0]) * scale, (route[0][1] - center[1]) * scale], np.float32) + ).unsqueeze(0) + _S.update( + lab=lab, + model=model, + field=field, + route=route, + meta=meta, + kw=kw, + nsub=nsub, + scale=scale, + center=center, + occ0=occ0, + ny=ny, + nx=nx, + bounds=bounds, + goal=goal, + sample=sample, + vpose=vpose, + chase=chase, + cam=CameraDriver(app, pycvc), + metrics=MetricsPublisher(app, pycvc), + o=o, + v=torch.zeros(1, 2), + ri=0, + stall=0, + done=False, + corr=0.35 / scale, + look=8, + init=len(route), + ) + lab.node("agent0").setTransform(vpose.update(start[0], start[1], 1.0 / 30.0)) + lab.pump() + print("grl_snam austin_learned: learned SDF policy driving the A* street corridor.", flush=True) + + +def _n2w(on, center, scale): + return np.array([on[0] / scale + center[0], on[1] / scale + center[1]], np.float32) + + +def _in_building(xw, yw): + mnx, mny, mxx, mxy = _S["bounds"] + c = int((xw - mnx) / (mxx - mnx) * (_S["nx"] - 1)) + r = int((yw - mny) / (mxy - mny) * (_S["ny"] - 1)) + return 0 <= r < _S["ny"] and 0 <= c < _S["nx"] and bool(_S["occ0"][r, c]) + + +@torch.no_grad() +def _nav_step() -> NavMetrics: + center, scale = _S["center"], _S["scale"] + route, goal = _S["route"], _S["goal"] + al = be = ga = 0.0 + if _S["done"]: + w = route[_S["ri"]] + else: + sub = route[min(_S["ri"] + _S["look"], len(route) - 1)] + gt = torch.from_numpy( + np.array([(sub[0] - center[0]) * scale, (sub[1] - center[1]) * scale], np.float32) + ).unsqueeze(0) + alt, bet, gat = _S["model"](sdf_nav.coef_feats(_S["field"], _S["o"], gt)) + al, be, ga = float(alt.mean()), float(bet.mean()), float(gat.mean()) + o2, v2, _ = sdf_nav.sdf_rollout( + _S["field"], _S["o"], _S["v"], gt, alt, bet, gat, 1, nsub=_S["nsub"], **_S["kw"] + ) + w = _n2w(o2[0].numpy(), center, scale) + lo = max(0, _S["ri"] - 2) + seg = route[lo : min(_S["ri"] + 3 * _S["look"], len(route))] + j = lo + int(np.argmin(np.linalg.norm(seg - w, axis=1))) + in_corr = (not _in_building(float(w[0]), float(w[1]))) and np.linalg.norm( + w - route[j] + ) < _S["corr"] + if in_corr: + _S["o"], _S["v"] = o2, v2 + _S["stall"] = _S["stall"] + 1 if j <= _S["ri"] else 0 + _S["ri"] = max(_S["ri"], j) + else: + _S["stall"] += 1 + if (not in_corr) or _S["stall"] > 25: + _S["ri"] = min(_S["ri"] + 1, len(route) - 1) + w = route[_S["ri"]].copy() + _S["o"] = torch.from_numpy( + np.array([(w[0] - center[0]) * scale, (w[1] - center[1]) * scale], np.float32) + ).unsqueeze(0) + _S["v"] = torch.zeros(1, 2) + _S["stall"] = 0 + mode = "route" if in_corr else "spine" + if np.linalg.norm(w - np.asarray(goal, np.float32)) < 40.0 or _S["ri"] >= len(route) - 1: + _S["done"] = True + phi, _ = _S["field"].sample(_S["o"]) + dg = float(np.linalg.norm(np.asarray(goal, np.float32) - w)) + return NavMetrics( + step=_S["ri"], + x=float(w[0]), + y=float(w[1]), + goal_x=goal[0], + goal_y=goal[1], + goal_dist_m=dg, + clearance_m=(float(phi[0]) - _S["kw"]["rr"]) / scale, + alpha=al, + beta=be, + gamma=ga, + mode=locals().get("mode", "route"), + stall=_S["stall"], + inside_building=_in_building(float(w[0]), float(w[1])), + progress=_S["ri"] / max(1, _S["init"] - 1), + ) + + +def step(dt: float) -> None: + if not _S: + setup() + m = None + for _ in range(3): + m = _nav_step() + _S["lab"].node("agent0").setTransform(_S["vpose"].update(m.x, m.y, dt)) + eye, tgt, up = _S["chase"].update((m.x, m.y, _S["sample"](m.x, m.y)), dt) + _S["cam"].look(eye, tgt, up) + _S["metrics"].publish(m) + _S["lab"].pump() diff --git a/grl_snam/demos/austin_patrol.py b/grl_snam/demos/austin_patrol.py new file mode 100644 index 0000000..bd24b9b --- /dev/null +++ b/grl_snam/demos/austin_patrol.py @@ -0,0 +1,95 @@ +"""Live demo: a grounded A* street patrol on real Austin, TX. The city mesh is +rasterized to an occupancy grid, an A* route threads a closed patrol loop through the +free space (streets), and a vehicle drives it draped on the terrain under a low chase +camera — a non-learned baseline (and a clean showcase of the scene + routing stack). +Env: GRL_SNAM_SCENE_BUNDLE. +""" + +from __future__ import annotations + +import math +import os + +from grl_snam.demos._common import CameraDriver, require_host, vehicle_box_mesh + +_S: dict = {} + + +def setup() -> None: + pycvc, vrhost = require_host() + from pycvc_gl.camera import ChaseCamera + from pycvc_gl.lab import Lab + from pycvc_gl.scenes import ( + building_occupancy, + load_geometry_bundle, + plan_ground_route, + resample_polyline, + terrain_grid, + ) + from pycvc_gl.vehicle import VehiclePose + + bundle = os.environ.get("GRL_SNAM_SCENE_BUNDLE", os.path.expanduser("~/scenes/austin_south")) + app = vrhost.app() + lab = Lab(app=app, scene=vrhost.scene()) + lab.set_axis_visible(False) + sample = load_geometry_bundle(lab, bundle) + bounds = terrain_grid(os.path.join(bundle, "terrain.json"))[1] + print("grl_snam austin_patrol: building occupancy + A* patrol route...", flush=True) + occ = building_occupancy( + os.path.join(bundle, "buildings.glb"), bounds, 512, 512, inflate_m=12.0 + ) + r = 430.0 + waypts = [ + (r * math.cos(a), r * math.sin(a)) + for a in [0.0, math.pi / 3, 2 * math.pi / 3, math.pi, 4 * math.pi / 3, 5 * math.pi / 3] + ] + route2d = resample_polyline( + plan_ground_route(occ, bounds, waypts, close_loop=True), spacing=4.0 + ) + if len(route2d) < 2: + route2d = [ + (r * math.cos(2 * math.pi * k / 240), r * math.sin(2 * math.pi * k / 240)) + for k in range(241) + ] + print("grl_snam austin_patrol: grounded route = %d pts along the streets." % len(route2d)) + vv, vt = vehicle_box_mesh() + lab.add_mesh("agent0", vv, vt, color=(0.90, 0.12, 0.12)) + lab.add_path( + "agent0_track", [(x, y, sample(x, y) + 0.6) for x, y in route2d], color=(0.95, 0.75, 0.10) + ) + vpose = VehiclePose(sample, lift=0.25) + chase = ChaseCamera(back=34.0, height=13.0, look_up=2.5, up=(0.0, 0.0, 1.0)) + lab.node("agent0").setTransform(vpose.update(route2d[0][0], route2d[0][1], 1.0 / 30.0)) + lab.pump() + print("grl_snam austin_patrol: the vehicle patrols the streets, routed around buildings.") + _S.update( + lab=lab, + route=route2d, + sample=sample, + vpose=vpose, + chase=chase, + cam=CameraDriver(app, pycvc), + t=0.0, + pps=3.0, + ) + + +def _pose(t): + route = _S["route"] + n = len(route) + f = (t * _S["pps"]) % n + i = int(f) + j = (i + 1) % n + a = f - i + return route[i][0] * (1 - a) + route[j][0] * a, route[i][1] * (1 - a) + route[j][1] * a + + +def step(dt: float) -> None: + if not _S: + setup() + _S["t"] += dt + x, y = _pose(_S["t"]) + _S["lab"].node("agent0").setTransform(_S["vpose"].update(x, y, dt)) + eye, tgt, up = _S["chase"].update((x, y, _S["sample"](x, y)), dt) + _S["cam"].look(eye, tgt, up) + _S["lab"].pump() diff --git a/grl_snam/demos/austin_planner.py b/grl_snam/demos/austin_planner.py new file mode 100644 index 0000000..ad3e53f --- /dev/null +++ b/grl_snam/demos/austin_planner.py @@ -0,0 +1,164 @@ +"""Live demo: the GRL-SNAM surrogate planner in its native sparse-obstacle regime. + +On an analytic terrain with a handful of circular obstacles, the differentiable +surrogate rollout (``grl_snam.dynamics.integrate_surrogate_v2`` — the same integrator +downstream training runs through) is iterated to a collision-free path A->B, and an +agent drives it under a chase camera. Coefficients are hand-set to a known-good regime +here; with a trained ``CoefEnergyNet`` you would predict them per step and adapt online +with ``grl_snam.adaptation.HistSecantController``. +""" + +from __future__ import annotations + +import torch + +from grl_snam.demos._common import CameraDriver, require_host +from grl_snam.demos.lab import TERRAIN_BOUNDS, _terrain_heights, terrain_height +from grl_snam.dynamics import integrate_surrogate_v2 + +START = (-70.0, -60.0) +GOAL = (70.0, 60.0) +OBSTACLES = [(10.0, 20.0, 14.0), (-20.0, 10.0, 12.0), (25.0, -15.0, 12.0), (-10.0, -25.0, 10.0)] +ROBOT_RADIUS = 3.0 +LIFT = 1.0 +_S: dict = {} + + +def _drape(x, y): + return terrain_height(x, y) + LIFT + + +def plan_path(): + """Iterate the surrogate rollout to a collision-free path A->B; return (path, min_clear).""" + f64 = torch.float64 + o = torch.tensor([[START[0], START[1]]], dtype=f64) + v = torch.zeros(1, 2, dtype=f64) + goal = torch.tensor([[GOAL[0], GOAL[1]]], dtype=f64) + c = torch.tensor([[[ox, oy] for ox, oy, _ in OBSTACLES]], dtype=f64) + r = torch.tensor([[rr for _, _, rr in OBSTACLES]], dtype=f64) + mask = torch.ones(1, len(OBSTACLES), dtype=torch.bool) + rr = torch.full((1,), ROBOT_RADIUS, dtype=f64) + alphas = torch.full((1, len(OBSTACLES)), 3.0, dtype=f64) + beta = torch.full((1,), 2.0, dtype=f64) + gamma = torch.full((1,), 4.0, dtype=f64) + d_hat = torch.full((1,), 16.0, dtype=f64) + dt = torch.full((1,), 0.06, dtype=f64) + h1 = torch.tensor([1], dtype=torch.long) + path = [START] + min_clear = float("inf") + for _ in range(9000): + o, v, clr = integrate_surrogate_v2( + o, + v, + goal, + c, + r, + mask, + alphas, + beta, + gamma, + d_hat, + dt, + h1, + robot_radius=rr, + margin_factor=0.5, + ) + min_clear = min(min_clear, float(clr[0])) + path.append((float(o[0, 0]), float(o[0, 1]))) + if torch.linalg.norm(o - goal).item() < 1.5: + break + return path, min_clear + + +def _obstacle_actor(x, y, r, height=40.0): + from vtkmodules.vtkCommonTransforms import vtkTransform + from vtkmodules.vtkFiltersGeneral import vtkTransformPolyDataFilter + from vtkmodules.vtkFiltersSources import vtkCylinderSource + from vtkmodules.vtkRenderingCore import vtkActor, vtkPolyDataMapper + + cyl = vtkCylinderSource() + cyl.SetRadius(r) + cyl.SetHeight(height) + cyl.SetResolution(28) + tr = vtkTransform() + tr.Translate(x, y, terrain_height(x, y) + height * 0.5) + tr.RotateX(90.0) + tf = vtkTransformPolyDataFilter() + tf.SetTransform(tr) + tf.SetInputConnection(cyl.GetOutputPort()) + tf.Update() + m = vtkPolyDataMapper() + m.SetInputData(tf.GetOutput()) + m.SetStatic(1) + a = vtkActor() + a.SetMapper(m) + a.GetProperty().SetColor(0.85, 0.20, 0.18) + a.GetProperty().SetOpacity(0.45) + return a, tf.GetOutput().GetBounds() + + +def setup() -> None: + pycvc, vrhost = require_host() + from pycvc_gl.camera import ChaseCamera + from pycvc_gl.lab import Lab + + path2d, clear = plan_path() + print( + "grl_snam planner: %d-waypoint path, min clearance %.1f — building scene..." + % (len(path2d), clear), + flush=True, + ) + app = vrhost.app() + lab = Lab(app=app, scene=vrhost.scene()) + lab.add_terrain(_terrain_heights(), bounds=TERRAIN_BOUNDS, color=(0.32, 0.40, 0.27)) + for i, (ox, oy, r) in enumerate(OBSTACLES): + actor, b = _obstacle_actor(ox, oy, r) + lab.add_prop( + "obstacle%d" % i, actor, (b[0], b[2], b[4], b[1], b[3], b[5]), parent="terrain" + ) + lab.add_markers("start", [(*START, _drape(*START))], color=(0.15, 0.85, 0.25)) + lab.add_markers("goal", [(*GOAL, _drape(*GOAL))], color=(0.95, 0.80, 0.10)) + lab.add_path("plan", [(x, y, _drape(x, y)) for x, y in path2d], color=(0.95, 0.85, 0.20)) + s = 5.0 + lab.add_mesh( + "agent0", + [0, 0, s, -s, -s, 0, s, -s, 0, 0, s, 0], + [0, 1, 2, 0, 2, 3, 0, 3, 1, 1, 3, 2], + color=(0.95, 0.15, 0.15), + ) + chase = ChaseCamera(back=55.0, height=42.0, look_up=3.0, up=(0.0, 0.0, 1.0)) + lab.move("agent0", *_pose(path2d, 0.0)) + lab.pump() + print("grl_snam planner: agent drives the planned path A->B around obstacles.", flush=True) + _S.update( + lab=lab, + path=path2d, + chase=chase, + cam=CameraDriver(app, pycvc, fov=55.0), + t=0.0, + speed=28.0, + t_end=(len(path2d) - 1) / 28.0, + ) + + +def _pose(path2d, t, speed=28.0): + f = t * speed + i = min(int(f), len(path2d) - 1) + j = min(i + 1, len(path2d) - 1) + a = f - i + x = path2d[i][0] * (1 - a) + path2d[j][0] * a + y = path2d[i][1] * (1 - a) + path2d[j][1] * a + return (x, y, _drape(x, y)) + + +def step(dt: float) -> None: + if not _S: + setup() + _S["t"] += dt + if _S["t"] > _S["t_end"] + 1.5: + _S["t"] = 0.0 + pos = _pose(_S["path"], _S["t"], _S["speed"]) + _S["lab"].move("agent0", *pos) + eye, tgt, up = _S["chase"].update(pos, dt) + _S["cam"].look(eye, tgt, up) + _S["lab"].pump() diff --git a/grl_snam/demos/lab.py b/grl_snam/demos/lab.py new file mode 100644 index 0000000..297b507 --- /dev/null +++ b/grl_snam/demos/lab.py @@ -0,0 +1,152 @@ +"""The general GRL-SNAM lab demo — an analytic terrain with an agent walking a draped +loop — static, animated, or offscreen, built on the CVC graphics stack (``pycvc_gl``). + +This is the home of the former ``grl_snam_lab.demo`` after consolidating onto +``pycvc_gl`` (which supplies ``Lab``, ``terrain_mesh``, cameras, etc.). The pure +geometry (``terrain_height``, ``TRACK``, ``agent_position``) needs no bindings and is +imported by the sparse-obstacle planner demo; the scene builders import ``pycvc_gl`` +lazily so this module imports anywhere. + +Entry points: :func:`run_standalone` (window / PNG, the ``grl-snam lab-demo`` command), +:func:`run_in_volrover` (into a running VolRover3 scene), and ``step(dt)`` (the animated +job loaded with ``grl-snam demo lab``). +""" + +from __future__ import annotations + +import math + +TERRAIN_BOUNDS = (-100.0, -100.0, 100.0, 100.0) +TRACK_RADIUS = 70.0 +LOOP_SECONDS = 22.0 +AGENT_LIFT = 0.4 +_N_TRACK = 180 + + +def terrain_height(x: float, y: float) -> float: + """Analytic terrain height at world ``(x, y)`` — gentle rolling hills. Single source + of truth: the mesh samples it and the agent/track drape onto it.""" + h = 14.0 * math.exp(-((x * x + y * y) / 3000.0)) + h += 8.0 * math.exp(-(((x - 55.0) ** 2 + (y + 35.0) ** 2) / 900.0)) + h += 6.0 * math.exp(-(((x + 45.0) ** 2 + (y - 50.0) ** 2) / 1200.0)) + h += 2.5 * math.sin(x * 0.06) * math.cos(y * 0.05) + return h + + +def _terrain_heights(n: int = 56): + """Sample ``terrain_height`` on an n x n grid over TERRAIN_BOUNDS (row -> y, col -> x).""" + min_x, min_y, max_x, max_y = TERRAIN_BOUNDS + sx = (max_x - min_x) / (n - 1) + sy = (max_y - min_y) / (n - 1) + return [[terrain_height(min_x + j * sx, min_y + i * sy) for j in range(n)] for i in range(n)] + + +def _loop_point(theta: float): + r = TRACK_RADIUS + 8.0 * math.sin(3.0 * theta) + x = r * math.cos(theta) + y = r * math.sin(theta) + return (x, y, terrain_height(x, y) + AGENT_LIFT) + + +TRACK = [_loop_point(2.0 * math.pi * k / _N_TRACK) for k in range(_N_TRACK + 1)] + + +def _agent_marker_mesh(size: float = 5.0): + s = size + verts = [0.0, 0.0, s, -s, -s, 0.0, s, -s, 0.0, 0.0, s, 0.0] + tris = [0, 1, 2, 0, 2, 3, 0, 3, 1, 1, 3, 2] + return verts, tris + + +def agent_position(t: float, speed: float = 1.0): + """The agent's ``(x, y, z)`` on the demo loop at time ``t`` (draped onto the terrain).""" + theta = 2.0 * math.pi * ((t * speed / LOOP_SECONDS) % 1.0) + return _loop_point(theta) + + +def demo_scene(lab): + """Add the canonical demo — terrain + the agent's track + the agent marker — to ``lab``.""" + lab.add_terrain(_terrain_heights(), bounds=TERRAIN_BOUNDS, color=(0.42, 0.53, 0.34)) + lab.add_path("agent0_track", TRACK, color=(0.95, 0.75, 0.10)) + verts, tris = _agent_marker_mesh() + lab.add_mesh("agent0", verts, tris, color=(1.0, 0.20, 0.20)) + lab.move("agent0", *agent_position(0.0)) + lab.pump() + return lab + + +def animate_agent(lab, t: float): + """Walk the ``agent0`` marker to its position at time ``t`` (transform mutate, no rebuild).""" + lab.move("agent0", *agent_position(t)) + lab.pump() + return lab + + +def run_standalone(png: str | None = None, width: int = 1024, height: int = 768): + """Static demo in a self-owned Lab; offscreen PNG when ``png`` is given, else a window.""" + from pycvc_gl.lab import Lab + + lab = demo_scene(Lab()) + if png: + lab.render_png(png, width, height) + print(f"grl-snam lab-demo: wrote {png} ({lab.num_nodes()} nodes)") + else: + print(f"grl-snam lab-demo: showing {lab.num_nodes()} nodes — close the window to exit") + lab.show("grl-snam lab demo", width, height) + return lab + + +def run_in_volrover(): + """Build the static demo in a running VolRover3 scene (call from the REPL / job).""" + from grl_snam.demos._common import require_host + + _pycvc, vrhost = require_host() + from pycvc_gl.lab import Lab + + lab = demo_scene(Lab(app=vrhost.app(), scene=vrhost.scene())) + lab.pump() + print( + f"grl-snam lab: live VolRover3 scene now has {lab.num_nodes()} node(s) — look at the viewport." + ) + return lab + + +_S: dict = {} + + +def setup() -> None: + from grl_snam.demos._common import require_host + + _pycvc, vrhost = require_host() + from pycvc_gl.lab import Lab + + lab = demo_scene(Lab(app=vrhost.app(), scene=vrhost.scene())) + _S.update(lab=lab, t=0.0) + print("grl-snam lab: animated agent walking the draped loop.", flush=True) + + +def step(dt: float) -> None: + """Animated job: walk the agent one tick (``grl-snam demo lab``).""" + if not _S: + setup() + _S["t"] += dt + animate_agent(_S["lab"], _S["t"]) + + +def main(argv: list[str] | None = None) -> int: + """Console entry (``grl-snam lab-demo``): standalone viz. ``lab-demo out.png`` => snapshot.""" + import argparse + + ap = argparse.ArgumentParser( + prog="grl-snam lab-demo", description="GRL-SNAM lab demo (terrain + agent track + marker)." + ) + ap.add_argument( + "png", nargs="?", help="write an offscreen PNG here instead of opening a window" + ) + args = ap.parse_args(argv) + try: + run_standalone(png=args.png) + except Exception as exc: # pycvc/pycvc_gl missing, or no GL display + print(f"grl-snam lab-demo: {exc}") + return 1 + return 0 diff --git a/grl_snam/metrics.py b/grl_snam/metrics.py new file mode 100644 index 0000000..14c686d --- /dev/null +++ b/grl_snam/metrics.py @@ -0,0 +1,92 @@ +"""Live navigation metrics + a HUD overlay — a running read-out of what the learned +policy is doing while it drives. + +``NavMetrics`` is one snapshot per navigation step (position, goal distance, wall +clearance, speed, the network's predicted coefficients alpha/beta/gamma, the escape +mode, whether the agent is penetrating a building, progress). ``SdfNavigator`` +(``grl_snam.nav``) fills one in every step; the demos print/emit them live and the +video capture draws them as an on-frame HUD via ``hud_lines``. + +The point is developer/operator insight: the coefficients are the policy's actual +output, and clearance + penetration + mode expose *why* it moves the way it does — +exactly the signals a HUD (still being designed) should surface in real time. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass +class NavMetrics: + """One per-step snapshot of the navigator's state and the network's output.""" + + step: int = 0 + x: float = 0.0 # agent world position + y: float = 0.0 + goal_x: float = 0.0 # active goal world position + goal_y: float = 0.0 + goal_dist_m: float = 0.0 # world metres to the active goal + clearance_m: float = 0.0 # SDF clearance to nearest wall minus robot radius (world m) + speed_mps: float = 0.0 # world metres/second + alpha: float = 0.0 # network output: obstacle-barrier weight + beta: float = 0.0 # network output: goal-spring stiffness + gamma: float = 0.0 # network output: velocity damping + mode: str = "seek" # "seek" (head to goal) or "wall" (bug-style escape) + stall: int = 0 # steps since last progress toward the goal + inside_building: bool = False # footprint penetration this step + progress: float = 0.0 # 0..1, 1 - dist/initial_dist toward the active goal + goal_wall_align: float = 0.0 # goal-direction . wall-normal (>0 = wall between agent and goal) + goal_index: int = 0 # which goal (multi-goal drives) + reached: bool = False # active goal reached this step + + def as_dict(self) -> dict: + return self.__dict__.copy() + + +@dataclass +class NavStats: + """Rolling aggregates over a drive — for an end-of-run summary and the HUD footer.""" + + steps: int = 0 + penetration_steps: int = 0 + goals_reached: int = 0 + min_clearance_m: float = field(default=1e9) + total_path_m: float = 0.0 + _prev: tuple | None = None + + def update(self, m: NavMetrics) -> None: + self.steps += 1 + if m.inside_building: + self.penetration_steps += 1 + if m.reached: + self.goals_reached += 1 + self.min_clearance_m = min(self.min_clearance_m, m.clearance_m) + if self._prev is not None: + dx, dy = m.x - self._prev[0], m.y - self._prev[1] + self.total_path_m += (dx * dx + dy * dy) ** 0.5 + self._prev = (m.x, m.y) + + @property + def penetration_pct(self) -> float: + return 100.0 * self.penetration_steps / max(1, self.steps) + + +def hud_lines(m: NavMetrics, stats: NavStats | None = None) -> list[str]: + """Human-readable HUD lines for a metrics snapshot (used as on-frame overlay text + and for live console output). Kept short so it fits a corner of the viewport.""" + clr = f"{m.clearance_m:5.1f} m" if m.clearance_m < 999 else " open" + lines = [ + "GRL-SNAM learned SDF navigator", + f"goal {m.goal_index} dist {m.goal_dist_m:6.1f} m {int(100 * m.progress):3d}%", + f"coeffs a {m.alpha:4.2f} b {m.beta:4.2f} g {m.gamma:4.2f}", + f"clearance {clr} speed {m.speed_mps:4.1f} m/s", + f"mode {m.mode:4s} stall {m.stall:3d}" + + (" [WALL CONTACT]" if m.inside_building else ""), + ] + if stats is not None: + lines.append( + f"reached {stats.goals_reached} penetration {stats.penetration_pct:4.1f}%" + f" min clr {stats.min_clearance_m:4.1f} m" + ) + return lines diff --git a/grl_snam/nav.py b/grl_snam/nav.py new file mode 100644 index 0000000..daac0db --- /dev/null +++ b/grl_snam/nav.py @@ -0,0 +1,269 @@ +"""The shared learned-SDF navigator — one implementation of the drive loop used by +every demo and the video capture (previously copy-pasted across four scripts). + +``SdfNavigator`` wraps a trained ``sdf_nav.CoefMLP`` + an ``SdfNavigator``-owned +``SDFField`` and drives a point agent toward a goal purely by reacting to the field: +a moving carrot toward the goal, deflected by the SDF wall barrier, with a +``bug``-style wall-follow escape when a potential-field dead-end stalls progress. +``step()`` advances one step and returns a :class:`~grl_snam.metrics.NavMetrics` +snapshot, so callers get the network's live coefficients and clearance for free. + +Goals are live-retargetable (``set_goal``) — the basis of the dynamic multi-goal +free-drive. ``select_reachable_goals`` picks corner-ward goals in the SDF's *own* +free space so every leg is actually reachable. +""" + +from __future__ import annotations + +import math + +import numpy as np +import torch + +import sdf_nav + +from .metrics import NavMetrics + + +class SdfNavigator: + """Drive a point agent across a trained SDF field toward a (retargetable) goal.""" + + def __init__(self, field, model, meta, *, reach_tol: float = 0.8): + self.field = field + self.model = model + self.meta = meta + self.S = float(meta["scale"]) + self.cx, self.cy = (float(c) for c in meta["center"]) + self.kw = dict( + rr=float(meta["rr"]), + d_hat=float(meta["d_hat"]), + dt=float(meta["dt"]), + vmax=float(meta["vmax"]), + ) + self.nsub = int(meta.get("nsub", 1)) + self.dt = float(meta["dt"]) + self.rr = float(meta["rr"]) + self.reach_tol = reach_tol + self.o = torch.zeros(1, 2) + self.v = torch.zeros(1, 2) + self._gn = np.zeros(2, np.float32) + self.step_i = 0 + self.goal_index = 0 + self._reset_goal_state(1e9) + + # ── world <-> normalized ──────────────────────────────────────────────── + def w2n(self, p): + return np.array([(p[0] - self.cx) * self.S, (p[1] - self.cy) * self.S], np.float32) + + def n2w(self, on): + return np.array([on[0] / self.S + self.cx, on[1] / self.S + self.cy], np.float32) + + def pos_world(self): + return self.n2w(self.o[0].numpy()) + + # ── setup ─────────────────────────────────────────────────────────────── + def start(self, start_world, goal_world): + self.o = torch.from_numpy(self.w2n(start_world)).unsqueeze(0).float() + self.v = torch.zeros(1, 2) + self.step_i = 0 + self.set_goal(goal_world) + return self + + def set_goal(self, goal_world, *, goal_index: int | None = None): + self._gn = self.w2n(goal_world) + if goal_index is not None: + self.goal_index = goal_index + p = self.o[0].numpy() + self._reset_goal_state(float(np.linalg.norm(self._gn - p))) + + def _reset_goal_state(self, dist0): + self._best = dist0 + self._init = max(dist0, 1e-6) + self._stall = 0 + self._mode = "seek" + self._turn = 1.0 + self._dhit = 0.0 + self.reached = False + + @torch.no_grad() + def _sdf_normal(self, on): + _, nrm = self.field.sample(torch.from_numpy(on).unsqueeze(0).float()) + return nrm[0].numpy() + + # ── one drive step ────────────────────────────────────────────────────── + @torch.no_grad() + def step(self) -> NavMetrics: + """Advance one navigation step; return the metrics snapshot after it.""" + prev_world = self.pos_world() + p = self.o[0].numpy() + dg = float(np.linalg.norm(self._gn - p)) + gdir = (self._gn - p) / (dg + 1e-6) + + if dg < self._best - 1e-3: + self._best = dg + self._stall = 0 + else: + self._stall += 1 + + # enter wall-follow when a potential-field dead-end stalls progress + if self._mode == "seek" and self._stall > 70: + nrm = self._sdf_normal(p) + tang = np.array([-nrm[1], nrm[0]], np.float32) + self._turn = 1.0 if np.dot(tang, gdir) >= 0 else -1.0 + self._dhit = dg + self._mode = "wall" + self._stall = 0 + + if self._mode == "wall": + nrm = self._sdf_normal(p) + tang = self._turn * np.array([-nrm[1], nrm[0]], np.float32) + carrot = (p + (0.6 * tang + 0.4 * nrm) * 1.6).astype(np.float32) + if dg < self._dhit - 1.2 or self._stall > 240: # rounded it, or give up escaping + self._mode = "seek" + self._best = dg + self._stall = 0 + else: + carrot = (p + gdir * min(1.8, dg)).astype(np.float32) + + gt = torch.from_numpy(carrot).unsqueeze(0) + al, be, ga = self.model(sdf_nav.coef_feats(self.field, self.o, gt)) + self.o, self.v, _ = sdf_nav.sdf_rollout( + self.field, self.o, self.v, gt, al, be, ga, 1, nsub=self.nsub, **self.kw + ) + self.step_i += 1 + + # metrics at the new position + onew = self.o[0].numpy() + phi, nrm = self.field.sample(self.o) + phi = float(phi[0]) + dg_new = float(np.linalg.norm(self._gn - onew)) + world = self.n2w(onew) + step_world = float(np.linalg.norm(world - prev_world)) + gwd = self._gn - onew + gwn = gwd / (np.linalg.norm(gwd) + 1e-6) + align = float(np.dot(gwn, nrm[0].numpy())) + self.reached = dg_new < self.reach_tol + return NavMetrics( + step=self.step_i, + x=float(world[0]), + y=float(world[1]), + goal_x=float(self._gn[0] / self.S + self.cx), + goal_y=float(self._gn[1] / self.S + self.cy), + goal_dist_m=dg_new / self.S, + clearance_m=(phi - self.rr) / self.S, + speed_mps=step_world / self.dt if self.dt > 0 else 0.0, + alpha=float(al.mean()), + beta=float(be.mean()), + gamma=float(ga.mean()), + mode=self._mode, + stall=self._stall, + inside_building=phi < 0.0, + progress=max(0.0, 1.0 - dg_new / self._init), + goal_wall_align=align, + goal_index=self.goal_index, + reached=self.reached, + ) + + def drive_to_goal(self, max_steps: int = 1300, *, stop_at_reach: bool = True): + """Run steps toward the current goal until reached / stuck / ``max_steps``. + Returns ``(metrics_list, closest_idx, o_at_closest, v_at_closest)`` — the + closest-approach index truncates an orbiting tail for clean video, and its + ``(o, v)`` lets a multi-goal driver continue the next leg from that point.""" + out = [] + best = 1e9 + best_i = 0 + since_best = 0 + best_o = self.o.clone() + best_v = self.v.clone() + for _ in range(max_steps): + m = self.step() + out.append(m) + if m.goal_dist_m < best - 1e-3: + best = m.goal_dist_m + best_i = len(out) - 1 + best_o = self.o.clone() + best_v = self.v.clone() + since_best = 0 + else: + since_best += 1 + if stop_at_reach and m.reached: + best_i = len(out) - 1 + best_o = self.o.clone() + best_v = self.v.clone() + break + if since_best > 340: # stuck / orbiting — stop; caller re-targets + break + return out, best_i, best_o, best_v + + +def select_reachable_goals( + field, model, meta, *, n_corners: int = 4, clear_world: float = 16.0, reach_world: float = 70.0 +): + """Pick ``n_corners`` corner-ward goals the navigator can actually REACH, in the + SDF's own free space, chained from an open centre cell in a perimeter loop. Each + candidate is *simulated* with the real model (not just checked geometrically), so + selection and navigation agree — the fix for far-corner wandering. Returns + ``(start_world, [goal_world, ...])``.""" + S = float(meta["scale"]) + cx, cy = (float(c) for c in meta["center"]) + region = float(meta["region"]) + mnx, mny, mxx, mxy = (float(b) for b in meta["bounds"]) + + def clearance(pw): + on = np.array([(pw[0] - cx) * S, (pw[1] - cy) * S], np.float32) + d, _ = field.sample(torch.from_numpy(on).unsqueeze(0).float()) + return float(d[0]) / S + + def w2n(pw): + return np.array([(pw[0] - cx) * S, (pw[1] - cy) * S], np.float32) + + start = np.array([cx, cy], np.float32) + for rad in (0, 20, 40, 60, 80, 120): + found = False + for a in range(0, 360, 30): + pw = np.array( + [cx + math.cos(math.radians(a)) * rad, cy + math.sin(math.radians(a)) * rad], + np.float32, + ) + if clearance(pw) > clear_world: + start = pw + found = True + break + if found: + break + + def candidates(sx, sy): + diag = math.atan2(sy, sx) + out = [] + for frac in (0.62, 0.56, 0.50, 0.44, 0.38, 0.32): + for da in (0, -18, 18, -32, 32): + ang = diag + math.radians(da) + r = region * frac + gw = np.array([cx + math.cos(ang) * r, cy + math.sin(ang) * r], np.float32) + if mnx < gw[0] < mxx and mny < gw[1] < mxy and clearance(gw) > clear_world: + out.append(gw) + return out + + nav = SdfNavigator(field, model, meta) + o = torch.from_numpy(w2n(start)).unsqueeze(0).float() + v = torch.zeros(1, 2) + goals = [] + for sx, sy in [(1, 1), (1, -1), (-1, -1), (-1, 1)][:n_corners]: + chosen = None + chosen_ba = 1e9 + chosen_o, chosen_v = o, v + for gw in candidates(sx, sy)[:8]: + nav.o = o.clone() + nav.v = v.clone() + nav.set_goal(gw) + ms, _bi, bo, bv = nav.drive_to_goal(max_steps=800) + ba = float(np.linalg.norm(bo[0].numpy() - w2n(gw))) / S + if ba < chosen_ba: + chosen_ba, chosen, chosen_o, chosen_v = ba, gw, bo, bv + if ba < reach_world: + break + if chosen is None: + chosen = np.array([cx + sx * region * 0.4, cy + sy * region * 0.4], np.float32) + o, v = chosen_o, chosen_v + goals.append(chosen) + return start, goals diff --git a/grl_snam/tools/__init__.py b/grl_snam/tools/__init__.py new file mode 100644 index 0000000..f98aa03 --- /dev/null +++ b/grl_snam/tools/__init__.py @@ -0,0 +1,14 @@ +"""grl_snam.tools — the importable core of the command-line pipeline. + +Each module holds the real logic (previously buried in a ``scripts/*.py`` ``main()``) +as a plain function the Click CLI (:mod:`grl_snam.cli`) and other code can call: + +* :func:`grl_snam.tools.obstacles.extract` — geometry bundle -> circular obstacles. +* :func:`grl_snam.tools.sdf.build` — geometry bundle -> navigation SDF (``nav_sdf.npz``). +* :func:`grl_snam.tools.train.train_sdf` — self-supervised SDF-coefficient training. +* :func:`grl_snam.tools.capture` — drive the learned policy and render an mp4 (HUD). +* :func:`grl_snam.tools.pipeline.run` — the whole world-model -> train -> video pipeline. + +Heavy, environment-specific deps (``pycvc_gl``, VTK) are imported lazily inside the +functions so importing this package never requires the compiled bindings. +""" diff --git a/grl_snam/tools/capture.py b/grl_snam/tools/capture.py new file mode 100644 index 0000000..89e804b --- /dev/null +++ b/grl_snam/tools/capture.py @@ -0,0 +1,410 @@ +"""Drive the learned SDF policy on a real scene and render it to an mp4 OFFSCREEN +(no window) via VTK — the same engine cvcGL/VolRover3 use — with an on-frame **HUD** +of live metrics (the network's coefficients, wall clearance, mode, penetration). + +Two entry points, both built on the one shared :class:`grl_snam.nav.SdfNavigator`: + +* :func:`capture_drive` — a single fixed A->B run. +* :func:`capture_multigoal` — the dynamic multi-goal free-drive (goals chosen in the + SDF's free space so every leg is reachable; goals re-targeted live; drone chase cam; + each leg truncated at closest approach so no orbiting tail shows). + +Needs the volrover env (``pycvc_gl`` scene helpers + VTK), a trained checkpoint, a +scene bundle (``terrain.json`` + ``buildings.glb``), and ffmpeg. +""" + +from __future__ import annotations + +import json +import math +import os +import subprocess + +import numpy as np +import torch + +import sdf_nav + +from ..metrics import NavStats, hud_lines +from ..nav import SdfNavigator, select_reachable_goals + +_GOAL_COLORS = [(0.20, 0.80, 0.35), (0.30, 0.65, 0.95), (0.95, 0.75, 0.15), (0.85, 0.35, 0.85)] + + +def _load(bundle: str, checkpoint: str, sdf_npz: str | None): + """Load the trained navigator + SDF field for a bundle.""" + from pycvc_gl.scenes import building_occupancy, terrain_grid + + ck = torch.load(checkpoint, map_location="cpu") + meta = ck["meta"] + model = sdf_nav.CoefMLP() + model.load_state_dict(ck["model_state_dict"]) + model.eval() + bounds = terrain_grid(os.path.join(bundle, "terrain.json"))[1] + if sdf_npz and os.path.exists(sdf_npz): + d = np.load(sdf_npz) + phi, nxg, nyg = d["phi"], d["normal_x"], d["normal_y"] + else: + occ = building_occupancy( + os.path.join(bundle, "buildings.glb"), bounds, 512, 512, inflate_m=0.0 + ) + phi, nxg, nyg = sdf_nav.build_sdf(occ, bounds, float(meta["scale"])) + field = sdf_nav.SDFField( + phi, nxg, nyg, bounds, meta["center"], float(meta["scale"]), device="cpu" + ) + return field, model, meta + + +def _hsamp_fn(bundle: str): + """A bilinear terrain-height sampler ``h(x, y)`` + the terrain polydata builder.""" + d = json.load(open(os.path.join(bundle, "terrain.json"))) + b = d["bounds"] + grid = list(reversed(d["grid"])) + rows, cols = d["rows"], d["cols"] + mnx, mny, mxx, mxy = b["min_x"], b["min_y"], b["max_x"], b["max_y"] + dx = (mxx - mnx) / (cols - 1) + dy = (mxy - mny) / (rows - 1) + + def hsamp(x, y): + fx = min(max((x - mnx) / dx, 0), cols - 1) + fy = min(max((y - mny) / dy, 0), rows - 1) + c0, r0 = int(fx), int(fy) + c1, r1 = min(c0 + 1, cols - 1), min(r0 + 1, rows - 1) + tx, ty = fx - c0, fy - r0 + return (grid[r0][c0] * (1 - tx) + grid[r0][c1] * tx) * (1 - ty) + ( + grid[r1][c0] * (1 - tx) + grid[r1][c1] * tx + ) * ty + + return hsamp, (grid, rows, cols, mnx, mny, dx, dy) + + +def _smooth(a, k=5): + a = np.asarray(a, np.float32) + out = a.copy() + for i in range(len(a)): + out[i] = a[max(0, i - k) : i + k + 1].mean(0) + return out + + +def render_drive( + bundle: str, + legs, + goals, + out_mp4: str, + *, + minutes: float = 3.0, + fps: int = 15, + hud: bool = True, + size=(960, 540), + cam="drone", +) -> int: + """Render a drive (``legs`` = list of ``(goal_index, [NavMetrics,...])``) to mp4. + + Paces by arc length (constant cruise + per-leg accel/decel), places a shaded + vehicle + goal spires, drives a drone chase camera, and draws the HUD from the + metrics carried by each leg. Returns the number of frames written.""" + import vtkmodules.vtkRenderingOpenGL2 # noqa: F401 + from vtkmodules.vtkCommonCore import vtkPoints + from vtkmodules.vtkCommonDataModel import vtkCellArray, vtkPolyData + from vtkmodules.vtkCommonTransforms import vtkTransform + from vtkmodules.vtkFiltersCore import vtkPolyDataNormals, vtkQuadricDecimation + from vtkmodules.vtkFiltersGeometry import vtkCompositeDataGeometryFilter + from vtkmodules.vtkFiltersSources import vtkConeSource, vtkCubeSource + from vtkmodules.vtkIOGeometry import vtkGLTFReader + from vtkmodules.vtkIOImage import vtkPNGWriter + from vtkmodules.vtkRenderingCore import ( + vtkActor, + vtkLight, + vtkPolyDataMapper, + vtkRenderer, + vtkRenderWindow, + vtkTextActor, + vtkWindowToImageFilter, + ) + + w, h = size + frames_dir = os.path.join(os.path.dirname(out_mp4) or ".", "_frames") + os.makedirs(frames_dir, exist_ok=True) + os.system("rm -f %s/f_*.png" % frames_dir) + hsamp, (grid, rows, cols, mnx, mny, dx, dy) = _hsamp_fn(bundle) + + pts = vtkPoints() + pts.SetNumberOfPoints(rows * cols) + for r in range(rows): + for c in range(cols): + pts.SetPoint(r * cols + c, mnx + c * dx, mny + r * dy, float(grid[r][c])) + tris = vtkCellArray() + for r in range(rows - 1): + for c in range(cols - 1): + vv = r * cols + c + for cell in ((vv, vv + 1, vv + cols), (vv + 1, vv + cols + 1, vv + cols)): + tris.InsertNextCell(3) + for idx in cell: + tris.InsertCellPoint(idx) + tpd = vtkPolyData() + tpd.SetPoints(pts) + tpd.SetPolys(tris) + tn = vtkPolyDataNormals() + tn.SetInputData(tpd) + tn.Update() + tm = vtkPolyDataMapper() + tm.SetInputConnection(tn.GetOutputPort()) + tm.ScalarVisibilityOff() + terrain = vtkActor() + terrain.SetMapper(tm) + terrain.GetProperty().SetColor(0.33, 0.39, 0.27) + terrain.GetProperty().SetAmbient(0.25) + + gr = vtkGLTFReader() + gr.SetFileName(os.path.join(bundle, "buildings.glb")) + gr.Update() + gf = vtkCompositeDataGeometryFilter() + gf.SetInputConnection(gr.GetOutputPort()) + gf.Update() + dec = vtkQuadricDecimation() + dec.SetInputData(gf.GetOutput()) + dec.SetTargetReduction(0.65) + dec.Update() + bm = vtkPolyDataMapper() + bm.SetInputData(dec.GetOutput()) + bm.SetStatic(1) + bm.ScalarVisibilityOff() + bld = vtkActor() + bld.SetMapper(bm) + bp = bld.GetProperty() + bp.SetColor(0.72, 0.72, 0.77) + bp.SetAmbient(0.4) + bp.SetDiffuse(0.72) + + cs = vtkCubeSource() + cs.SetXLength(4.6) + cs.SetYLength(2.0) + cs.SetZLength(1.6) + cs.Update() + vm = vtkPolyDataMapper() + vm.SetInputConnection(cs.GetOutputPort()) + veh = vtkActor() + veh.SetMapper(vm) + veh.GetProperty().SetColor(0.90, 0.13, 0.12) + veh.GetProperty().SetAmbient(0.3) + veh.GetProperty().SetDiffuse(0.8) + + def shaded_cone(height, rad, col, down=True): + c = vtkConeSource() + c.SetHeight(height) + c.SetRadius(rad) + c.SetResolution(26) + c.SetDirection(0, 0, -1 if down else 1) + c.Update() + n = vtkPolyDataNormals() + n.SetInputConnection(c.GetOutputPort()) + n.Update() + m = vtkPolyDataMapper() + m.SetInputConnection(n.GetOutputPort()) + a = vtkActor() + a.SetMapper(m) + pr = a.GetProperty() + pr.SetColor(*col) + pr.SetAmbient(0.22) + pr.SetDiffuse(0.85) + pr.SetSpecular(0.25) + pr.SetSpecularPower(18) + return a + + veh_beacon = shaded_cone(24, 5.0, (0.98, 0.42, 0.10), down=True) + goal_beacons = [ + shaded_cone(150, 9.0, _GOAL_COLORS[i % len(_GOAL_COLORS)], down=False) + for i in range(len(goals)) + ] + + ren = vtkRenderer() + ren.SetBackground(0.16, 0.19, 0.13) + ren.SetBackground2(0.55, 0.68, 0.82) + ren.GradientBackgroundOn() + for a in (terrain, bld, veh, veh_beacon): + ren.AddActor(a) + for gi, gb in enumerate(goal_beacons): + t = vtkTransform() + t.Translate(goals[gi][0], goals[gi][1], hsamp(*goals[gi]) + 74) + gb.SetUserTransform(t) + ren.AddActor(gb) + sun = vtkLight() + sun.SetPosition(mnx - 500, mny - 500, 2500) + sun.SetFocalPoint(mnx + (cols - 1) * dx / 2, mny + (rows - 1) * dy / 2, 0) + sun.SetIntensity(1.0) + sun.SetLightTypeToSceneLight() + ren.AddLight(sun) + + hud_actor = None + if hud: + hud_actor = vtkTextActor() + hud_actor.SetDisplayPosition(16, h - 24) + tp = hud_actor.GetTextProperty() + tp.SetFontFamilyToCourier() + tp.SetFontSize(15) + tp.SetColor(0.92, 0.97, 0.85) + tp.SetLineOffset(-2) + tp.SetVerticalJustificationToTop() + tp.SetBackgroundColor(0.05, 0.07, 0.03) + tp.SetBackgroundOpacity(0.55) + ren.AddActor2D(hud_actor) + + rw = vtkRenderWindow() + rw.SetOffScreenRendering(1) + rw.AddRenderer(ren) + rw.SetSize(w, h) + camera = ren.GetActiveCamera() + camera.SetClippingRange(1.0, 9000.0) + + # arc-length pacing so speed is constant (per-leg sin ease = accel/decel); total + # frames set to `minutes` for a steady cruise regardless of path length. + total_frames = int(minutes * 60 * fps) + smoothed = [] + lens = [] + for gi, ms in legs: + seg = np.array([[m.x, m.y] for m in ms], np.float32) + if len(seg) < 2: + continue + seg_s = _smooth(seg) + dd = np.linalg.norm(np.diff(seg_s, axis=0), axis=1) + arcl = np.concatenate([[0.0], np.cumsum(dd)]) + smoothed.append((gi, seg_s, arcl, ms)) + lens.append(float(arcl[-1])) + tot = sum(lens) or 1.0 + seq = [] # (x, y, goal_index, NavMetrics) + for (gi, seg_s, arcl, ms), leg_len in zip(smoothed, lens): + if leg_len < 1.0: + continue + n = max(120, int(round(total_frames * leg_len / tot))) + u = np.linspace(0, 1, n) + spd = 0.35 + 0.65 * np.sin(np.pi * u) + cum = np.cumsum(spd) + cum = cum / cum[-1] * leg_len + xs = np.interp(cum, arcl, seg_s[:, 0]) + ys = np.interp(cum, arcl, seg_s[:, 1]) + midx = np.clip(np.searchsorted(arcl, cum), 0, len(ms) - 1) # nearest nav-step for the HUD + for k in range(n): + seq.append((float(xs[k]), float(ys[k]), gi, ms[int(midx[k])])) + + stats = NavStats() + w2i = vtkWindowToImageFilter() + w2i.SetInput(rw) + head = np.array([1.0, 0.0]) + back, ht, ahead = (74.0, 52.0, 24.0) if cam == "drone" else (34.0, 16.0, 22.0) + prev = np.array([seq[0][0], seq[0][1]]) + for fi, (px, py, ag, m) in enumerate(seq): + p = np.array([px, py]) + dvec = p - prev + prev = p + nn = np.linalg.norm(dvec) + if nn > 1e-4: + head = 0.6 * head + 0.4 * (dvec / nn) + head /= np.linalg.norm(head) + 1e-9 + zt = hsamp(px, py) + 0.9 + tf = vtkTransform() + tf.Translate(px, py, zt + 0.8) + tf.RotateZ(math.degrees(math.atan2(head[1], head[0]))) + veh.SetUserTransform(tf) + bt = vtkTransform() + bt.Translate(px, py, zt + 24.0) + veh_beacon.SetUserTransform(bt) + for gi, gb in enumerate(goal_beacons): + gb.GetProperty().SetAmbient(0.62 if gi == ag else 0.22) + camera.SetPosition(px - head[0] * back, py - head[1] * back, zt + ht) + camera.SetFocalPoint(px + head[0] * ahead, py + head[1] * ahead, zt - 2.0) + camera.SetViewUp(0, 0, 1) + if hud_actor is not None: + stats.update(m) + hud_actor.SetInput("\n".join(hud_lines(m, stats))) + rw.Render() + w2i.Modified() + w2i.Update() + wr = vtkPNGWriter() + wr.SetFileName("%s/f_%05d.png" % (frames_dir, fi)) + wr.SetInputConnection(w2i.GetOutputPort()) + wr.Write() + subprocess.run( + [ + "ffmpeg", + "-y", + "-framerate", + str(fps), + "-i", + "%s/f_%%05d.png" % frames_dir, + "-c:v", + "libx264", + "-pix_fmt", + "yuv420p", + "-vf", + "scale=%d:%d" % (w, h), + out_mp4, + ], + check=True, + ) + return len(seq) + + +def capture_drive( + bundle: str, + checkpoint: str, + start, + goal, + out: str = "drive.mp4", + *, + sdf_npz: str | None = None, + minutes: float = 1.0, + fps: int = 15, + hud: bool = True, + cam: str = "chase", +) -> str: + """Drive a single fixed A->B run and render it (with HUD) to ``out``.""" + field, model, meta = _load(bundle, checkpoint, sdf_npz) + nav = SdfNavigator(field, model, meta, reach_tol=0.5) + nav.start(start, goal) + ms, best_i, _bo, _bv = nav.drive_to_goal(max_steps=3000) + print( + "drive: %d steps, best-approach %.1f m, penetration %d" + % (len(ms), ms[best_i].goal_dist_m, sum(1 for m in ms if m.inside_building)) + ) + n = render_drive( + bundle, [(0, ms[: best_i + 1])], [goal], out, minutes=minutes, fps=fps, hud=hud, cam=cam + ) + print("wrote %s (%d frames)" % (out, n)) + return out + + +def capture_multigoal( + bundle: str, + checkpoint: str, + out: str = "multigoal.mp4", + *, + sdf_npz: str | None = None, + minutes: float = 3.0, + fps: int = 15, + hud: bool = True, + n_goals: int = 4, +) -> str: + """Dynamic multi-goal free-drive (goals chosen in the SDF's free space, re-targeted + live, drone chase cam) rendered with HUD to ``out``.""" + field, model, meta = _load(bundle, checkpoint, sdf_npz) + start, goals = select_reachable_goals(field, model, meta, n_corners=n_goals) + print("start %s goals %s" % (np.round(start).tolist(), [np.round(g).tolist() for g in goals])) + nav = SdfNavigator(field, model, meta, reach_tol=0.8) + o = torch.from_numpy(nav.w2n(start)).unsqueeze(0).float() + v = torch.zeros(1, 2) + legs = [] + total_pen = 0 + for gi, g in enumerate(goals): + nav.o = o.clone() + nav.v = v.clone() + nav.set_goal(g, goal_index=gi) + ms, best_i, o, v = nav.drive_to_goal(max_steps=1300) + legs.append((gi, ms[: best_i + 1])) + total_pen += sum(1 for m in ms[: best_i + 1] if m.inside_building) + print( + "multigoal drive: %d steps over %d legs, penetration=%d" + % (sum(len(ms) for _, ms in legs), len(legs), total_pen) + ) + n = render_drive(bundle, legs, goals, out, minutes=minutes, fps=fps, hud=hud, cam="drone") + print("wrote %s (%d frames, %.0fs)" % (out, n, n / fps)) + return out diff --git a/grl_snam/tools/obstacles.py b/grl_snam/tools/obstacles.py new file mode 100644 index 0000000..2de1fb3 --- /dev/null +++ b/grl_snam/tools/obstacles.py @@ -0,0 +1,69 @@ +"""Geometry bundle -> circular-obstacle set for the CoefEnergyNet (circle-surrogate) +training track. See :func:`extract`. (The SDF track uses :mod:`grl_snam.tools.sdf` +instead.)""" + +from __future__ import annotations + +import os + +import numpy as np + + +def extract( + bundle_dir: str, + grid: int = 512, + block: int = 2, + radius_frac: float = 0.6, + robot_radius_world: float = 3.0, +) -> dict: + """Bundle -> (obstacle centers, radius, free-point pool), all in WORLD units. + + The city footprint (``buildings.glb``) is rasterized over the ``terrain.json`` + bounds to a ``grid x grid`` solid occupancy mask, coarsened by ``block`` (each + obstacle stands in for a ``block x block`` wall patch), and every occupied coarse + cell becomes one circular obstacle at the cell centre with radius + ``radius_frac * cell_size``. Free cells become the drivable pool. Needs the + ``pycvc_gl`` scene helpers (VTK rasterization); imported lazily.""" + from pycvc_gl.scenes import building_occupancy, terrain_grid + + terrain = os.path.join(bundle_dir, "terrain.json") + glb = os.path.join(bundle_dir, "buildings.glb") + _, bounds, _, _ = terrain_grid(terrain) + mnx, mny, mxx, mxy = bounds + + occ = building_occupancy(glb, bounds, grid, grid, inflate_m=0.0) + ny, nx = occ.shape + + cny, cnx = ny // block, nx // block + cocc = occ[: cny * block, : cnx * block].reshape(cny, block, cnx, block).any(axis=(1, 3)) + csx = (mxx - mnx) / cnx + csy = (mxy - mny) / cny + ys, xs = np.where(cocc) + centers = np.stack([mnx + (xs + 0.5) * csx, mny + (ys + 0.5) * csy], 1).astype(np.float32) + radius_world = float(radius_frac * 0.5 * (csx + csy)) + + gx = np.linspace(mnx, mxx, nx, dtype=np.float32) + gy = np.linspace(mny, mxy, ny, dtype=np.float32) + free_r, free_c = np.where(~occ) + free_pool = np.stack([gx[free_c], gy[free_r]], 1).astype(np.float32) + + return { + "centers": centers, + "radius_world": np.float32(radius_world), + "robot_radius_world": np.float32(robot_radius_world), + "bounds": np.asarray(bounds, np.float32), + "free_pool": free_pool, + "cell_size": np.float32(0.5 * (csx + csy)), + } + + +def extract_to_npz(bundle_dir: str, out: str | None = None, grid: int = 512, block: int = 2) -> str: + """Run :func:`extract` and save an ``obstacles.npz`` (default ``/obstacles.npz``).""" + data = extract(bundle_dir, grid=grid, block=block) + out = out or os.path.join(bundle_dir, "obstacles.npz") + np.savez_compressed(out, **data) + print( + "extracted %d obstacles (r=%.1fm) + %d free points -> %s" + % (len(data["centers"]), float(data["radius_world"]), len(data["free_pool"]), out) + ) + return out diff --git a/grl_snam/tools/pipeline.py b/grl_snam/tools/pipeline.py new file mode 100644 index 0000000..5d11e56 --- /dev/null +++ b/grl_snam/tools/pipeline.py @@ -0,0 +1,63 @@ +"""The whole learned-nav pipeline in one call: from nothing but a **world model** +(a scene bundle = ``terrain.json`` + ``buildings.glb``) through SDF construction and +self-supervised training to a rendered demo video. This is ``grl-snam pipeline``. +""" + +from __future__ import annotations + +import os + +from . import capture, sdf, train + + +def run( + bundle: str, + *, + source: str = "edt", + region: float = 430.0, + steps: int = 1500, + out_dir: str | None = None, + minutes: float = 3.0, + fps: int = 15, + hud: bool = True, + drive: str = "multigoal", +) -> dict: + """world model -> SDF -> trained policy -> demo video. Returns the artifact paths. + + ``drive`` selects the demo video: ``multigoal`` (dynamic re-targeting) or + ``single`` (a fixed A->B). Reuses an existing ``nav_sdf.npz`` / checkpoint if + present in ``out_dir`` unless absent.""" + out_dir = out_dir or bundle + os.makedirs(out_dir, exist_ok=True) + sdf_npz = os.path.join(out_dir, "nav_sdf.npz") + ckpt = os.path.join(out_dir, "coef_sdf.pt") + video = os.path.join(out_dir, "grl_snam_%s.mp4" % drive) + + print("== [1/3] build SDF from the world model ==") + sdf.build(bundle, source=source, region=region, out=sdf_npz) + print("== [2/3] train the navigation policy (self-supervised) ==") + train.train_sdf(sdf_npz, ckpt, steps=steps) + print("== [3/3] drive + render the demo video (%s) ==" % drive) + if drive == "single": + # a reachable A->B is scene-specific; reuse the multigoal selector to find one + from ..nav import select_reachable_goals + + field, model, meta = capture._load(bundle, ckpt, sdf_npz) + start, goals = select_reachable_goals(field, model, meta, n_corners=1) + capture.capture_drive( + bundle, + ckpt, + tuple(start), + tuple(goals[0]), + video, + sdf_npz=sdf_npz, + minutes=minutes, + fps=fps, + hud=hud, + ) + else: + capture.capture_multigoal( + bundle, ckpt, video, sdf_npz=sdf_npz, minutes=minutes, fps=fps, hud=hud + ) + print("== pipeline done ==") + return {"sdf": sdf_npz, "checkpoint": ckpt, "video": video} diff --git a/grl_snam/tools/sdf.py b/grl_snam/tools/sdf.py new file mode 100644 index 0000000..a9e49bc --- /dev/null +++ b/grl_snam/tools/sdf.py @@ -0,0 +1,93 @@ +"""Geometry bundle -> navigation SDF (``nav_sdf.npz``) — the primary obstacle model +for the learned-nav demo. Two interchangeable sources: + +* ``edt`` (default) — grl-snam's own exact 2-D footprint distance transform. +* ``cvc`` — CVC's mesh-exact 3-D SDF via ``pycvc.sdf`` (SDF_V2), sliced to the ground + plane; leans on the CVC compute layer and is the substrate for 3-D GRL-SNAM later. + +Both emit the same normalized ``phi`` + unit-normal grids + meta that +:mod:`grl_snam.tools.train` and the navigator consume source-agnostically. +""" + +from __future__ import annotations + +import os + +import numpy as np + +import sdf_nav + +TARGET_EXTENT = 10.0 # the surrogate's normalized working scale (~10 units across the region) + + +def _gltf_mesh(glb_path: str): + """Flat ``(verts[x,y,z,...], tris[i,j,k,...])`` from a glTF/GLB via VTK (lazy import).""" + import vtkmodules.vtkRenderingOpenGL2 # noqa: F401 (register factories) + from vtkmodules.util.numpy_support import vtk_to_numpy + from vtkmodules.vtkFiltersGeometry import vtkCompositeDataGeometryFilter + from vtkmodules.vtkIOGeometry import vtkGLTFReader + + reader = vtkGLTFReader() + reader.SetFileName(glb_path) + reader.Update() + geom = vtkCompositeDataGeometryFilter() + geom.SetInputConnection(reader.GetOutputPort()) + geom.Update() + pd = geom.GetOutput() + pts = vtk_to_numpy(pd.GetPoints().GetData()).astype(np.float64) + polys = vtk_to_numpy(pd.GetPolys().GetData()) # [n0,i,j,k, n1,...]; triangulated -> n0==3 + tris: list[int] = [] + i = 0 + while i < len(polys): + n = int(polys[i]) + if n == 3: + tris += [int(polys[i + 1]), int(polys[i + 2]), int(polys[i + 3])] + i += n + 1 + return pts.reshape(-1).tolist(), tris + + +def build( + bundle_dir: str, + source: str = "edt", + region: float = 430.0, + grid: int = 512, + cvc_dim=(512, 512, 48), + out: str | None = None, +) -> str: + """Build the navigation SDF for a scene bundle and save ``nav_sdf.npz``. Returns the + output path. Needs the ``pycvc_gl`` scene helpers (imported lazily).""" + from pycvc_gl.scenes import building_occupancy, terrain_grid + + terrain = os.path.join(bundle_dir, "terrain.json") + glb = os.path.join(bundle_dir, "buildings.glb") + _, bounds, _, _ = terrain_grid(terrain) + mnx, mny, mxx, mxy = bounds + cx, cy = 0.5 * (mnx + mxx), 0.5 * (mny + mxy) + scale = TARGET_EXTENT / (2.0 * region) + + if source == "edt": + occ = building_occupancy(glb, bounds, grid, grid, inflate_m=0.0) + phi, nxg, nyg = sdf_nav.build_sdf(occ, bounds, scale) + elif source == "cvc": + verts, tris = _gltf_mesh(glb) + phi, nxg, nyg = sdf_nav.build_sdf_cvc(verts, tris, bounds, scale, dim=tuple(cvc_dim)) + else: + raise ValueError(f"unknown SDF source {source!r} (expected 'edt' or 'cvc')") + + out = out or os.path.join(bundle_dir, "nav_sdf.npz") + np.savez_compressed( + out, + phi=phi, + normal_x=nxg, + normal_y=nyg, + bounds=np.asarray(bounds, np.float32), + center=np.asarray([cx, cy], np.float32), + scale=np.float32(scale), + region=np.float32(region), + source=source, + ) + print( + "built %s SDF %s (phi[%.2f,%.2f]) -> %s" + % (source, phi.shape, float(phi.min()), float(phi.max()), out) + ) + return out diff --git a/grl_snam/tools/train.py b/grl_snam/tools/train.py new file mode 100644 index 0000000..d524183 --- /dev/null +++ b/grl_snam/tools/train.py @@ -0,0 +1,138 @@ +"""Self-supervised training of the SDF navigation coefficients. + +Reads a ``nav_sdf.npz`` (from :mod:`grl_snam.tools.sdf`) and trains +``sdf_nav.CoefMLP`` to predict ``(alpha, beta, gamma)`` for the differentiable SDF +surrogate by back-propagating a reach-goal + no-collision objective through the +rollout — no labels. The net is biased toward the known-good navigating regime, so +it starts near-optimal and converges in a few hundred to a few thousand steps. Runs +on CPU (fast — the SDF removes the per-step obstacle search) or CUDA if present. +""" + +from __future__ import annotations + +import json +import os +import time + +import numpy as np +import torch +import torch.nn.functional as F + +import sdf_nav + + +def train_sdf( + sdf_npz: str, + out: str = "coef_sdf.pt", + *, + steps: int = 1500, + batch: int = 128, + horizon: int = 28, + dt: float = 0.06, + d_hat: float = 0.35, + robot_radius_world: float = 3.0, + vmax: float = 0.9, + nsub_infer: int = 3, + lr: float = 3e-4, + threads: int = 6, + seed: int = 0, + log_every: int = 300, +) -> str: + """Train the SDF coefficient net on ``sdf_npz`` and save a checkpoint ``.pt`` + (state dict + physics ``meta``) plus a ``_result.json``. Returns the checkpoint path.""" + dev = torch.device("cuda" if torch.cuda.is_available() else "cpu") + torch.set_num_threads(threads) + torch.manual_seed(seed) + np.random.seed(seed) + + d = np.load(sdf_npz) + phi, nxg, nyg = d["phi"], d["normal_x"], d["normal_y"] + bounds = [float(v) for v in d["bounds"]] + center = [float(v) for v in d["center"]] + scale = float(d["scale"]) + region = float(d["region"]) + rr = robot_radius_world * scale + field = sdf_nav.SDFField(phi, nxg, nyg, bounds, center, scale, device=dev) + mnx, mny, mxx, mxy = bounds + print( + "device=%s field=%s phi[%.2f,%.2f] rr=%.3f S=%.5f" + % (dev, phi.shape, phi.min(), phi.max(), rr, scale) + ) + + # drivable free-point pool (normalized, centered) within the working region + ny_, nx_ = phi.shape + gxs = np.linspace(mnx, mxx, nx_) + gys = np.linspace(mny, mxy, ny_) + grid_x, grid_y = np.meshgrid(gxs, gys) + free = phi > (rr + 0.02) + reg = (np.abs(grid_x - center[0]) < region) & (np.abs(grid_y - center[1]) < region) + sel = free & reg + pool = np.stack([(grid_x[sel] - center[0]) * scale, (grid_y[sel] - center[1]) * scale], 1) + poolt = torch.from_numpy(pool.astype(np.float32)).to(dev) + print("pool=%d drivable points" % len(pool)) + + model = sdf_nav.CoefMLP().to(dev) + opt = torch.optim.Adam(model.parameters(), lr=lr) + kw = dict(rr=rr, d_hat=d_hat, dt=dt, vmax=vmax) + t0 = time.time() + for it in range(steps): + si = torch.randint(0, len(poolt), (batch,), device=dev) + o0 = poolt[si] + ang = torch.rand(batch, device=dev) * 6.2832 + dd = 1.0 + torch.rand(batch, device=dev) # 1..2 local goal + goal = o0 + torch.stack([torch.cos(ang), torch.sin(ang)], -1) * dd.unsqueeze(-1) + al, be, ga = model(sdf_nav.coef_feats(field, o0, goal)) + oT, _vT, clr = sdf_nav.sdf_rollout( + field, o0, torch.zeros_like(o0), goal, al, be, ga, horizon, nsub=1, **kw + ) + loss_goal = ((oT - goal) ** 2).sum(-1).mean() + loss_col = F.softplus((0.02 - clr) / 0.02).mean() + loss_reg = ((be - 3.0) ** 2).mean() + ((ga - 4.0) ** 2).mean() + ((al - 1.0) ** 2).mean() + loss = loss_goal + 3.0 * loss_col + 0.1 * loss_reg + opt.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0) + opt.step() + if it % log_every == 0: + print( + "it=%d/%d L_goal=%.3f L_col=%.3f al=%.2f be=%.2f ga=%.2f (%.0fs)" + % ( + it, + steps, + float(loss_goal), + float(loss_col), + float(al.mean()), + float(be.mean()), + float(ga.mean()), + time.time() - t0, + ), + flush=True, + ) + dt_total = time.time() - t0 + print( + "TIME %d steps %.1fs = %.1f ms/step on %s" + % (steps, dt_total, 1000 * dt_total / max(1, steps), dev) + ) + + meta = { + "scale": scale, + "center": center, + "region": region, + "rr": rr, + "d_hat": d_hat, + "dt": dt, + "horizon": horizon, + "nsub": nsub_infer, + "vmax": vmax, + "steps": steps, + "bounds": bounds, + "kind": "sdf", + "sdf_npz": os.path.abspath(sdf_npz), + } + torch.save({"model_state_dict": model.state_dict(), "meta": meta}, out) + print("SAVED %s" % out) + with open(os.path.splitext(out)[0] + "_result.json", "w") as fh: + json.dump( + {"steps": steps, "ms_per_step": 1000 * dt_total / max(1, steps), "device": str(dev)}, fh + ) + return out diff --git a/grl_snam_lab/README.md b/grl_snam_lab/README.md deleted file mode 100644 index 9cbd6dc..0000000 --- a/grl_snam_lab/README.md +++ /dev/null @@ -1,34 +0,0 @@ -# grl_snam_lab — a general GRL-SNAM visualization laboratory - -A domain-general 3D lab for GRL-SNAM built on the CVC graphics stack -(`pycvc` + `pycvc_gl`, the Python bindings for libcvc / cvcGL). It speaks -general GRL-SNAM vocabulary — **terrain, obstacles, agent paths, agents, -scalar fields (e.g. risk / SDF)** — and turns each into a live 3D -scene node. It has **no** mission- or dataset-specific knowledge. - -```python -from grl_snam_lab import Lab -lab = Lab() -lab.add_terrain(heights, bounds=(-100, -100, 100, 100)) -lab.add_mesh("bldg", verts, faces, color=(0.6, 0.6, 0.6)) -lab.add_path("agent0", [(x0, y0, z0), (x1, y1, z1), ...]) -lab.add_field("risk", grid, dims=(nx, ny, nz), bounds=(-100,-100,0, 100,100,20)) -lab.show() # interactive window (needs a display) -# lab.render_png("s.png") # offscreen snapshot -``` - -## Adapting for a dataset - -Keep dataset-specific code **in the downstream project**. An adapter reads that -project's own track format and drives the same general API — e.g. one `add_path` -per agent track, `add_field` for a scalar raster, `add_mesh` for the scene's -buildings. The lab stays general; the mapping lives with the data. - -## Requirements - -`pycvc` and `pycvc_gl` (from **libcvc** — install libcvc with its Python -bindings into your prefix, e.g. via cvcpkg, and put them on `PYTHONPATH`). -The pure-Python geometry helpers (`terrain_mesh`, …) work without them; the -scene/display methods require them. `show()`/`render_png()` need a GL display. - -Run the demo: `python examples/lab_demo.py` (or `… out.png` for a snapshot). diff --git a/grl_snam_lab/__init__.py b/grl_snam_lab/__init__.py deleted file mode 100644 index 8438468..0000000 --- a/grl_snam_lab/__init__.py +++ /dev/null @@ -1,43 +0,0 @@ -"""grl_snam_lab — a general GRL-SNAM visualization laboratory. - -A thin, **domain-general** layer over the CVC graphics stack (``pycvc`` + -``pycvc_gl``, the SWIG bindings for libcvc / cvcGL). It speaks the vocabulary -of GRL-SNAM — *terrain*, *obstacles*, *agent paths*, *agents*, and *scalar -fields* (e.g. risk / SDF) — and turns each into a live 3D scene node. -It knows nothing about any specific mission or dataset. - -Downstream projects adapt it by feeding their own data through this API — e.g. -an adapter reads that project's own track format and calls :meth:`Lab.add_path` -once per agent track. Keep dataset-specific code in the downstream project; -keep this package general. - -The graphics backend (``pycvc``/``pycvc_gl``) is imported lazily so the module -and its pure-Python geometry helpers import even where the compiled bindings -are not installed; the display/scene methods require them. - -Example:: - - from grl_snam_lab import Lab - lab = Lab() - lab.add_terrain(heights, bounds=(-100, -100, 100, 100)) - lab.add_mesh("bldg", verts, faces, color=(0.6, 0.6, 0.6)) - lab.add_path("agent0", [(x0, y0, z0), (x1, y1, z1), ...]) - lab.add_field("risk", risk_grid, bounds=(-100, -100, 0, 100, 100, 20)) - lab.show() # or lab.render_png("scene.png") - -Ready-to-run demo — inside a running volrover3 (Python Console REPL):: - - >>> import grl_snam_lab - >>> grl_snam_lab.run_in_volrover() # terrain + track + marker in the live window - -...or standalone (own window / PNG) via ``run_standalone()`` or the -``grl-snam-lab-demo`` console script. -""" - -from __future__ import annotations - -from .demo import demo_scene, run_in_volrover, run_standalone -from .lab import Lab, terrain_mesh - -__all__ = ["Lab", "terrain_mesh", "demo_scene", "run_in_volrover", "run_standalone"] -__version__ = "0.1.0" diff --git a/grl_snam_lab/camera.py b/grl_snam_lab/camera.py deleted file mode 100644 index 719da01..0000000 --- a/grl_snam_lab/camera.py +++ /dev/null @@ -1,4 +0,0 @@ -"""Chase camera — moved to :mod:`pycvc_gl.camera` (a generic scene utility that -ships with the pycvc_gl bindings). Re-exported here for backward compatibility.""" -from pycvc_gl.camera import * # noqa: F401,F403 -from pycvc_gl.camera import ChaseCamera # noqa: F401 diff --git a/grl_snam_lab/demo.py b/grl_snam_lab/demo.py deleted file mode 100644 index f6d53a3..0000000 --- a/grl_snam_lab/demo.py +++ /dev/null @@ -1,200 +0,0 @@ -"""Runnable demos for the GRL-SNAM lab — static and animated. - -Entry points: - -* ``run_in_volrover()`` — build the STATIC demo (terrain + full track + agent) in a - RUNNING volrover3's live scene. Call it from the volrover3 Python Console REPL:: - - >>> import grl_snam_lab - >>> grl_snam_lab.run_in_volrover() - -* ``run_standalone(png=...)`` — the same static scene in a self-owned window / PNG - (the ``grl-snam-lab-demo`` console script). - -* ANIMATED job — ``examples/volrover_lab_animated.py`` is a JobScheduler ``step(dt)`` - script that walks the agent along the track in the live volrover3 window (Load - Script -> Run as Job). It reuses ``demo_scene`` + ``animate_agent`` below. - -``demo_scene(lab)`` populates any ``Lab`` (the agent marker is added ONCE); -``animate_agent(lab, t)`` moves the agent to its position at time ``t`` by mutating -the node's transform in place — no destroy/recreate, just ``lab.move(...)`` — the -shared core for the static, animated, and frame-capture paths. volrover3 has no -GRL-SNAM dependency; these run only because GRL-SNAM is importable in the -interpreter's environment. -""" - -from __future__ import annotations - -import math - -from .lab import Lab - -__all__ = [ - "demo_scene", - "agent_position", - "animate_agent", - "terrain_height", - "run_in_volrover", - "run_standalone", - "TRACK", -] - -# ── World geometry ─────────────────────────────────────────────────────────── -# The terrain is an ANALYTIC height field h(x,y) over a 200x200 world box; the -# drawn terrain mesh samples it, and the agent + track DRAPE onto it (so the -# robot sits on the ground, not floating). The agent walks a SMOOTH analytic -# loop (a real circle, not a coarse polygon), so its position AND its heading -# vary continuously — the earlier 32-gon made the chase camera jerk at each -# vertex because the velocity jumped there. - -TERRAIN_BOUNDS = (-100.0, -100.0, 100.0, 100.0) # (min_x, min_y, max_x, max_y) -TRACK_RADIUS = 70.0 # loop radius (world units) -LOOP_SECONDS = 22.0 # time for the agent to walk one full lap at speed 1.0 -AGENT_LIFT = 0.4 # small lift so the marker base clears z-fighting with terrain -_N_TRACK = 180 # drawn-polyline resolution (fine => looks like a smooth circle) - - -def terrain_height(x: float, y: float) -> float: - """Analytic terrain height at world ``(x, y)`` — gentle rolling hills (a central - rise, two bumps, a low ripple). The single source of truth for the terrain: the - mesh samples it and the agent/track drape onto it.""" - h = 14.0 * math.exp(-((x * x + y * y) / 3000.0)) # central hill - h += 8.0 * math.exp(-(((x - 55.0) ** 2 + (y + 35.0) ** 2) / 900.0)) # bump - h += 6.0 * math.exp(-(((x + 45.0) ** 2 + (y - 50.0) ** 2) / 1200.0)) # bump - h += 2.5 * math.sin(x * 0.06) * math.cos(y * 0.05) # ripple - return h - - -def _terrain_heights(n: int = 56): - """Sample ``terrain_height`` on an n x n grid over TERRAIN_BOUNDS (row-major, - row index -> y, col index -> x), matching ``terrain_mesh``'s layout.""" - min_x, min_y, max_x, max_y = TERRAIN_BOUNDS - sx = (max_x - min_x) / (n - 1) - sy = (max_y - min_y) / (n - 1) - return [[terrain_height(min_x + j * sx, min_y + i * sy) for j in range(n)] for i in range(n)] - - -def _loop_point(theta: float): - """A point on the agent's demo loop at angle ``theta``, draped onto the terrain. - - The radius WANDERS (a 3-lobed modulation) so the heading changes at a varying - rate — this is a stand-in for a real GRL-SNAM path, and it gives the - position-driven chase camera genuine direction changes to smooth. (The camera - never sees this function; it only sees the emitted positions.)""" - r = TRACK_RADIUS + 8.0 * math.sin(3.0 * theta) - x = r * math.cos(theta) - y = r * math.sin(theta) - return (x, y, terrain_height(x, y) + AGENT_LIFT) - - -# The drawn track: a fine, closed, draped ring (looks like a smooth circle on the -# terrain). Shared by the static path (drawn whole) and the animation. -TRACK = [_loop_point(2.0 * math.pi * k / _N_TRACK) for k in range(_N_TRACK + 1)] - - -def _agent_marker_mesh(size: float = 5.0): - """A small upward tetrahedron with its BASE at local z=0 (apex at +size), so - placing it via setPosition(x, y, terrain_height) sits it ON the ground. Built - ONCE at the origin and then MOVED via the node transform (see animate_agent).""" - s = size - verts = [0.0, 0.0, s, -s, -s, 0.0, s, -s, 0.0, 0.0, s, 0.0] - tris = [0, 1, 2, 0, 2, 3, 0, 3, 1, 1, 3, 2] - return verts, tris - - -def _theta(t: float, speed: float) -> float: - return 2.0 * math.pi * ((t * speed / LOOP_SECONDS) % 1.0) - - -def agent_position(t: float, speed: float = 1.0): - """The agent's (x, y, z) at time ``t`` on the demo loop, draped onto the - terrain. ``speed`` scales time (1.0 => one lap per LOOP_SECONDS). This is only - a stand-in path for the demo — the live GRL-SNAM planner will supply real - positions; feed whichever stream you have to ``ChaseCamera`` (grl_snam_lab.camera), - which derives the camera heading from the positions themselves.""" - return _loop_point(_theta(t, speed)) - - -def animate_agent(lab: Lab, t: float) -> Lab: - """Move the ``agent0`` marker to its position at time ``t`` by mutating the - node's transform — the marker walks the track with NO destroy/recreate (the - geometry + VTK actor are reused; only the transform changes). ``demo_scene`` - must have added ``agent0`` first. Returns ``lab``.""" - lab.move("agent0", *agent_position(t)) - lab.pump() - return lab - - -def demo_scene(lab: Lab) -> Lab: - """Add the canonical demo — terrain + the agent's track + the agent marker - (added ONCE at the origin, then placed at its t=0 position) — to ``lab``. - Returns ``lab`` so calls chain.""" - lab.add_terrain( - _terrain_heights(), bounds=(-100.0, -100.0, 100.0, 100.0), color=(0.42, 0.53, 0.34) - ) - lab.add_path("agent0_track", TRACK, color=(0.95, 0.75, 0.10)) - # Build the agent marker ONCE (at the origin); animate_agent MOVES it. - verts, tris = _agent_marker_mesh() - lab.add_mesh("agent0", verts, tris, color=(1.0, 0.20, 0.20)) - animate_agent(lab, 0.0) # place it at the track start - return lab - - -def run_in_volrover() -> Lab: - """Build the STATIC demo in volrover3's RUNNING scene (call from the REPL). Adopts - the host's live app + scene via ``vrhost``; raises a clear error outside volrover3. - For the ANIMATED version, load ``examples/volrover_lab_animated.py`` as a job.""" - try: - import vrhost # host-injected inside volrover3; not a GRL-SNAM dependency - except ImportError as exc: - raise RuntimeError( - "grl_snam_lab.run_in_volrover(): `vrhost` not found — this must run " - "INSIDE volrover3's embedded Python console. For a standalone window " - "use grl_snam_lab.run_standalone()." - ) from exc - - lab = demo_scene(Lab(app=vrhost.app(), scene=vrhost.scene())) - lab.pump() - print( - f"grl_snam_lab: live volrover3 scene now has {lab.num_nodes()} node(s) " - "(terrain + agent0_track + agent0) — look at the viewport. For the animated " - "walk, Load Script -> Run as Job: examples/volrover_lab_animated.py" - ) - return lab - - -def run_standalone(png: str | None = None, width: int = 1024, height: int = 768) -> Lab: - """Build the static demo in a self-owned Lab; render an offscreen PNG when ``png`` - is given (headless), else open a blocking interactive window (needs a display).""" - lab = demo_scene(Lab()) - if png: - lab.render_png(png, width, height) - print(f"grl_snam_lab: wrote {png} ({lab.num_nodes()} nodes)") - else: - print(f"grl_snam_lab: showing {lab.num_nodes()} nodes — close the window to exit") - lab.show("grl_snam_lab demo", width, height) - return lab - - -def main(argv: list[str] | None = None) -> int: - """Console entry point (``grl-snam-lab-demo``): standalone viz demo. - - ``grl-snam-lab-demo`` opens a window; ``grl-snam-lab-demo out.png`` renders an - offscreen snapshot. Inside volrover3 use the REPL: - ``import grl_snam_lab; grl_snam_lab.run_in_volrover()``.""" - import argparse - - ap = argparse.ArgumentParser( - prog="grl-snam-lab-demo", - description="GRL-SNAM lab demo (terrain + agent track + marker), standalone.", - ) - ap.add_argument( - "png", nargs="?", help="write an offscreen PNG here instead of opening a window" - ) - args = ap.parse_args(argv) - try: - run_standalone(png=args.png) - except Exception as exc: # pycvc/pycvc_gl missing, or no GL display - print(f"grl-snam-lab-demo: {exc}") - return 1 - return 0 diff --git a/grl_snam_lab/lab.py b/grl_snam_lab/lab.py deleted file mode 100644 index 6ae1fc1..0000000 --- a/grl_snam_lab/lab.py +++ /dev/null @@ -1,8 +0,0 @@ -"""Scene Lab + pure-Python geometry helpers — moved to :mod:`pycvc_gl.lab` -(a generic scene utility that ships with the pycvc_gl bindings). Re-exported here -for backward compatibility.""" -from pycvc_gl.lab import * # noqa: F401,F403 - -# `import *` skips underscore-prefixed names, but tests/test_lab.py imports -# `_flatten_points` from here (its old home), so surface the helpers by name too. -from pycvc_gl.lab import Lab, terrain_mesh, polyline_indices, _flatten_points # noqa: F401 diff --git a/grl_snam_lab/scenes.py b/grl_snam_lab/scenes.py deleted file mode 100644 index 5005bb2..0000000 --- a/grl_snam_lab/scenes.py +++ /dev/null @@ -1,7 +0,0 @@ -"""Scene loaders + grounded routing — moved to :mod:`pycvc_gl.scenes`. Re-exported -here for backward compatibility.""" -from pycvc_gl.scenes import * # noqa: F401,F403 -from pycvc_gl.scenes import ( # noqa: F401 - terrain_grid, terrain_sampler, add_terrain_json, add_gltf, load_geometry_bundle, - building_occupancy, plan_ground_route, resample_polyline, -) diff --git a/grl_snam_lab/vehicle.py b/grl_snam_lab/vehicle.py deleted file mode 100644 index 246b44c..0000000 --- a/grl_snam_lab/vehicle.py +++ /dev/null @@ -1,4 +0,0 @@ -"""Vehicle pose smoother — moved to :mod:`pycvc_gl.vehicle`. Re-exported here for -backward compatibility.""" -from pycvc_gl.vehicle import * # noqa: F401,F403 -from pycvc_gl.vehicle import VehiclePose # noqa: F401 diff --git a/pyproject.toml b/pyproject.toml index ce108b0..658c9c6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ + "click>=8.1", "torch>=2.1", "numpy>=1.24", "matplotlib>=3.7", @@ -29,24 +30,25 @@ Repository = "https://github.com/CVC-Lab/GRL-SNAM" "Bug Tracker" = "https://github.com/CVC-Lab/GRL-SNAM/issues" [project.scripts] -# Numerical correctness demo — exercises the surrogate navigation dynamics and -# asserts goal-seeking / finiteness / determinism (needs torch). ~1s, no data. +# The unified command-line interface — one entry point for every workflow +# (selftest / obstacles / build-sdf / train / capture / pipeline / demo / eval). +grl-snam = "grl_snam.cli:main" +# Back-compat aliases (also reachable as `grl-snam selftest` / `grl-snam lab-demo`): grl-snam-selftest = "grl_snam.selftest:main" -# Standalone visual lab demo (needs pycvc/pycvc_gl + a GL context): opens a window, -# or `grl-snam-lab-demo out.png` renders offscreen. Inside volrover3 use the REPL -# instead: `import grl_snam_lab; grl_snam_lab.run_in_volrover()`. -grl-snam-lab-demo = "grl_snam_lab.demo:main" +grl-snam-lab-demo = "grl_snam.demos.lab:main" [tool.poetry] packages = [ { include = "grl_snam" }, - { include = "grl_snam_lab" }, { include = "scripts" }, { include = "experiments" }, { include = "src" }, + # Research core kept as top-level modules (imported by grl_snam + downstreams); + # sdf_nav.py was previously omitted here — a real packaging bug now fixed. { include = "train_coef_energy.py" }, { include = "eval_coef_energy.py" }, { include = "surrogate_robust.py" }, + { include = "sdf_nav.py" }, ] [tool.poetry.group.dev.dependencies] @@ -86,4 +88,5 @@ select = ["E", "F", "W", "I", "UP"] ignore = [ "E501", # line length handled by formatter "E741", # ambiguous variable names tolerated in math-heavy research code + "UP031", # printf-style %-format is idiomatic in the ported research/print code ] diff --git a/scripts/build_sdf.py b/scripts/build_sdf.py deleted file mode 100644 index 4318760..0000000 --- a/scripts/build_sdf.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Build a navigation SDF for a scene, from either source (configurable): - - --source edt (default) grl-snam's own exact footprint distance transform - (2-D, top-down, no extra deps beyond the occupancy) - --source cvc CVC's mesh-exact 3-D SDF via pycvc.sdf (SDF_V2), sliced - to the ground plane — leans on the CVC compute layer and - is the substrate for extending GRL-SNAM to 3-D later - -Both emit the same ``/nav_sdf.npz`` (normalized ``phi`` + unit ``normal_x/y`` -grids + meta), which ``train_sdf.py`` and the demo consume source-agnostically. - -Usage: - python scripts/build_sdf.py [--source edt|cvc] [--region 430] [--grid 512] -""" -from __future__ import annotations - -import argparse -import os - -import numpy as np - -from pycvc_gl.scenes import building_occupancy, terrain_grid - -import sdf_nav - -TARGET_EXTENT = 10.0 - - -def _gltf_mesh(glb_path): - """Flat (verts[x,y,z,...], tris[i,j,k,...]) from a glTF/GLB via VTK.""" - import vtkmodules.vtkRenderingOpenGL2 # noqa: F401 (register factories) - from vtkmodules.vtkIOGeometry import vtkGLTFReader - from vtkmodules.vtkFiltersGeometry import vtkCompositeDataGeometryFilter - from vtkmodules.util.numpy_support import vtk_to_numpy - - reader = vtkGLTFReader(); reader.SetFileName(glb_path); reader.Update() - geom = vtkCompositeDataGeometryFilter(); geom.SetInputConnection(reader.GetOutputPort()); geom.Update() - pd = geom.GetOutput() - pts = vtk_to_numpy(pd.GetPoints().GetData()).astype(np.float64) - polys = vtk_to_numpy(pd.GetPolys().GetData()) # [n0,i,j,k, n1,...]; triangulated -> n0==3 - tris = [] - i = 0 - while i < len(polys): - n = int(polys[i]) - if n == 3: - tris += [int(polys[i + 1]), int(polys[i + 2]), int(polys[i + 3])] - i += n + 1 - return pts.reshape(-1).tolist(), tris - - -def main() -> None: - ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - ap.add_argument("bundle_dir") - ap.add_argument("--source", choices=["edt", "cvc"], default="edt") - ap.add_argument("--region", type=float, default=430.0, help="working-region half-extent (world)") - ap.add_argument("--grid", type=int, default=512, help="2-D field resolution") - ap.add_argument("--cvc-dim", type=int, nargs=3, default=(512, 512, 48), help="cvc 3-D SDF dims") - ap.add_argument("-o", "--out", default=None) - args = ap.parse_args() - - terrain = os.path.join(args.bundle_dir, "terrain.json") - glb = os.path.join(args.bundle_dir, "buildings.glb") - _, bounds, _, _ = terrain_grid(terrain) - mnx, mny, mxx, mxy = bounds - cx, cy = 0.5 * (mnx + mxx), 0.5 * (mny + mxy) - S = TARGET_EXTENT / (2.0 * args.region) - - if args.source == "edt": - occ = building_occupancy(glb, bounds, args.grid, args.grid, inflate_m=0.0) - phi, nxg, nyg = sdf_nav.build_sdf(occ, bounds, S) - else: - verts, tris = _gltf_mesh(glb) - phi, nxg, nyg = sdf_nav.build_sdf_cvc(verts, tris, bounds, S, dim=tuple(args.cvc_dim)) - - out = args.out or os.path.join(args.bundle_dir, "nav_sdf.npz") - np.savez_compressed(out, phi=phi, normal_x=nxg, normal_y=nyg, - bounds=np.asarray(bounds, np.float32), center=np.asarray([cx, cy], np.float32), - scale=np.float32(S), region=np.float32(args.region), source=args.source) - print("built %s SDF %s (phi[%.2f,%.2f]) -> %s" - % (args.source, phi.shape, float(phi.min()), float(phi.max()), out)) - - -if __name__ == "__main__": - main() diff --git a/scripts/capture_drive_video.py b/scripts/capture_drive_video.py deleted file mode 100644 index a15ccee..0000000 --- a/scripts/capture_drive_video.py +++ /dev/null @@ -1,187 +0,0 @@ -"""Capture a learned SDF drive to an mp4: run the end-to-end navigator (with the -wall-follow local-minimum escape) from START to GOAL on a real scene, then render -the ACTUAL 3-D scene — terrain relief + glTF buildings + the vehicle — from a chase -camera OFFSCREEN via VTK (the same engine cvcGL/VolRover3 use), and ffmpeg the -frames into a video. - -This is how the demo video is produced without a live window: same geometry, same -renderer, a scripted chase camera. Needs the volrover env (pycvc_gl/VTK), torch, -the GRL-SNAM repo, a scene bundle, a trained SDF checkpoint, and ffmpeg. - - python scripts/capture_drive_video.py checkpoints/coef_sdf.pt \ - --start -361 114 --goal 185 50 -o drive.mp4 -""" -from __future__ import annotations - -import argparse -import math -import os -import subprocess - -import numpy as np -import torch - -import sdf_nav -from pycvc_gl.scenes import building_occupancy, terrain_grid - - -def navigate(field, occ0, bounds, model, meta, start, goal, maxst=3000): - """End-to-end SDF nav with wall-follow escape; returns the world trajectory.""" - S = meta["scale"]; cx, cy = meta["center"]; rr = meta["rr"] - mnx, mny, mxx, mxy = bounds; ny, nx = occ0.shape - kw = dict(rr=rr, d_hat=meta["d_hat"], dt=meta["dt"], vmax=meta["vmax"]) - - def w2n(p): return np.array([(p[0] - cx) * S, (p[1] - cy) * S], np.float32) - - def n2w(o): return np.array([o[0] / S + cx, o[1] / S + cy], np.float32) - - def nrm_at(on): - _, n = field.sample(torch.from_numpy(on).unsqueeze(0).float()); return n[0].numpy() - - o = torch.from_numpy(w2n(start)).unsqueeze(0); v = torch.zeros(1, 2); gn = w2n(goal) - tr = [np.asarray(start, np.float32)]; best = 1e9; stall = 0; mode = "seek"; turn = 1.0; dhit = 0.0 - with torch.no_grad(): - for _ in range(maxst): - p = o[0].numpy(); dg = float(np.linalg.norm(p - gn)); gdir = (gn - p) / (dg + 1e-6) - if dg < best - 1e-3: - best = dg; stall = 0 - else: - stall += 1 - if mode == "seek" and stall > 70: - t = np.array([-nrm_at(p)[1], nrm_at(p)[0]], np.float32) - turn = 1.0 if np.dot(t, gdir) >= 0 else -1.0; dhit = dg; mode = "wall"; stall = 0 - if mode == "wall": - n = nrm_at(p); t = turn * np.array([-n[1], n[0]], np.float32) - carrot = (p + (0.6 * t + 0.4 * n) * 1.6).astype(np.float32) - if dg < dhit - 1.2 or stall > 240: - mode = "seek"; best = dg; stall = 0 - else: - carrot = (p + gdir * min(1.8, dg)).astype(np.float32) - al, be, ga = model(sdf_nav.coef_feats(field, o, torch.from_numpy(carrot).unsqueeze(0))) - o, v, _ = sdf_nav.sdf_rollout(field, o, v, torch.from_numpy(carrot).unsqueeze(0), al, be, ga, 1, - nsub=meta["nsub"], **kw) - tr.append(n2w(o[0].numpy())); - if dg < 0.4: - break - return np.asarray(tr) - - -def render_video(bundle, traj, start, goal, out_mp4, frames=170, size=(960, 540)): - import vtkmodules.vtkRenderingOpenGL2 # noqa: F401 - from vtkmodules.vtkCommonCore import vtkPoints - from vtkmodules.vtkCommonDataModel import vtkPolyData, vtkCellArray - from vtkmodules.vtkCommonTransforms import vtkTransform - from vtkmodules.vtkFiltersCore import vtkPolyDataNormals - from vtkmodules.vtkFiltersSources import vtkCubeSource, vtkConeSource - from vtkmodules.vtkFiltersGeometry import vtkCompositeDataGeometryFilter - from vtkmodules.vtkIOGeometry import vtkGLTFReader - from vtkmodules.vtkIOImage import vtkPNGWriter - from vtkmodules.vtkRenderingCore import (vtkRenderer, vtkRenderWindow, vtkActor, vtkPolyDataMapper, - vtkWindowToImageFilter, vtkLight) - import json - W, H = size - fr = os.path.join(os.path.dirname(out_mp4) or ".", "_frames"); os.makedirs(fr, exist_ok=True) - os.system("rm -f %s/f_*.png" % fr) - d = json.load(open(os.path.join(bundle, "terrain.json"))); b = d["bounds"] - grid = list(reversed(d["grid"])); rows, cols = d["rows"], d["cols"] - mnx, mny, mxx, mxy = b["min_x"], b["min_y"], b["max_x"], b["max_y"] - dx = (mxx - mnx) / (cols - 1); dy = (mxy - mny) / (rows - 1) - - def hsamp(x, y): - fx = min(max((x - mnx) / dx, 0), cols - 1); fy = min(max((y - mny) / dy, 0), rows - 1) - c0, r0 = int(fx), int(fy); c1 = min(c0 + 1, cols - 1); r1 = min(r0 + 1, rows - 1) - tx, ty = fx - c0, fy - r0 - return ((grid[r0][c0] * (1 - tx) + grid[r0][c1] * tx) * (1 - ty) - + (grid[r1][c0] * (1 - tx) + grid[r1][c1] * tx) * ty) - - pts = vtkPoints(); pts.SetNumberOfPoints(rows * cols) - for r in range(rows): - for c in range(cols): - pts.SetPoint(r * cols + c, mnx + c * dx, mny + r * dy, float(grid[r][c])) - tris = vtkCellArray() - for r in range(rows - 1): - for c in range(cols - 1): - v = r * cols + c - for cell in ((v, v + 1, v + cols), (v + 1, v + cols + 1, v + cols)): - tris.InsertNextCell(3); [tris.InsertCellPoint(i) for i in cell] - tpd = vtkPolyData(); tpd.SetPoints(pts); tpd.SetPolys(tris) - tn = vtkPolyDataNormals(); tn.SetInputData(tpd); tn.Update() - tm = vtkPolyDataMapper(); tm.SetInputConnection(tn.GetOutputPort()); tm.ScalarVisibilityOff() - terrain = vtkActor(); terrain.SetMapper(tm) - terrain.GetProperty().SetColor(0.33, 0.39, 0.27); terrain.GetProperty().SetAmbient(0.28) - - gr = vtkGLTFReader(); gr.SetFileName(os.path.join(bundle, "buildings.glb")); gr.Update() - gf = vtkCompositeDataGeometryFilter(); gf.SetInputConnection(gr.GetOutputPort()); gf.Update() - bm = vtkPolyDataMapper(); bm.SetInputData(gf.GetOutput()); bm.SetStatic(1); bm.ScalarVisibilityOff() - bld = vtkActor(); bld.SetMapper(bm) - bp = bld.GetProperty(); bp.SetColor(0.72, 0.72, 0.77); bp.SetAmbient(0.42); bp.SetDiffuse(0.72) - - cs = vtkCubeSource(); cs.SetXLength(4.6); cs.SetYLength(2.0); cs.SetZLength(1.6); cs.Update() - vm = vtkPolyDataMapper(); vm.SetInputConnection(cs.GetOutputPort()) - veh = vtkActor(); veh.SetMapper(vm); veh.GetProperty().SetColor(0.90, 0.13, 0.12); veh.GetProperty().SetAmbient(0.4) - cn = vtkConeSource(); cn.SetHeight(16); cn.SetRadius(3.2); cn.SetResolution(20); cn.SetDirection(0, 0, -1); cn.Update() - bcm = vtkPolyDataMapper(); bcm.SetInputConnection(cn.GetOutputPort()) - beacon = vtkActor(); beacon.SetMapper(bcm); beacon.GetProperty().SetColor(0.98, 0.32, 0.12); beacon.GetProperty().SetAmbient(0.9) - - def pillar(x, y, col): - c = vtkCubeSource(); c.SetXLength(6); c.SetYLength(6); c.SetZLength(60); c.Update() - m = vtkPolyDataMapper(); m.SetInputConnection(c.GetOutputPort()) - a = vtkActor(); a.SetMapper(m); a.GetProperty().SetColor(*col); a.GetProperty().SetOpacity(0.55) - tf = vtkTransform(); tf.Translate(x, y, hsamp(x, y) + 30); a.SetUserTransform(tf); return a - - ren = vtkRenderer(); ren.SetBackground(0.16, 0.19, 0.13); ren.SetBackground2(0.55, 0.68, 0.82); ren.GradientBackgroundOn() - for a in (terrain, bld, veh, beacon, pillar(*start, (0.15, 0.85, 0.25)), pillar(*goal, (0.95, 0.78, 0.10))): - ren.AddActor(a) - sun = vtkLight(); sun.SetPosition(mnx, mny, 2500); sun.SetFocalPoint((mnx + mxx) / 2, (mny + mxy) / 2, 0) - sun.SetIntensity(0.9); sun.SetLightTypeToSceneLight(); ren.AddLight(sun) - rw = vtkRenderWindow(); rw.SetOffScreenRendering(1); rw.AddRenderer(ren); rw.SetSize(W, H) - cam = ren.GetActiveCamera(); cam.SetClippingRange(1.0, 9000.0) - - T = traj.copy() - for i in range(len(T)): - lo = max(0, i - 4); T[i] = traj[lo:i + 5].mean(0) - idx = np.linspace(0, len(T) - 2, frames).astype(int) - w2i = vtkWindowToImageFilter(); w2i.SetInput(rw) - head = np.array([1.0, 0.0]); BACK, HT, AHEAD = 74.0, 52.0, 24.0 - for fi, i in enumerate(idx): - p = T[i]; dvec = T[min(i + 3, len(T) - 1)] - p; nn = np.linalg.norm(dvec) - if nn > 1e-3: - head = 0.7 * head + 0.3 * (dvec / nn); head /= (np.linalg.norm(head) + 1e-9) - z = hsamp(p[0], p[1]) + 0.9 - tf = vtkTransform(); tf.Translate(p[0], p[1], z + 0.8) - tf.RotateZ(math.degrees(math.atan2(head[1], head[0]))); veh.SetUserTransform(tf) - bt = vtkTransform(); bt.Translate(p[0], p[1], z + 20.0); beacon.SetUserTransform(bt) - cam.SetPosition(p[0] - head[0] * BACK, p[1] - head[1] * BACK, z + HT) - cam.SetFocalPoint(p[0] + head[0] * AHEAD, p[1] + head[1] * AHEAD, z - 2.0); cam.SetViewUp(0, 0, 1) - rw.Render(); w2i.Modified(); w2i.Update() - wr = vtkPNGWriter(); wr.SetFileName("%s/f_%04d.png" % (fr, fi)); wr.SetInputConnection(w2i.GetOutputPort()); wr.Write() - subprocess.run(["ffmpeg", "-y", "-framerate", "30", "-i", "%s/f_%%04d.png" % fr, - "-c:v", "libx264", "-pix_fmt", "yuv420p", "-vf", "scale=%d:%d" % (W, H), out_mp4], check=True) - - -def main(): - ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - ap.add_argument("bundle"); ap.add_argument("checkpoint") - ap.add_argument("--sdf", default=None, help="prebuilt nav_sdf.npz (else built from occupancy)") - ap.add_argument("--start", type=float, nargs=2, required=True) - ap.add_argument("--goal", type=float, nargs=2, required=True) - ap.add_argument("--frames", type=int, default=170) - ap.add_argument("-o", "--out", default="drive.mp4") - args = ap.parse_args() - ck = torch.load(args.checkpoint, map_location="cpu"); meta = ck["meta"] - model = sdf_nav.CoefMLP(); model.load_state_dict(ck["model_state_dict"]); model.eval() - bounds = terrain_grid(os.path.join(args.bundle, "terrain.json"))[1] - occ0 = building_occupancy(os.path.join(args.bundle, "buildings.glb"), bounds, 512, 512, inflate_m=0.0) - if args.sdf and os.path.exists(args.sdf): - d = np.load(args.sdf); phi, nxg, nyg = d["phi"], d["normal_x"], d["normal_y"] - else: - phi, nxg, nyg = sdf_nav.build_sdf(occ0, bounds, meta["scale"]) - field = sdf_nav.SDFField(phi, nxg, nyg, bounds, meta["center"], meta["scale"], device="cpu") - traj = navigate(field, occ0, bounds, model, meta, np.asarray(args.start), np.asarray(args.goal)) - print("navigated %d steps; rendering %d frames -> %s" % (len(traj), args.frames, args.out)) - render_video(args.bundle, traj, args.start, args.goal, args.out, frames=args.frames) - print("wrote %s (%d bytes)" % (args.out, os.path.getsize(args.out))) - - -if __name__ == "__main__": - main() diff --git a/scripts/capture_multigoal_video.py b/scripts/capture_multigoal_video.py deleted file mode 100644 index 816132e..0000000 --- a/scripts/capture_multigoal_video.py +++ /dev/null @@ -1,292 +0,0 @@ -"""Capture the DYNAMIC multi-goal free-drive to an mp4: the learned SDF policy starts -near the centre of a real scene and heads for a goal that is RE-TARGETED live through -several points toward the corners — no precomputed route. Renders the actual 3-D scene -(terrain relief + glTF buildings + vehicle + shaded goal spires) from a drone chase -camera OFFSCREEN via VTK (the engine cvcGL/VolRover3 use), and ffmpeg's it to video. - -Two robustness details make the drive clean (no circling, no far-corner wandering): - * goals are chosen in the SDF's OWN free space (so selection and navigation agree — - every goal is actually reachable), corner-ward but reachable, in a perimeter loop; - * each leg is rendered only up to its CLOSEST APPROACH to the goal, so any - unproductive orbit/back-track after the last real progress is never shown (and the - goal simply re-targets — on-narrative for a dynamic free-drive). - -Needs the volrover env (pycvc_gl/VTK), torch, the GRL-SNAM repo, a scene bundle -(terrain.json + buildings.glb), a trained SDF checkpoint, and ffmpeg. - - python scripts/capture_multigoal_video.py checkpoints/coef_sdf.pt \ - --sdf /nav_sdf.npz --minutes 3 -o multigoal.mp4 -""" -from __future__ import annotations - -import argparse -import json -import math -import os -import subprocess - -import numpy as np -import torch - -import sdf_nav -from pycvc_gl.scenes import building_occupancy, terrain_grid - - -def _leg(field, model, meta, o, v, goal_w, kw, maxst=1300): - """Drive toward goal_w; return the state and trajectory truncated at the CLOSEST - approach to the goal (so an orbit/back-track tail is never rendered).""" - S = meta["scale"]; cx, cy = meta["center"] - - def w2n(p): return np.array([(p[0] - cx) * S, (p[1] - cy) * S], np.float32) - - def n2w(on): return np.array([on[0] / S + cx, on[1] / S + cy], np.float32) - - gn = w2n(goal_w); best = 1e9; stall = 0; mode = "seek"; turn = 1.0; dhit = 0.0 - out = []; best_o = o.clone(); best_v = v.clone(); besti = 0; bg = 1e9; since = 0 - with torch.no_grad(): - for _ in range(maxst): - p = o[0].numpy(); dg = float(np.linalg.norm(p - gn)); gdir = (gn - p) / (dg + 1e-6) - if dg < best - 1e-3: - best = dg; stall = 0 - else: - stall += 1 - if mode == "seek" and stall > 70: - _, nn = field.sample(torch.from_numpy(p).unsqueeze(0)); nn = nn[0].numpy() - t = np.array([-nn[1], nn[0]], np.float32) - turn = 1.0 if np.dot(t, gdir) >= 0 else -1.0; dhit = dg; mode = "wall"; stall = 0 - if mode == "wall": - _, nn = field.sample(torch.from_numpy(p).unsqueeze(0)); nn = nn[0].numpy() - t = turn * np.array([-nn[1], nn[0]], np.float32) - carrot = (p + (0.6 * t + 0.4 * nn) * 1.6).astype(np.float32) - if dg < dhit - 1.2 or stall > 240: - mode = "seek"; best = dg; stall = 0 - else: - carrot = (p + gdir * min(1.8, dg)).astype(np.float32) - al, be, ga = model(sdf_nav.coef_feats(field, o, torch.from_numpy(carrot).unsqueeze(0))) - o, v, _ = sdf_nav.sdf_rollout(field, o, v, torch.from_numpy(carrot).unsqueeze(0), al, be, ga, 1, - nsub=meta["nsub"], **kw) - out.append(n2w(o[0].numpy())) - pg = float(np.linalg.norm(o[0].numpy() - gn)) - if pg < bg - 1e-3: - bg = pg; besti = len(out) - 1; best_o = o.clone(); best_v = v.clone(); since = 0 - else: - since += 1 - if pg < 0.8: - besti = len(out) - 1; best_o = o.clone(); best_v = v.clone(); break - if since > 340: - break - return best_o, best_v, out[: besti + 1] - - -def select_goals(field, meta, bounds, clear_m=16.0): - """Pick 4 corner-ward goals in the SDF's free space, chained from an open centre - cell in a perimeter loop (NE, SE, SW, NW) so consecutive legs are short/reachable.""" - S = meta["scale"]; cx, cy = meta["center"]; region = meta["region"] - mnx, mny, mxx, mxy = bounds - kw = dict(rr=meta["rr"], d_hat=meta["d_hat"], dt=meta["dt"], vmax=meta["vmax"]) - - def w2n(p): return np.array([(p[0] - cx) * S, (p[1] - cy) * S], np.float32) - - def clearance(pw): - d, _ = field.sample(torch.from_numpy(w2n(pw)).unsqueeze(0)); return float(d[0]) / S - - start = None - for rad in (0, 20, 40, 60, 80, 120): - for a in range(0, 360, 30): - pw = np.array([cx + math.cos(math.radians(a)) * rad, cy + math.sin(math.radians(a)) * rad], np.float32) - if clearance(pw) > clear_m: - start = pw; break - if start is not None: - break - - def candidates(sx, sy): - diag = math.atan2(sy, sx); out = [] - for frac in (0.62, 0.56, 0.50, 0.44, 0.38, 0.32): - for da in (0, -18, 18, -32, 32): - ang = diag + math.radians(da); r = region * frac - gw = np.array([cx + math.cos(ang) * r, cy + math.sin(ang) * r], np.float32) - if mnx < gw[0] < mxx and mny < gw[1] < mxy and clearance(gw) > clear_m: - out.append(gw) - return out - - o = torch.from_numpy(w2n(start)).unsqueeze(0); v = torch.zeros(1, 2); goals = [] - for sx, sy in ((1, 1), (1, -1), (-1, -1), (-1, 1)): - chosen = None; chosen_ba = 1e9; chosen_state = None - for gw in candidates(sx, sy)[:8]: - bo, bv, seg = _leg(field, model_ref["m"], meta, o.clone(), v.clone(), gw, kw, maxst=800) - ba = float(np.linalg.norm(bo[0].numpy() - w2n(gw))) / S - if ba < chosen_ba: - chosen_ba = ba; chosen = gw; chosen_state = (bo, bv) - if ba < 70: - chosen = gw; chosen_state = (bo, bv); break - o, v = chosen_state; goals.append(chosen) - return start, goals - - -model_ref = {"m": None} # module-level handle so select_goals can reuse the loaded model - - -def render_video(bundle, legs, goals, out_mp4, minutes=3.0, fps=15, size=(960, 540)): - import vtkmodules.vtkRenderingOpenGL2 # noqa: F401 - from vtkmodules.vtkCommonCore import vtkPoints - from vtkmodules.vtkCommonDataModel import vtkPolyData, vtkCellArray - from vtkmodules.vtkCommonTransforms import vtkTransform - from vtkmodules.vtkFiltersCore import vtkPolyDataNormals, vtkQuadricDecimation - from vtkmodules.vtkFiltersSources import vtkCubeSource, vtkConeSource - from vtkmodules.vtkFiltersGeometry import vtkCompositeDataGeometryFilter - from vtkmodules.vtkIOGeometry import vtkGLTFReader - from vtkmodules.vtkIOImage import vtkPNGWriter - from vtkmodules.vtkRenderingCore import (vtkRenderer, vtkRenderWindow, vtkActor, vtkPolyDataMapper, - vtkWindowToImageFilter, vtkLight) - W, H = size - fr = os.path.join(os.path.dirname(out_mp4) or ".", "_frames"); os.makedirs(fr, exist_ok=True) - os.system("rm -f %s/f_*.png" % fr) - d = json.load(open(os.path.join(bundle, "terrain.json"))); b = d["bounds"] - grid = list(reversed(d["grid"])); rows, cols = d["rows"], d["cols"] - mnx, mny, mxx, mxy = b["min_x"], b["min_y"], b["max_x"], b["max_y"] - dx = (mxx - mnx) / (cols - 1); dy = (mxy - mny) / (rows - 1) - - def hsamp(x, y): - fx = min(max((x - mnx) / dx, 0), cols - 1); fy = min(max((y - mny) / dy, 0), rows - 1) - c0, r0 = int(fx), int(fy); c1 = min(c0 + 1, cols - 1); r1 = min(r0 + 1, rows - 1) - tx, ty = fx - c0, fy - r0 - return ((grid[r0][c0] * (1 - tx) + grid[r0][c1] * tx) * (1 - ty) - + (grid[r1][c0] * (1 - tx) + grid[r1][c1] * tx) * ty) - - pts = vtkPoints(); pts.SetNumberOfPoints(rows * cols) - for r in range(rows): - for c in range(cols): - pts.SetPoint(r * cols + c, mnx + c * dx, mny + r * dy, float(grid[r][c])) - tris = vtkCellArray() - for r in range(rows - 1): - for c in range(cols - 1): - vv = r * cols + c - for cell in ((vv, vv + 1, vv + cols), (vv + 1, vv + cols + 1, vv + cols)): - tris.InsertNextCell(3); [tris.InsertCellPoint(i) for i in cell] - tpd = vtkPolyData(); tpd.SetPoints(pts); tpd.SetPolys(tris) - tn = vtkPolyDataNormals(); tn.SetInputData(tpd); tn.Update() - tm = vtkPolyDataMapper(); tm.SetInputConnection(tn.GetOutputPort()); tm.ScalarVisibilityOff() - terrain = vtkActor(); terrain.SetMapper(tm) - terrain.GetProperty().SetColor(0.33, 0.39, 0.27); terrain.GetProperty().SetAmbient(0.25) - - gr = vtkGLTFReader(); gr.SetFileName(os.path.join(bundle, "buildings.glb")); gr.Update() - gf = vtkCompositeDataGeometryFilter(); gf.SetInputConnection(gr.GetOutputPort()); gf.Update() - dec = vtkQuadricDecimation(); dec.SetInputData(gf.GetOutput()); dec.SetTargetReduction(0.65); dec.Update() - bm = vtkPolyDataMapper(); bm.SetInputData(dec.GetOutput()); bm.SetStatic(1); bm.ScalarVisibilityOff() - bld = vtkActor(); bld.SetMapper(bm) - bp = bld.GetProperty(); bp.SetColor(0.72, 0.72, 0.77); bp.SetAmbient(0.4); bp.SetDiffuse(0.72) - - cs = vtkCubeSource(); cs.SetXLength(4.6); cs.SetYLength(2.0); cs.SetZLength(1.6); cs.Update() - vm = vtkPolyDataMapper(); vm.SetInputConnection(cs.GetOutputPort()) - veh = vtkActor(); veh.SetMapper(vm) - veh.GetProperty().SetColor(0.90, 0.13, 0.12); veh.GetProperty().SetAmbient(0.3); veh.GetProperty().SetDiffuse(0.8) - - def shaded_cone(h, rad, col, down=True): - c = vtkConeSource(); c.SetHeight(h); c.SetRadius(rad); c.SetResolution(26) - c.SetDirection(0, 0, -1 if down else 1); c.Update() - n = vtkPolyDataNormals(); n.SetInputConnection(c.GetOutputPort()); n.Update() - m = vtkPolyDataMapper(); m.SetInputConnection(n.GetOutputPort()) - a = vtkActor(); a.SetMapper(m); p = a.GetProperty() - p.SetColor(*col); p.SetAmbient(0.22); p.SetDiffuse(0.85); p.SetSpecular(0.25); p.SetSpecularPower(18) - return a - - veh_beacon = shaded_cone(24, 5.0, (0.98, 0.42, 0.10), down=True) - goal_cols = [(0.20, 0.80, 0.35), (0.30, 0.65, 0.95), (0.95, 0.75, 0.15), (0.85, 0.35, 0.85)] - goal_beacons = [shaded_cone(150, 9.0, goal_cols[i % len(goal_cols)], down=False) for i in range(len(goals))] - - ren = vtkRenderer(); ren.SetBackground(0.16, 0.19, 0.13); ren.SetBackground2(0.55, 0.68, 0.82); ren.GradientBackgroundOn() - ren.AddActor(terrain); ren.AddActor(bld); ren.AddActor(veh); ren.AddActor(veh_beacon) - for gi, gb in enumerate(goal_beacons): - t = vtkTransform(); t.Translate(goals[gi][0], goals[gi][1], hsamp(*goals[gi]) + 74) - gb.SetUserTransform(t); ren.AddActor(gb) - sun = vtkLight(); sun.SetPosition(mnx - 500, mny - 500, 2500); sun.SetFocalPoint((mnx + mxx) / 2, (mny + mxy) / 2, 0) - sun.SetIntensity(1.0); sun.SetLightTypeToSceneLight(); ren.AddLight(sun) - rw = vtkRenderWindow(); rw.SetOffScreenRendering(1); rw.AddRenderer(ren); rw.SetSize(W, H) - cam = ren.GetActiveCamera(); cam.SetClippingRange(1.0, 9000.0) - - def smooth(a, k=5): - a = np.asarray(a, np.float32); o = a.copy() - for i in range(len(a)): - o[i] = a[max(0, i - k): i + k + 1].mean(0) - return o - - # Pace by ARC LENGTH so speed is constant (fixes 'slow then fast'); a per-leg - # sin ease decelerates into each goal and accelerates out. Total frames are set - # to `minutes` so the film runs the requested length at a steady cruise. - total_frames = int(minutes * 60 * fps) - lens = [] - smoothed = [] - for gi, seg in legs: - segS = smooth(seg); dd = np.linalg.norm(np.diff(segS, axis=0), axis=1) - arcl = np.concatenate([[0.0], np.cumsum(dd)]); smoothed.append((gi, segS, arcl)); lens.append(float(arcl[-1])) - tot = sum(lens) or 1.0 - seq = [] - for (gi, segS, arcl), L in zip(smoothed, lens): - if L < 1.0: - continue - n = max(120, int(round(total_frames * L / tot))) - u = np.linspace(0, 1, n); spd = 0.35 + 0.65 * np.sin(np.pi * u) - cum = np.cumsum(spd); cum = cum / cum[-1] * L - xs = np.interp(cum, arcl, segS[:, 0]); ys = np.interp(cum, arcl, segS[:, 1]) - for k in range(n): - seq.append((float(xs[k]), float(ys[k]), gi)) - - w2i = vtkWindowToImageFilter(); w2i.SetInput(rw) - head = np.array([1.0, 0.0]); BACK, HT, AHEAD = 74.0, 52.0, 24.0 - prevp = np.array([seq[0][0], seq[0][1]]) - for fi, (px, py, ag) in enumerate(seq): - p = np.array([px, py]); dvec = p - prevp; prevp = p; nn = np.linalg.norm(dvec) - if nn > 1e-4: - head = 0.6 * head + 0.4 * (dvec / nn); head /= (np.linalg.norm(head) + 1e-9) - zt = hsamp(px, py) + 0.9 - tf = vtkTransform(); tf.Translate(px, py, zt + 0.8) - tf.RotateZ(math.degrees(math.atan2(head[1], head[0]))); veh.SetUserTransform(tf) - bt = vtkTransform(); bt.Translate(px, py, zt + 24.0); veh_beacon.SetUserTransform(bt) - for gi, gb in enumerate(goal_beacons): - gb.GetProperty().SetAmbient(0.62 if gi == ag else 0.22) # active goal spire glows - cam.SetPosition(px - head[0] * BACK, py - head[1] * BACK, zt + HT) - cam.SetFocalPoint(px + head[0] * AHEAD, py + head[1] * AHEAD, zt - 2.0); cam.SetViewUp(0, 0, 1) - rw.Render(); w2i.Modified(); w2i.Update() - wr = vtkPNGWriter(); wr.SetFileName("%s/f_%05d.png" % (fr, fi)); wr.SetInputConnection(w2i.GetOutputPort()); wr.Write() - subprocess.run(["ffmpeg", "-y", "-framerate", str(fps), "-i", "%s/f_%%05d.png" % fr, - "-c:v", "libx264", "-pix_fmt", "yuv420p", "-vf", "scale=%d:%d" % (W, H), out_mp4], check=True) - return len(seq) - - -def main(): - ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - ap.add_argument("bundle"); ap.add_argument("checkpoint") - ap.add_argument("--sdf", default=None, help="prebuilt nav_sdf.npz (else built from occupancy)") - ap.add_argument("--minutes", type=float, default=3.0, help="target film length") - ap.add_argument("--fps", type=int, default=15) - ap.add_argument("-o", "--out", default="multigoal.mp4") - args = ap.parse_args() - ck = torch.load(args.checkpoint, map_location="cpu"); meta = ck["meta"] - model = sdf_nav.CoefMLP(); model.load_state_dict(ck["model_state_dict"]); model.eval() - model_ref["m"] = model - bounds = terrain_grid(os.path.join(args.bundle, "terrain.json"))[1] - occ0 = building_occupancy(os.path.join(args.bundle, "buildings.glb"), bounds, 512, 512, inflate_m=0.0) - if args.sdf and os.path.exists(args.sdf): - dd = np.load(args.sdf); phi, nxg, nyg = dd["phi"], dd["normal_x"], dd["normal_y"] - else: - phi, nxg, nyg = sdf_nav.build_sdf(occ0, bounds, meta["scale"]) - field = sdf_nav.SDFField(phi, nxg, nyg, bounds, meta["center"], meta["scale"], device="cpu") - kw = dict(rr=meta["rr"], d_hat=meta["d_hat"], dt=meta["dt"], vmax=meta["vmax"]) - - start, goals = select_goals(field, meta, bounds) - print("start %s goals %s" % (np.round(start).tolist(), [np.round(g).tolist() for g in goals])) - S = meta["scale"]; cx, cy = meta["center"] - - def w2n(p): return np.array([(p[0] - cx) * S, (p[1] - cy) * S], np.float32) - - o = torch.from_numpy(w2n(start)).unsqueeze(0); v = torch.zeros(1, 2); legs = [] - for gi, g in enumerate(goals): - o, v, seg = _leg(field, model, meta, o, v, g, kw) - pts = ([start.tolist()] if gi == 0 else []) + [p.tolist() for p in seg] - legs.append((gi, np.asarray(pts, np.float32))) - n = render_video(args.bundle, legs, goals, args.out, minutes=args.minutes, fps=args.fps) - print("wrote %s (%d bytes, %d frames, %.0fs)" % (args.out, os.path.getsize(args.out), n, n / args.fps)) - - -if __name__ == "__main__": - main() diff --git a/scripts/check_dataset.py b/scripts/check_dataset.py deleted file mode 100644 index 2ba60bc..0000000 --- a/scripts/check_dataset.py +++ /dev/null @@ -1,208 +0,0 @@ -# plot_dataset_trajectories.py -import os -import random -from typing import List, Optional, Sequence -import numpy as np - -import torch -import matplotlib -matplotlib.use("Agg") # headless-friendly -import matplotlib.pyplot as plt -from matplotlib.patches import Circle - -def _as_tensor(x, dtype=torch.float32): - return x if torch.is_tensor(x) else torch.as_tensor(x, dtype=dtype) - -def _get_xy(t: torch.Tensor) -> torch.Tensor: - """ - Take first two dims as (x,y). Supports [T,2+] or [2+]. - Returns shape [T,2]. - """ - t = _as_tensor(t).detach().cpu().float() - if t.ndim == 1: - if t.shape[0] < 2: - raise ValueError("Need at least 2D to plot (x,y).") - return t[:2].unsqueeze(0) - if t.ndim >= 2: - if t.shape[-1] < 2: - raise ValueError(f"Last dim < 2: {tuple(t.shape)}") - return t[..., :2] - raise ValueError(f"Unexpected tensor shape: {tuple(t.shape)}") - -def _draw_obstacles(ax, centers: torch.Tensor, radii: torch.Tensor, - facecolor=(0.95,0.4,0.4), edgecolor=(0.35,0.05,0.05), alpha=0.25): - centers = _as_tensor(centers).cpu().float() - radii = _as_tensor(radii).cpu().float() - if centers.ndim != 2 or centers.shape[-1] != 2: - raise ValueError(f"obstacle_centers should be (C,2), got {tuple(centers.shape)}") - if radii.ndim != 1 or radii.shape[0] != centers.shape[0]: - raise ValueError("obstacle_radii should be (C,) matching centers.") - for (cx, cy), r in zip(centers.numpy(), radii.numpy()): - circ = Circle((cx, cy), r, facecolor=facecolor, edgecolor=edgecolor, linewidth=1.2, alpha=alpha) - ax.add_patch(circ) - -def _compute_limits(items: Sequence[torch.Tensor], pad_ratio: float = 0.08): - xs, ys = [], [] - for it in items: - a = _as_tensor(it).cpu().float() - if a.ndim == 1 and a.numel() >= 2: - xs.append(a[0].item()); ys.append(a[1].item()) - elif a.ndim == 2 and a.shape[-1] >= 2: - xs.extend(a[:, 0].tolist()); ys.extend(a[:, 1].tolist()) - if not xs: - return (-1, 1), (-1, 1) - xmin, xmax = min(xs), max(xs) - ymin, ymax = min(ys), max(ys) - dx = max(1e-6, xmax - xmin) - dy = max(1e-6, ymax - ymin) - px, py = dx * pad_ratio, dy * pad_ratio - return (xmin - px, xmax + px), (ymin - py, ymax + py) - -def plot_single_trajectory( - traj: dict, - ax=None, - show_velocity_every: int = 5, - velocity_scale: float = 1.0, - path_kwargs: Optional[dict] = None, - quiver_kwargs: Optional[dict] = None, - show_goal: bool = True, - show_start: bool = True, - draw_obstacles: bool = True, -): - """ - Plot one trajectory dict from the dataset on an axes. - """ - if ax is None: - _, ax = plt.subplots(figsize=(6, 6)) - - # ---- meta - goal = _get_xy(traj["goal_position"]) - centers = traj.get("obstacle_centers", None) - radii = traj.get("obstacle_radii", None) - if draw_obstacles and centers is not None and radii is not None: - _draw_obstacles(ax, centers, radii) - - # ---- time series: q_frame & p_frame - states: List[dict] = traj["trajectory_states"] - q_seq = torch.stack([_get_xy(s["q_frame"]).squeeze(0) for s in states], dim=0) # [T,2] - p_seq = torch.stack([_get_xy(s["p_frame"]).squeeze(0) for s in states], dim=0) # [T,2] - - # ---- path - pk = dict(color="tab:blue", linewidth=2.0, alpha=0.95) - if path_kwargs: pk.update(path_kwargs) - ax.plot(q_seq[:, 0].numpy(), q_seq[:, 1].numpy(), **pk) - - # ---- start & goal - if show_start: - ax.scatter(q_seq[0, 0].item(), q_seq[0, 1].item(), marker="o", s=40, color="tab:blue", zorder=3, label="start") - if show_goal: - ax.scatter(goal[0, 0].item(), goal[0, 1].item(), marker="*", s=120, color="tab:green", edgecolor="k", zorder=4, label="goal") - - # ---- velocity arrows - if show_velocity_every is not None and show_velocity_every > 0: - idx = torch.arange(0, q_seq.shape[0], show_velocity_every) - qv = q_seq[idx] - pv = p_seq[idx] * float(velocity_scale) - qk = dict(angles="xy", scale_units="xy", scale=1.0, width=0.004, alpha=0.9) - if quiver_kwargs: qk.update(quiver_kwargs) - ax.quiver(qv[:, 0].numpy(), qv[:, 1].numpy(), pv[:, 0].numpy(), pv[:, 1].numpy(), **qk) - - # ---- cosmetics - ax.set_aspect("equal", adjustable="datalim") - ax.set_xlabel("x") - ax.set_ylabel("y") - ax.grid(True, linestyle=":", linewidth=0.7, alpha=0.6) - - return ax, q_seq, p_seq - -def plot_sampled_trajectories( - pt_path: str, - num_trajs: int = 4, - seed: Optional[int] = 0, - velocity_every: int = 5, - velocity_scale: float = 1.0, - save_path: Optional[str] = None, - suptitle: Optional[str] = None, -): - """ - Draw num_trajs random trajectories with scene layout (obstacles + goal). - """ - data = torch.load(pt_path, map_location="cpu") - assert isinstance(data, list) and len(data) > 0, "Top-level should be a non-empty list of trajectories." - - rng = random.Random(seed) if seed is not None else random - idxs = list(range(len(data))) - rng.shuffle(idxs) - idxs = idxs[:min(num_trajs, len(data))] - - # Prepare grid - n = len(idxs) - ncols = min(3, n) - nrows = (n + ncols - 1) // ncols - fig, axes = plt.subplots(nrows, ncols, figsize=(6*ncols, 6*nrows)) - fig, axes = plt.subplots(nrows, ncols, figsize=(6*ncols, 6*nrows)) - - # Flatten axes to a simple list - if isinstance(axes, np.ndarray): - axes = axes.ravel().tolist() - else: - axes = [axes] - axes = [a for row in (axes if isinstance(axes, list) else [axes]) for a in (row if isinstance(row, (list, tuple)) else [row])] - axes = axes[:n] - - # First pass: plot and track limits - all_items_for_limits = [] - for ax, i in zip(axes, idxs): - traj = data[i] - ax.set_title(f"Trajectory {i} | success={traj.get('success', 'NA')}") - ax, q_seq, _ = plot_single_trajectory( - traj, ax=ax, - show_velocity_every=velocity_every, - velocity_scale=velocity_scale, - ) - all_items_for_limits.append(q_seq) - if "obstacle_centers" in traj: - all_items_for_limits.append(_get_xy(traj["obstacle_centers"])) - if "goal_position" in traj: - all_items_for_limits.append(_get_xy(traj["goal_position"]).squeeze(0)) - - # Harmonize axis limits across subplots - xlim, ylim = _compute_limits(all_items_for_limits, pad_ratio=0.12) - for ax in axes: - ax.set_xlim(*xlim); ax.set_ylim(*ylim) - - if suptitle: - fig.suptitle(suptitle) - fig.tight_layout(rect=(0, 0, 1, 0.98)) - - if save_path is not None: - os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True) - fig.savefig(save_path, dpi=160, bbox_inches="tight") - - return fig - -# ---- Optional: log to TensorBoard ---- -def log_sampled_trajectories_tb( - pt_path: str, - log_dir: str, - tag: str = "Eval/Trajectories", - global_step: int = 0, - **kwargs, -): - """Create the plot and push it to TensorBoard.""" - from torch.utils.tensorboard import SummaryWriter - fig = plot_sampled_trajectories(pt_path, **kwargs) - w = SummaryWriter(log_dir=log_dir) - w.add_figure(tag, fig, global_step=global_step, close=True) - w.flush(); w.close() - - -# Basic: draw 4 random trajectories and save -fig = plot_sampled_trajectories( - "./complete_dpo_navigation_dataset/complete_trajectories_torch.pt", - num_trajs=30, - velocity_every=6, # put arrows every 6 steps - velocity_scale=0.8, # shorten/lenghten arrows - save_path="sampled_trajs.png", - suptitle="Dataset samples", -) diff --git a/scripts/extract_obstacles.py b/scripts/extract_obstacles.py deleted file mode 100644 index 5060d5b..0000000 --- a/scripts/extract_obstacles.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Ingest a real-world geometry bundle into an obstacle set for navigation training. - -Turns a ``geometry_bundle`` (a ``terrain.json`` heightfield + a ``buildings.glb`` -city mesh) into a compact ``obstacles.npz`` the trainer consumes: the building -footprints become a set of CIRCULAR obstacles (the shape the GRL-SNAM surrogate -repels from), plus a pool of free (drivable) world points to sample start/goal -positions from. Run this ONCE per scene; the trainer then reads the ``.npz`` with -no graphics dependency, and a pre-generated obstacle set can be shipped as-is. - -This step needs the ``pycvc_gl`` scene helpers (they rasterize the city mesh with -VTK to a solid occupancy grid), so run it in an environment where volrover3's -Python / cvcGL bindings are importable. The rasterized occupancy is cached next to -the ``.glb`` by ``pycvc_gl.scenes.building_occupancy``, so re-runs are instant. - -Usage: - python scripts/extract_obstacles.py [-o obstacles.npz] [--grid 512] - -```` holds ``terrain.json`` and ``buildings.glb`` (e.g. a scene export). -""" -from __future__ import annotations - -import argparse -import os - -import numpy as np - -from pycvc_gl.scenes import building_occupancy, terrain_grid - - -def extract(bundle_dir: str, grid: int = 512, block: int = 4, radius_frac: float = 0.6, - robot_radius_world: float = 3.0): - """Bundle -> (obstacle centers, radius, free-point pool), all in WORLD units. - - The city footprint is rasterized to a ``grid x grid`` solid occupancy mask, - coarsened by ``block`` (so each obstacle stands in for a ``block x block`` - patch of wall), and every occupied coarse cell becomes one circular obstacle - at the cell center with radius ``radius_frac * cell_size``. Free cells become - the drivable pool. Nothing here is scaled to the surrogate's ~10-unit regime - yet — the trainer normalizes per working-region so one bundle can be trained - at several zooms. - """ - terrain = os.path.join(bundle_dir, "terrain.json") - glb = os.path.join(bundle_dir, "buildings.glb") - _, bounds, _, _ = terrain_grid(terrain) # (min_x, min_y, max_x, max_y) - mnx, mny, mxx, mxy = bounds - - # Solid top-down occupancy (True = inside a building). inflate_m=0: raw - # footprints; the trainer's clearance margin keeps the robot off the walls. - occ = building_occupancy(glb, bounds, grid, grid, inflate_m=0.0) - ny, nx = occ.shape - - cny, cnx = ny // block, nx // block - cocc = occ[: cny * block, : cnx * block].reshape(cny, block, cnx, block).any(axis=(1, 3)) - csx = (mxx - mnx) / cnx # coarse cell size (world), x - csy = (mxy - mny) / cny # ... y - ys, xs = np.where(cocc) - centers = np.stack([mnx + (xs + 0.5) * csx, mny + (ys + 0.5) * csy], 1).astype(np.float32) - radius_world = float(radius_frac * 0.5 * (csx + csy)) - - # Free-point pool over the SAME frame the terrain sampler uses (row 0 = min_y). - gx = np.linspace(mnx, mxx, nx, dtype=np.float32) - gy = np.linspace(mny, mxy, ny, dtype=np.float32) - freeR, freeC = np.where(~occ) - free_pool = np.stack([gx[freeC], gy[freeR]], 1).astype(np.float32) - - return { - "centers": centers, # [M,2] obstacle centers (world) - "radius_world": np.float32(radius_world), # scalar obstacle radius (world) - "robot_radius_world": np.float32(robot_radius_world), - "bounds": np.asarray(bounds, np.float32), # (min_x,min_y,max_x,max_y) - "free_pool": free_pool, # [K,2] drivable points (world) - "cell_size": np.float32(0.5 * (csx + csy)), - } - - -def main() -> None: - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument("bundle_dir", help="dir with terrain.json + buildings.glb") - ap.add_argument("-o", "--out", default=None, help="output .npz (default /obstacles.npz)") - ap.add_argument("--grid", type=int, default=512, help="occupancy raster resolution") - ap.add_argument("--block", type=int, default=2, - help="coarsen factor (obstacle granularity). Finer (2) gives ~7 m circles that " - "leave streets navigable; coarser (4) gives ~14 m circles that crowd them.") - args = ap.parse_args() - - data = extract(args.bundle_dir, grid=args.grid, block=args.block) - out = args.out or os.path.join(args.bundle_dir, "obstacles.npz") - np.savez_compressed(out, **data) - print( - "extracted %d obstacles (r=%.1fm) + %d free points over %s -> %s" - % (len(data["centers"]), float(data["radius_world"]), - len(data["free_pool"]), tuple(data["bounds"].tolist()), out) - ) - - -if __name__ == "__main__": - main() diff --git a/scripts/large_scale_ex.py b/scripts/large_scale_ex.py deleted file mode 100644 index cf58036..0000000 --- a/scripts/large_scale_ex.py +++ /dev/null @@ -1,738 +0,0 @@ -#!/usr/bin/env python3 -# One large-scale example + MP4 movie with frame updates. -# Assumes 'spline_stagewise4' (ssi) is available on PYTHONPATH. - -import sys, os, math, json -import numpy as np -import torch -import matplotlib -matplotlib.use("Agg") -import matplotlib.pyplot as plt -from matplotlib.animation import FFMpegWriter, writers -import matplotlib as mpl -mpl.rcParams['path.simplify'] = True -mpl.rcParams['path.simplify_threshold'] = 0.8 -mpl.rcParams['agg.path.chunksize'] = 10000 - -# If you need to add your scripts folder (adjust if different): -sys.path.append("/mnt/data/adityas/DPO/scripts") -import spline_stagewise4 as ssi - - -# ========= Base config you provided (trimmed replica for inheritance) ========= -class GenCfg: - out_root: str = "./nav_dataset_maximin" - seed: int = 2025 - - # simulation - total_steps: int = 2000 - dt: float = 0.03 - snapshot_every: int = 5 # we will also use this as video frame stride - - # start/goal and stage geometry - start = np.array([-1.0, -0.5], float) - goal = np.array([ 9.0, 0.2], float) - stage_size = (2.6, 2.0) - overlap = 0.30 - - # tube inflation (planning clearance) - inflate = 0.05 - - # robot base params - n_ctrl: int = 20 - K: int = 240 - lam_reg: float = 3.0 - mu: float = 0.25 - d_hat: float = 1.0 - seed_robot: int = 7 - radius: float = 0.35 - - # success thresholds - goal_tol: float = 0.30 - min_clear_tol: float = 1e-6 - - # obstacle field (Case 1 tight corridor generator) - env_x_bounds = (-0.5, 9.0) - env_y_bounds = (-1.8, 1.8) - radius_range = (0.28, 0.80) - weight_range = (0.6, 1.5) - n_obs_range = (12, 22) - poisson_max_tries: int = 6000 - min_separation_slack: float = 0.08 - min_start_clear: float = 0.9 - min_goal_clear: float = 0.9 - - # tight gap controls - gap_mult_range = (0.70, 0.80) - gap_margin_range = (0.01, 0.03) - corridor_wiggle_amp = 0.35 - corridor_wiggle_k = 0.6 - n_gates_range = (3, 5) - n_intruders_range = (2, 4) - intruder_depth_frac_range = (0.60, 0.95) - - # graph parameters - grid_res: int = 12 - knn_k: int = 7 - nsamp_edge: int = 20 - min_node_clear: float = 1e-3 - follow_corridor: bool = True - - # decoupled-search ROI - roi_pad_scale_x: float = 0.8 - roi_pad_scale_y: float = 0.8 - search_union_current_next: bool = True - - # hard extra margins beyond tube - min_edge_margin: float = 1e-3 - min_exit_margin: float = 1e-3 - - # viz - overlay_inflated_in_snapshot: bool = True - - -# ========= "Bigger scale, more obstacles; same core configs" ========= -class BigCfg(GenCfg): - # Output folder for this single example - out_root = "./nav_large_example" - - # Make the world ~4x longer and ~2x taller - start = np.array([-2.0, -1.2], float) - goal = np.array([ 32.0, 1.0], float) - env_x_bounds = (-3.0, 34.0) - env_y_bounds = (-4.0, 4.0) - - gap_mult_range = (0.55, 0.65) # lower => narrower nominal gap - gap_margin_range = (0.005, 0.02) # smaller extra margin - intruder_depth_frac_range = (0.85, 0.98) # intruders push deeper in - - # Stage window scaled up to cover longer path comfortably - stage_size = (5.2, 3.8) # ~2x original - overlap = 0.30 # keep same overlap fraction - - # Simulation: a bit longer to let it traverse - total_steps = 3600 - snapshot_every = 8 # render every 3 sim steps for smoother video (dt=0.03 => ≈10 fps) - dt = 0.03 - - # Much denser obstacle field - n_obs_range = (70, 120) - radius_range = (0.30, 1.20) - corridor_wiggle_amp = 0.60 # wider wiggle to force route finding - n_gates_range = (6, 9) - n_intruders_range = (5, 8) - - # Slightly denser graph for exit selection - grid_res = 20 - knn_k = 9 - nsamp_edge = 28 - - # Keep the planning tube inflation and robot params the same spirit - inflate = 0.05 - n_ctrl = 20 - K = 240 - lam_reg = 3.0 - mu = 0.25 - d_hat = 1.0 - radius = 0.35 - - # overlay - overlay_inflated_in_snapshot = True - - -# ========= Helpers (directly in this file so it’s one runnable script) ========= -def set_all_seeds(seed: int): - import random - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - -def rand_uniform(a: float, b: float) -> float: - return float(np.random.uniform(a, b)) - - -def sample_obstacles_case1_tight(cfg: GenCfg): - """Tight, solvable corridor with enforced narrow gates + intruders (scaled by cfg).""" - rng = np.random.default_rng() - - robot_diam = 2.0 * cfg.radius - gap_mult = rng.uniform(*cfg.gap_mult_range) - gap_margin = rng.uniform(*cfg.gap_margin_range) - target_gap = max(gap_mult * robot_diam + gap_margin, 2.0 * cfg.inflate + 0.01) - half_gap = 0.5 * target_gap - eps = 0.012 - corridor_half_w = half_gap + 0.01 - - p0 = cfg.start.astype(float); p1 = cfg.goal.astype(float) - dir01 = p1 - p0; dir01 = dir01 / (np.linalg.norm(dir01) + 1e-12) - ortho = np.array([-dir01[1], dir01[0]]) - amp = cfg.corridor_wiggle_amp - kfreq = cfg.corridor_wiggle_k - ts = np.linspace(0.0, 1.0, 120) - - def centerline_point(t): - base = (1 - t) * p0 + t * p1 - return base + amp * np.sin(2 * np.pi * kfreq * t) * ortho - - centerline = np.stack([centerline_point(t) for t in ts], axis=0) - - def dist_to_polyline(pt): - return float(np.linalg.norm(centerline - pt[None, :], axis=1).min()) - - def disk_outside_corridor(c, r): # forbid clutter inside corridor - return dist_to_polyline(c) >= (corridor_half_w + r) - - def non_overlap(c, r, Cs, Rs, slack): - if Cs.size == 0: return True - d2 = np.sum((Cs - c[None, :])**2, axis=1) - min_allow2 = (Rs + r + slack)**2 - return np.all(d2 >= min_allow2) - - def clear_of_pts(c, r, pts, clear): - if pts.size == 0: return True - d = np.linalg.norm(pts - c[None, :], axis=1) - return np.all(d >= (r + clear)) - - xmin, xmax = cfg.env_x_bounds - ymin, ymax = cfg.env_y_bounds - - Cs = np.zeros((0, 2), float) - Rs = np.zeros((0,), float) - Ws = np.zeros((0,), float) - - # Gates - n_gates = int(rng.integers(cfg.n_gates_range[0], cfg.n_gates_range[1] + 1)) - for _ in range(n_gates): - t_gate = float(rng.uniform(0.10, 0.90)) - g = centerline_point(t_gate) - rL = float(rng.uniform(0.40, 0.85)) - rR = float(rng.uniform(0.40, 0.85)) - cL = g - ortho * (half_gap + rL + eps) - cR = g + ortho * (half_gap + rR + eps) - jitter = rng.uniform(-0.25, 0.25) - cL = cL + jitter * dir01 - cR = cR + jitter * dir01 - if not non_overlap(cL, rL, Cs, Rs, cfg.min_separation_slack): continue - if not non_overlap(cR, rR, Cs, Rs, cfg.min_separation_slack): continue - if not clear_of_pts(cL, rL, np.stack([p0, p1]), min(cfg.min_start_clear, cfg.min_goal_clear)): continue - if not clear_of_pts(cR, rR, np.stack([p0, p1]), min(cfg.min_start_clear, cfg.min_goal_clear)): continue - Cs = np.vstack([Cs, cL, cR]) - Rs = np.concatenate([Rs, [rL, rR]]) - Ws = np.concatenate([Ws, [float(rng.uniform(*cfg.weight_range)), - float(rng.uniform(*cfg.weight_range))]]) - - # Intruders (push partially into corridor) - n_intr = int(rng.integers(cfg.n_intruders_range[0], cfg.n_intruders_range[1] + 1)) - for _ in range(n_intr): - t_int = float(rng.uniform(0.15, 0.85)) - g = centerline_point(t_int) - side = 1.0 if rng.uniform() < 0.5 else -1.0 - rI = float(rng.uniform(0.35, 0.75)) - depth = float(rng.uniform(0.60, 0.95)) * half_gap - depth = min(depth, half_gap - rI - 2*eps) - cI = g + side * ortho * (half_gap - depth + rI + eps) - cI = cI + rng.uniform(-0.18, 0.18) * dir01 - if not non_overlap(cI, rI, Cs, Rs, cfg.min_separation_slack): continue - if not clear_of_pts(cI, rI, np.stack([p0, p1]), min(cfg.min_start_clear, cfg.min_goal_clear)): continue - Cs = np.vstack([Cs, cI]); Rs = np.concatenate([Rs, [rI]]) - Ws = np.concatenate([Ws, [float(rng.uniform(*cfg.weight_range))]]) - - # Clutter strictly outside the corridor - n_target = int(rng.integers(max(50, cfg.n_obs_range[0]), cfg.n_obs_range[1] + 1)) - tries = 0 - while Cs.shape[0] < n_target and tries < cfg.poisson_max_tries: - tries += 1 - c = np.array([rng.uniform(xmin, xmax), rng.uniform(ymin, ymax)], float) - r = float(rng.uniform(cfg.radius_range[0], cfg.radius_range[1])) - if not disk_outside_corridor(c, r): continue - if not non_overlap(c, r, Cs, Rs, cfg.min_separation_slack): continue - if not clear_of_pts(c, r, np.stack([p0, p1]), min(cfg.min_start_clear, cfg.min_goal_clear)): continue - Cs = np.vstack([Cs, c]); Rs = np.concatenate([Rs, [r]]) - Ws = np.concatenate([Ws, [float(rng.uniform(*cfg.weight_range))]]) - - # Purge too close to start/goal - keep = np.ones(Cs.shape[0], dtype=bool) - for k in range(Cs.shape[0]): - if (np.linalg.norm(Cs[k] - p0) < (Rs[k] + cfg.min_start_clear)) or \ - (np.linalg.norm(Cs[k] - p1) < (Rs[k] + cfg.min_goal_clear)): - keep[k] = False - Cs, Rs, Ws = Cs[keep], Rs[keep], Ws[keep] - return Cs.astype(np.float64), Rs.astype(np.float64), Ws.astype(np.float64) - - -def planner_from_cfg(cfg, world_obs, lam_reg, mu, d_hat, radius): - """Small helper: make planner + scale the initial loop radius.""" - # Stage manager + deformable spline system - planner = PlannerMaximin( - start_xy=cfg.start, - goal_xy=cfg.goal, - stage_size=cfg.stage_size, - overlap=cfg.overlap, - n_ctrl=cfg.n_ctrl, - K=cfg.K, - d_hat=d_hat, - lam_reg=lam_reg, - mu=mu, - seed=cfg.seed_robot, - obstacles=(world_obs.C_np, world_obs.R_np), - inflate=cfg.inflate, - graph_cfg=cfg, - ) - # Scale initial loop to requested radius prior - with torch.no_grad(): - sys = planner.sys - P0 = sys.P0_loc - cur_norm = torch.mean(torch.linalg.norm(P0, dim=-1)).item() + 1e-8 - scale = float(radius / cur_norm) - sys.P0_loc *= scale - sys.Ploc *= scale - return planner - - -# ========= Planner/Stage classes (thin wrapper around your logic) ========= -from dataclasses import dataclass -from typing import Optional, Tuple, List, Any, Dict - -@dataclass -class Stage: - center: np.ndarray - size: Tuple[float, float] - entry_point: np.ndarray - exit_point: np.ndarray - stage_id: int - def __post_init__(self): - W, H = self.size; cx, cy = self.center - self.bounds = (cx - W/2, cx + W/2, cy - H/2, cy + H/2) - -def _union_bounds(b1, b2): - x1a,x1b,y1a,y1b = b1; x2a,x2b,y2a,y2b = b2 - return (min(x1a,x2a), max(x1b,x2b), min(y1a,y2a), max(y1b,y2b)) - -def _pad_bounds(b, px, py): - xmin,xmax,ymin,ymax = b - return (xmin - px, xmax + px, ymin - py, ymax + py) - -# --- (The graph helpers are lengthy; we’ll reuse your widest-path routine inline to keep behavior identical.) --- -# To keep this script focused, we import your previously defined helpers where possible. -# If those helpers live in your dataset script, you can paste them here, or import as a module. - -# Minimal imports of your maximin graph helpers (paste your originals if needed) -from math import inf -import heapq - -def _sample_rect_boundary(bounds, n_per_edge: int = 18) -> np.ndarray: - xmin, xmax, ymin, ymax = bounds - xs = np.linspace(xmin, xmax, n_per_edge, endpoint=True) - ys = np.linspace(ymin, ymax, n_per_edge, endpoint=True) - top = np.stack([xs, np.full_like(xs, ymax)], axis=1) - bottom = np.stack([xs, np.full_like(xs, ymin)], axis=1) - left = np.stack([np.full_like(ys, xmin), ys], axis=1) - right = np.stack([np.full_like(ys, xmax), ys], axis=1) - B = np.concatenate([top, bottom, left, right], axis=0) - return np.unique(B, axis=0) - -def _clearance_points(P: np.ndarray, C: np.ndarray, R: np.ndarray) -> np.ndarray: - if C.size == 0: return np.full((P.shape[0],), np.inf, dtype=float) - diff = P[:, None, :] - C[None, :, :] - dist = np.linalg.norm(diff, axis=-1) - d = dist - R[None, :] - return d.min(axis=1) - -def _segment_clearance(a: np.ndarray, b: np.ndarray, C: np.ndarray, R: np.ndarray, nsamp: int = 20) -> float: - ts = np.linspace(0.0, 1.0, nsamp) - P = (1.0 - ts)[:, None] * a[None, :] + ts[:, None] * b[None, :] - return float(_clearance_points(P, C, R).min()) - -def _los_segment(a: np.ndarray, b: np.ndarray, C: np.ndarray, R: np.ndarray) -> bool: - if C.size == 0: return True - ab = b - a - L2 = float(np.dot(ab, ab)) + 1e-12 - ac = C - a[None, :] - t = np.clip((ac @ ab) / L2, 0.0, 1.0) - p = a[None, :] + t[:, None] * ab[None, :] - d2 = np.sum((p - C)**2, axis=1) - return bool(np.all(d2 > (R * R))) - -def _sample_interior(bounds, grid_res: int = 12) -> np.ndarray: - xmin, xmax, ymin, ymax = bounds - xs = np.linspace(xmin, xmax, grid_res) - ys = np.linspace(ymin, ymax, grid_res) - XX, YY = np.meshgrid(xs, ys, indexing="xy") - return np.stack([XX.reshape(-1), YY.reshape(-1)], axis=1) - -def _knn_indices(P: np.ndarray, k: int = 6) -> List[List[int]]: - n = P.shape[0] - nbrs: List[List[int]] = [[] for _ in range(n)] - D2 = ((P[:, None, :] - P[None, :, :])**2).sum(axis=-1) - np.fill_diagonal(D2, np.inf) - kk = min(k, n-1) - idxs = np.argpartition(D2, kth=kk-1, axis=1)[:, :kk] - for i in range(n): - nbrs[i] = [int(j) for j in idxs[i]] - return nbrs - -def _maximin_widest_path(entry_idx: int, target_idx_set: set, - P: np.ndarray, neighbors: List[List[int]], - edge_clear_fn): - n = P.shape[0] - best = np.full(n, -np.inf, dtype=float) - parent = np.full(n, -1, dtype=int) - best[entry_idx] = np.inf - heap = [(-best[entry_idx], entry_idx)] - visited = np.zeros(n, dtype=bool) - while heap: - _, u = heapq.heappop(heap) - if visited[u]: continue - visited[u] = True - for v in neighbors[u]: - if visited[v]: continue - ce = edge_clear_fn(u, v) - cand = min(best[u], ce) - if cand > best[v]: - best[v] = cand - parent[v] = u - heapq.heappush(heap, (-best[v], v)) - if not target_idx_set: return -1, parent, {"best": best} - tgt_list = list(target_idx_set) - scores = best[tgt_list] - if np.all(scores == -np.inf): return -1, parent, {"best": best} - best_tgt = tgt_list[int(np.argmax(scores))] - return best_tgt, parent, {"best": best} - -def _reconstruct_path(parent: np.ndarray, tgt: int) -> np.ndarray: - if tgt < 0: return np.zeros((0,2), float) - path = [] - cur = tgt - while cur != -1: - path.append(cur) - cur = parent[cur] - path = path[::-1] - return np.array(path, int) - -def pick_maximin_boundary_exit_decoupled(entry, next_bounds, search_bounds, C_inf, R_inf, - grid_res=12, k=7, min_node_clear=1e-3, - nsamp_edge=20, min_exit_clearance=0.0, min_edge_clearance=0.0): - """Same decoupled-ROI widest-exit search as your dataset script.""" - B_next = _sample_rect_boundary(next_bounds, n_per_edge=max(14, grid_res)) - I_roi = _sample_interior(search_bounds, grid_res=grid_res) - P = np.vstack([entry[None, :], B_next, I_roi]) - - clr = _clearance_points(P, C_inf, R_inf) - keep = clr > max(min_node_clear, 1e-4) - keep[0] = True - P, clr = P[keep], clr[keep] - if P.shape[0] < 3: - return None, None, -np.inf - - entry_idx = int(np.argmin(np.linalg.norm(P - entry[None, :], axis=1))) - - xmin,xmax,ymin,ymax = next_bounds - epsb = 1e-6 - on_next = (np.isclose(P[:,0], xmin, atol=epsb) | np.isclose(P[:,0], xmax, atol=epsb) | - np.isclose(P[:,1], ymin, atol=epsb) | np.isclose(P[:,1], ymax, atol=epsb)) - ok_exit = on_next & (clr >= (min_exit_clearance + 1e-6)) - targets = set(np.nonzero(ok_exit)[0].tolist()) - targets.discard(entry_idx) - if not targets: - return None, None, -np.inf - - nbrs = _knn_indices(P, k=k) - for u in range(P.shape[0]): - nbrs[u] = [v for v in nbrs[u] if _los_segment(P[u], P[v], C_inf, R_inf)] - - req_edge = float(min_edge_clearance) - edge_cache: Dict[Tuple[int,int], float] = {} - def edge_clear(u, v): - key = (u, v) if u < v else (v, u) - if key in edge_cache: return edge_cache[key] - ce = _segment_clearance(P[u], P[v], C_inf, R_inf, nsamp=nsamp_edge) - edge_cache[key] = (ce if ce >= req_edge - 1e-6 else -np.inf) - return edge_cache[key] - - tgt_idx, parent, aux = _maximin_widest_path(entry_idx, targets, P, nbrs, edge_clear) - if tgt_idx < 0 or not np.isfinite(aux["best"][tgt_idx]) or aux["best"][tgt_idx] <= 0.0: - return None, None, -np.inf - - path_idx = _reconstruct_path(parent, tgt_idx) - return P[tgt_idx].copy(), P[path_idx], float(aux["best"][tgt_idx]) - - -class StageManagerMaximin: - def __init__(self, start_xy, goal_xy, stage_size=(2.6,2.0), overlap_ratio=0.3, - obstacles=None, inflate=0.35, graph_cfg: Optional[GenCfg]=None): - self.start = np.array(start_xy, float) - self.goal = np.array(goal_xy, float) - self.stage_size = tuple(stage_size) - self.overlap_ratio = float(overlap_ratio) - self.C = np.zeros((0, 2), float); self.R = np.zeros((0,), float) - if obstacles is not None: - C, R = obstacles - self.C = np.array(C, float) - self.R = np.array(R, float) + float(inflate) - self.gcfg = graph_cfg or GenCfg() - self.inflate = float(inflate) - self.stages: List[Stage] = [] - self.current_stage_idx = 0 - self._build() - - def _build(self): - d = self.goal - self.start - L = np.linalg.norm(d) - if L < 1e-6: - self.stages = [Stage((self.start + self.goal) / 2, self.stage_size, - self.start.copy(), self.goal.copy(), 0)] - return - step = self.stage_size[0] * (1.0 - self.overlap_ratio) - n = max(1, int(np.ceil(L / step))) - centers = [] - for i in range(n): - t = 0.0 if n == 1 else i / (n - 1) - centers.append(self.start + t * d) - centers = [np.array(c, float) for c in centers] - - W, H = self.stage_size - pad_x = self.gcfg.roi_pad_scale_x * W - pad_y = self.gcfg.roi_pad_scale_y * H - - for i in range(n): - c_i = centers[i] - entry = self.start if i == 0 else self.stages[-1].exit_point - if i == n - 1: - exitp = self.goal - else: - nxt_center = centers[i + 1] - next_bounds = (nxt_center[0] - W/2, nxt_center[0] + W/2, - nxt_center[1] - H/2, nxt_center[1] + H/2) - cur_bounds = (c_i[0] - W/2, c_i[0] + W/2, - c_i[1] - H/2, c_i[1] + H/2) - roi = _union_bounds(cur_bounds, next_bounds) if self.gcfg.search_union_current_next else next_bounds - search_bounds = _pad_bounds(roi, pad_x, pad_y) - - cand, poly, _ = pick_maximin_boundary_exit_decoupled( - entry=np.array(entry, float), - next_bounds=next_bounds, - search_bounds=search_bounds, - C_inf=self.C, R_inf=self.R, - grid_res=self.gcfg.grid_res, - k=self.gcfg.knn_k, - min_node_clear=self.gcfg.min_node_clear, - nsamp_edge=self.gcfg.nsamp_edge, - min_exit_clearance=self.gcfg.min_exit_margin, - min_edge_clearance=self.gcfg.min_edge_margin, - ) - if cand is not None: - exitp = cand - if self.gcfg.follow_corridor and poly is not None and poly.shape[0] >= 2: - seg_len = np.linalg.norm(np.diff(poly, axis=0), axis=1) - if seg_len.size > 0 and np.sum(seg_len) > 1e-9: - s = np.cumsum(seg_len) - j = int(np.searchsorted(s, s[-1]*0.33)) - j = min(max(j, 0), poly.shape[0]-1) - desired_center = poly[j] - centers[i + 1] = 0.7 * centers[i + 1] + 0.3 * desired_center - else: - exitp = self.start + (i + 1) / n * d - - self.stages.append(Stage(c_i, self.stage_size, np.array(entry, float), np.array(exitp, float), i)) - - def current(self) -> Stage: - return self.stages[self.current_stage_idx] - - def advance_if_needed(self, robot_xy: np.ndarray, thresh=0.45) -> bool: - st = self.current() - if (np.linalg.norm(robot_xy - st.exit_point) < thresh) and (self.current_stage_idx < len(self.stages) - 1): - self.current_stage_idx += 1 - return True - if self.current_stage_idx < len(self.stages) - 1: - nxt = self.stages[self.current_stage_idx + 1] - path_dir = self.goal - self.start - if np.dot(robot_xy - st.center, path_dir) > np.dot(nxt.center - st.center, path_dir): - self.current_stage_idx += 1 - return True - return False - - -class PlannerMaximin: - def __init__(self, start_xy, goal_xy, stage_size=(2.5,2.0), overlap=0.3, - n_ctrl=20, K=240, d_hat=1.0, lam_reg=3.0, mu=0.25, seed=7, - obstacles=None, inflate=0.35, graph_cfg: Optional[GenCfg]=None): - - self.sm = StageManagerMaximin( - np.array(start_xy), np.array(goal_xy), - stage_size=stage_size, overlap_ratio=overlap, - obstacles=obstacles, inflate=inflate, graph_cfg=graph_cfg - ) - - self.sys = ssi.DeformableSplineSystem(n_ctrl=n_ctrl, K=K, d_hat=d_hat, lam_reg=lam_reg, mu=mu, seed=seed) - # initialize pose - self.sys.o = torch.tensor(start_xy, dtype=self.sys.dtype) - d = np.array(goal_xy)-np.array(start_xy) - self.sys.theta = torch.tensor(math.atan2(d[1],d[0]), dtype=self.sys.dtype) - - self.stage_field = ssi.StageForceFieldTorch(r_safe=0.30, r_contact=0.16, - w_goal=1.0, w_radial=2.0, w_tangential=1.2, w_flow=0.8, w_boundary=0.6) - - def _stage_slice(self, C,R,W): - st = self.sm.current(); xmin,xmax,ymin,ymax = st.bounds - m = (C[:,0] >= xmin-0.5) & (C[:,0] <= xmax+0.5) & (C[:,1] >= ymin-0.5) & (C[:,1] <= ymax+0.5) - if np.any(m): return C[m],R[m],W[m] - return C,R,W - - def step(self, dt, world_obs: ssi.WorldObstacles): - st = self.sm.current() - C,R,W = self._stage_slice(world_obs.C_np, world_obs.R_np, world_obs.W_np) - obs_t = ssi.ObstacleProviderTorch(C,R,W, dtype=self.sys.dtype) - barrier = ssi.IPCBarrier(obs_t, d_hat=world_obs.d_hat) - info = self.sys.step(dt, obs_t, barrier, self.stage_field, st.bounds, tuple(st.exit_point)) - self.sm.advance_if_needed(np.array(info["center"],float)) - return info - - -# ========= Movie runner ========= -def run_large_example_and_movie(): - cfg = BigCfg() - set_all_seeds(cfg.seed) - os.makedirs(cfg.out_root, exist_ok=True) - - # Build a single, large environment - C, R, W = sample_obstacles_case1_tight(cfg) - world = ssi.WorldObstacles(C, R, W, d_hat=cfg.d_hat) - - # Planner with same core robot/configs (just scaled world) - planner = planner_from_cfg(cfg, world, - lam_reg=cfg.lam_reg, mu=cfg.mu, d_hat=cfg.d_hat, radius=cfg.radius) - - # Figure & axes - fig = plt.figure(figsize=(14, 6), dpi=100) - ax = fig.add_subplot(111) - ax.set_aspect("equal", "box") - - # MP4 path - mp4_path = os.path.join(cfg.out_root, "large_example.mp4") - - # FPS ~ 1/(dt*snapshot_every); clamp to >= 8 for a smoother look - fps = 24 # max(8, int(round(1.0 / (cfg.dt * cfg.snapshot_every)))) - - # Fallback: if ffmpeg not available, we’ll dump PNG frames instead - ffmpeg_ok = writers.is_available("ffmpeg") - png_dir = os.path.join(cfg.out_root, "frames_png") - - def draw_frame(t_step: int): - ax.clear() - ax.set_aspect("equal", "box") - - # Obstacles (raw + inflated) - ssi.draw_obstacles(ax, world.C_np, world.R_np, ec="k") - if cfg.overlay_inflated_in_snapshot: - ssi.draw_obstacles(ax, world.C_np, world.R_np + cfg.inflate, ec="#00cfdc", alpha=0.25) - - # Stage windows and entry/exit markers - for i, st in enumerate(planner.sm.stages): - xmin, xmax, ymin, ymax = st.bounds - ax.add_patch(plt.Rectangle( - (xmin, ymin), xmax - xmin, ymax - ymin, - fill=False, lw=1.4, - ec="#00cfdc" if i == planner.sm.current_stage_idx else (0.4,0.6,0.6,0.4) - )) - ax.plot([st.entry_point[0]],[st.entry_point[1]], "o", ms=3, color="#85ff85") - ax.plot([st.exit_point[0]],[st.exit_point[1]], "o", ms=3, color="#ff8585") - - # Current spline shape - sys = planner.sys - with torch.no_grad(): - Pw = sys.world_points().detach().cpu().numpy() - # ssi.render_curve(ax, planner.sys.B, Pw, label=None, color="#1f77b4") - Pw = planner.sys.world_points().detach().cpu().numpy() - B = planner.sys.B.detach().cpu().numpy() - X = (B @ Pw) # (K,2) curve samples - - (line,) = ax.plot(X[:,0], X[:,1], - lw=5.0, # <-- make the spline look bigger - alpha=0.98, - solid_capstyle="round", - solid_joinstyle="round", - zorder=5) - - # add a light stroke for contrast (optional but helps visibility) - import matplotlib.patheffects as pe - line.set_path_effects([ - pe.Stroke(linewidth=7.0, foreground="white"), - pe.Normal() - ]) - - # Robot COM, start/goal - c = planner.sys.o.detach().cpu().numpy() - ax.plot([c[0]],[c[1]], "o", ms=4, color="black") - ax.plot([cfg.start[0]],[cfg.start[1]], "o", color="green", ms=6) - ax.plot([cfg.goal[0]],[cfg.goal[1]], "*", color="gold", ms=10) - - # Nice bounds - ax.set_xlim(cfg.env_x_bounds[0]-1.0, cfg.env_x_bounds[1]+1.0) - ax.set_ylim(cfg.env_y_bounds[0]-1.0, cfg.env_y_bounds[1]+1.0) - - ax.set_title(f"Large-scale navigation | t={t_step*cfg.dt:.2f}s " - f"stage {planner.sm.current_stage_idx+1}/{len(planner.sm.stages)}") - ax.set_xlabel("x") - ax.set_ylabel("y") - - if ffmpeg_ok: - writer = FFMpegWriter( - fps=fps, - codec="mpeg4", - bitrate=-1, - extra_args=["-vtag", "mp4v", "-pix_fmt", "yuv420p"] # ensure MP4 compatibility - ) - with writer.saving(fig, mp4_path, dpi=120): - for t in range(cfg.total_steps): - info = planner.step(cfg.dt, world) - # Only draw at snapshot cadence - if (t % cfg.snapshot_every) != 0 and t not in (0, cfg.total_steps - 1): - continue - draw_frame(t) - writer.grab_frame() - plt.close(fig) - print(f"[OK] MP4 saved to: {mp4_path}") - else: - # PNG fallback - os.makedirs(png_dir, exist_ok=True) - frame_id = 0 - for t in range(cfg.total_steps): - info = planner.step(cfg.dt, world) - if (t % cfg.snapshot_every) != 0 and t not in (0, cfg.total_steps - 1): - continue - draw_frame(t) - png_path = os.path.join(png_dir, f"frame_{frame_id:05d}.png") - fig.savefig(png_path, bbox_inches="tight") - frame_id += 1 - plt.close(fig) - print(f"[WARN] ffmpeg not available. Wrote PNG frames to: {png_dir}") - print(" To make an mp4 manually you can run something like:") - print(f" ffmpeg -y -framerate {fps} -i {png_dir}/frame_%05d.png -c:v libx264 -pix_fmt yuv420p {mp4_path}") - - # Also dump a small manifest for reproducibility - manifest = { - "seed": cfg.seed, - "dt": cfg.dt, - "snapshot_every": cfg.snapshot_every, - "fps": fps, - "start": cfg.start.tolist(), - "goal": cfg.goal.tolist(), - "stage_size": list(cfg.stage_size), - "env_bounds": [*cfg.env_x_bounds, *cfg.env_y_bounds], - "inflate": cfg.inflate, - "robot": {"n_ctrl": cfg.n_ctrl, "K": cfg.K, "lam_reg": cfg.lam_reg, "mu": cfg.mu, "d_hat": cfg.d_hat, "radius": cfg.radius}, - "graph": { - "grid_res": cfg.grid_res, "knn_k": cfg.knn_k, "nsamp_edge": cfg.nsamp_edge, - "min_node_clear": cfg.min_node_clear, "follow_corridor": bool(cfg.follow_corridor), - "roi_pad_scale_x": cfg.roi_pad_scale_x, "roi_pad_scale_y": cfg.roi_pad_scale_y, - "search_union_current_next": bool(cfg.search_union_current_next), - "min_edge_margin": cfg.min_edge_margin, "min_exit_margin": cfg.min_exit_margin, - }, - "outputs": {"mp4": mp4_path if ffmpeg_ok else None, "png_dir": png_dir if not ffmpeg_ok else None} - } - with open(os.path.join(cfg.out_root, "large_example_manifest.json"), "w") as f: - json.dump(manifest, f, indent=2) - - -if __name__ == "__main__": - run_large_example_and_movie() diff --git a/scripts/spline_dataset_maximin.py b/scripts/spline_dataset_maximin.py deleted file mode 100644 index 7e64d97..0000000 --- a/scripts/spline_dataset_maximin.py +++ /dev/null @@ -1,913 +0,0 @@ - -#!/usr/bin/env python3 -# Generate a dataset using a max–min (widest-path) exit selection per stage -# with decoupled search region, hard tube clearance, and solvable tight corridors. - -import sys -import os, json, math, random, heapq -from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple - -import numpy as np -import torch -from torch.utils.data import Dataset - -import matplotlib -matplotlib.use("Agg") -import matplotlib.pyplot as plt - -import scripts.spline_stagewise4 as ssi - - -# ============================== -# Config -# ============================== - -class GenCfg: - out_root: str = "./nav_dataset_maximin" - seed: int = 2025 - - # simulation - total_steps: int = 2000 - dt: float = 0.03 - snapshot_every: int = 5 - - # start/goal and stage geometry (stage centers may be corridor-followed) - start = np.array([-1.0, -0.5], float) - goal = np.array([ 9.0, 0.2], float) - stage_size = (2.6, 2.0) - overlap = 0.30 - - # ---- tube radius (robot radius + margin) ---- - # We inflate obstacles by this much; guarantees a free tube of radius = inflate. - inflate = 0.05 - - # robot base params (perturbed for Case 2) - n_ctrl: int = 20 - K: int = 240 - lam_reg: float = 3.0 - mu: float = 0.25 - d_hat: float = 1.0 - seed_robot: int = 7 - radius: float = 0.35 # initial loop radius (shape prior for the spline) - - # success thresholds - goal_tol: float = 0.30 - min_clear_tol: float = 1e-6 - - # dataset sizes - num_env_variations: int = 10 - num_robot_variations: int = 10 - - # *** Case 1 (tight corridor) obstacle generation *** - env_x_bounds = (-0.5, 9.0) - env_y_bounds = (-1.8, 1.8) - radius_range = (0.28, 0.80) - weight_range = (0.6, 1.5) - n_obs_range = (12, 22) # total count target (gates + intruders + clutter) - poisson_max_tries: int = 6000 - min_separation_slack: float = 0.08 - min_start_clear: float = 0.9 - min_goal_clear: float = 0.9 - - # tight gap controls (Case 1) — we clamp to >= 2*inflate + eps to ensure solvable - gap_mult_range = (0.70, 0.80) # gap ≈ mult * (2*radius) + margin - gap_margin_range = (0.01, 0.03) - corridor_wiggle_amp = 0.35 - corridor_wiggle_k = 0.6 - n_gates_range = (3, 5) - n_intruders_range = (2, 4) - - # *** Case 2 (robot perturbations) *** - lam_reg_mult_range = (0.92, 1.08) - mu_mult_range = (0.92, 1.08) - d_hat_mult_range = (0.95, 1.05) - radius_abs_range = (0.30, 0.50) - - # widest-path graph parameters - grid_res: int = 12 # interior grid samples per axis - knn_k: int = 7 # neighbors per node - nsamp_edge: int = 20 # samples to estimate edge clearance - min_node_clear: float = 1e-3 - follow_corridor: bool = True # recenter next stage toward corridor - - # decoupled-search ROI (bigger than next stage) - roi_pad_scale_x: float = 0.8 # pad ~0.8*W on both sides - roi_pad_scale_y: float = 0.8 # pad ~0.8*H on both sides - search_union_current_next: bool = True # use union(current, next) before padding - - # hard extra margins beyond tube (set small >0 to avoid grazing) - min_edge_margin: float = 1e-3 # extra beyond inflate along edges - min_exit_margin: float = 1e-3 # extra beyond inflate at exit - - # viz - overlay_inflated_in_snapshot: bool = True - - -# ============================== -# Utility / RNG -# ============================== - -def set_all_seeds(seed: int): - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - -def rand_uniform(a: float, b: float) -> float: - return float(np.random.uniform(a, b)) - - -# ============================== -# Tight obstacle generator (Case 1) -# ============================== - -def sample_obstacles_case1_tight(cfg: GenCfg) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: - """Tight, solvable corridor with enforced narrow gates + intruders.""" - rng = np.random.default_rng() - - # target gap tied to robot size, but clamp so it's always >= tube diameter - robot_diam = 2.0 * cfg.radius - gap_mult = rng.uniform(*cfg.gap_mult_range) - gap_margin = rng.uniform(*cfg.gap_margin_range) - target_gap = gap_mult * robot_diam + gap_margin - - # enforce solvability w.r.t. the tube used in planning - min_gap = 2.0 * cfg.inflate + 0.01 # tiny slack - if target_gap < min_gap: - target_gap = min_gap - half_gap = 0.5 * target_gap - - eps = 0.012 - corridor_half_w = half_gap + 0.01 - - p0 = cfg.start.astype(float); p1 = cfg.goal.astype(float) - dir01 = p1 - p0; dir01 = dir01 / (np.linalg.norm(dir01) + 1e-12) - ortho = np.array([-dir01[1], dir01[0]]) - amp = cfg.corridor_wiggle_amp - kfreq = cfg.corridor_wiggle_k - ts = np.linspace(0.0, 1.0, 80) - - def centerline_point(t): - base = (1 - t) * p0 + t * p1 - return base + amp * np.sin(2 * np.pi * kfreq * t) * ortho - - centerline = np.stack([centerline_point(t) for t in ts], axis=0) - - def dist_to_polyline(pt): - return float(np.linalg.norm(centerline - pt[None, :], axis=1).min()) - - def disk_outside_corridor(c, r): - return dist_to_polyline(c) >= (corridor_half_w + r) - - def non_overlap(c, r, Cs, Rs, slack): - if Cs.size == 0: return True - d2 = np.sum((Cs - c[None, :])**2, axis=1) - min_allow2 = (Rs + r + slack)**2 - return np.all(d2 >= min_allow2) - - def clear_of_pts(c, r, pts, clear): - if pts.size == 0: return True - d = np.linalg.norm(pts - c[None, :], axis=1) - return np.all(d >= (r + clear)) - - xmin, xmax = cfg.env_x_bounds - ymin, ymax = cfg.env_y_bounds - - Cs = np.zeros((0, 2), float) - Rs = np.zeros((0,), float) - Ws = np.zeros((0,), float) - - # gates - n_gates = int(rng.integers(cfg.n_gates_range[0], cfg.n_gates_range[1] + 1)) - for _ in range(n_gates): - t_gate = float(rng.uniform(0.18, 0.82)) - g = centerline_point(t_gate) - rL = float(rng.uniform(0.32, 0.55)) - rR = float(rng.uniform(0.32, 0.55)) - cL = g - ortho * (half_gap + rL + eps) - cR = g + ortho * (half_gap + rR + eps) - jitter = rng.uniform(-0.12, 0.12) - cL = cL + jitter * dir01 - cR = cR + jitter * dir01 - if not non_overlap(cL, rL, Cs, Rs, cfg.min_separation_slack): continue - if not non_overlap(cR, rR, Cs, Rs, cfg.min_separation_slack): continue - if not clear_of_pts(cL, rL, np.stack([p0, p1]), min(cfg.min_start_clear, cfg.min_goal_clear)): continue - if not clear_of_pts(cR, rR, np.stack([p0, p1]), min(cfg.min_start_clear, cfg.min_goal_clear)): continue - Cs = np.vstack([Cs, cL, cR]) - Rs = np.concatenate([Rs, [rL, rR]]) - Ws = np.concatenate([Ws, [float(rng.uniform(*cfg.weight_range)), - float(rng.uniform(*cfg.weight_range))]]) - - # intruders - n_intr = int(rng.integers(cfg.n_intruders_range[0], cfg.n_intruders_range[1] + 1)) - for _ in range(n_intr): - t_int = float(rng.uniform(0.20, 0.80)) - g = centerline_point(t_int) - side = 1.0 if rng.uniform() < 0.5 else -1.0 - rI = float(rng.uniform(0.28, 0.48)) - depth = float(rng.uniform(0.65, 0.90)) * half_gap - depth = min(depth, half_gap - rI - 2*eps) - cI = g + side * ortho * (half_gap - depth + rI + eps) - cI = cI + rng.uniform(-0.10, 0.10) * dir01 - if not non_overlap(cI, rI, Cs, Rs, cfg.min_separation_slack): continue - if not clear_of_pts(cI, rI, np.stack([p0, p1]), min(cfg.min_start_clear, cfg.min_goal_clear)): continue - Cs = np.vstack([Cs, cI]); Rs = np.concatenate([Rs, [rI]]) - Ws = np.concatenate([Ws, [float(rng.uniform(*cfg.weight_range))]]) - - # clutter strictly outside corridor - n_target = int(rng.integers(max(10, cfg.n_obs_range[0]), cfg.n_obs_range[1] + 1)) - tries = 0 - while Cs.shape[0] < n_target and tries < cfg.poisson_max_tries: - tries += 1 - c = np.array([rng.uniform(xmin, xmax), rng.uniform(ymin, ymax)], float) - r = float(rng.uniform(cfg.radius_range[0], cfg.radius_range[1])) - if not disk_outside_corridor(c, r): continue - if not non_overlap(c, r, Cs, Rs, cfg.min_separation_slack): continue - if not clear_of_pts(c, r, np.stack([p0, p1]), min(cfg.min_start_clear, cfg.min_goal_clear)): continue - Cs = np.vstack([Cs, c]); Rs = np.concatenate([Rs, [r]]) - Ws = np.concatenate([Ws, [float(rng.uniform(*cfg.weight_range))]]) - - # purge near start/goal - keep = np.ones(Cs.shape[0], dtype=bool) - for k in range(Cs.shape[0]): - if (np.linalg.norm(Cs[k] - p0) < (Rs[k] + cfg.min_start_clear)) or \ - (np.linalg.norm(Cs[k] - p1) < (Rs[k] + cfg.min_goal_clear)): - keep[k] = False - Cs, Rs, Ws = Cs[keep], Rs[keep], Ws[keep] - return Cs.astype(np.float64), Rs.astype(np.float64), Ws.astype(np.float64) - - -# ============================== -# Widest-path graph helpers -# ============================== - -def _sample_rect_boundary(bounds, n_per_edge: int = 18) -> np.ndarray: - xmin, xmax, ymin, ymax = bounds - xs = np.linspace(xmin, xmax, n_per_edge, endpoint=True) - ys = np.linspace(ymin, ymax, n_per_edge, endpoint=True) - top = np.stack([xs, np.full_like(xs, ymax)], axis=1) - bottom = np.stack([xs, np.full_like(xs, ymin)], axis=1) - left = np.stack([np.full_like(ys, xmin), ys], axis=1) - right = np.stack([np.full_like(ys, xmax), ys], axis=1) - B = np.concatenate([top, bottom, left, right], axis=0) - B = np.unique(B, axis=0) - return B - -def _clearance_points(P: np.ndarray, C: np.ndarray, R: np.ndarray) -> np.ndarray: - if C.size == 0: - return np.full((P.shape[0],), np.inf, dtype=float) - diff = P[:, None, :] - C[None, :, :] - dist = np.linalg.norm(diff, axis=-1) - d = dist - R[None, :] - return d.min(axis=1) - -def _segment_clearance(a: np.ndarray, b: np.ndarray, C: np.ndarray, R: np.ndarray, nsamp: int = 20) -> float: - ts = np.linspace(0.0, 1.0, nsamp) - P = (1.0 - ts)[:, None] * a[None, :] + ts[:, None] * b[None, :] - return float(_clearance_points(P, C, R).min()) - -def _los_segment(a: np.ndarray, b: np.ndarray, C: np.ndarray, R: np.ndarray) -> bool: - if C.size == 0: return True - ab = b - a - L2 = float(np.dot(ab, ab)) + 1e-12 - ac = C - a[None, :] - t = np.clip((ac @ ab) / L2, 0.0, 1.0) - p = a[None, :] + t[:, None] * ab[None, :] - d2 = np.sum((p - C)**2, axis=1) - return bool(np.all(d2 > (R * R))) - -def _sample_interior(bounds, grid_res: int = 12) -> np.ndarray: - xmin, xmax, ymin, ymax = bounds - xs = np.linspace(xmin, xmax, grid_res) - ys = np.linspace(ymin, ymax, grid_res) - XX, YY = np.meshgrid(xs, ys, indexing="xy") - return np.stack([XX.reshape(-1), YY.reshape(-1)], axis=1) - -def _knn_indices(P: np.ndarray, k: int = 6) -> List[List[int]]: - n = P.shape[0] - nbrs: List[List[int]] = [[] for _ in range(n)] - D2 = ((P[:, None, :] - P[None, :, :])**2).sum(axis=-1) - np.fill_diagonal(D2, np.inf) - kk = min(k, n-1) - idxs = np.argpartition(D2, kth=kk-1, axis=1)[:, :kk] - for i in range(n): - nbrs[i] = [int(j) for j in idxs[i]] - return nbrs - -def _maximin_widest_path(entry_idx: int, target_idx_set: set, - P: np.ndarray, neighbors: List[List[int]], - edge_clear_fn) -> Tuple[int, np.ndarray, Dict[str, np.ndarray]]: - n = P.shape[0] - best = np.full(n, -np.inf, dtype=float) - parent = np.full(n, -1, dtype=int) - best[entry_idx] = np.inf - heap = [(-best[entry_idx], entry_idx)] - visited = np.zeros(n, dtype=bool) - while heap: - negb, u = heapq.heappop(heap) - if visited[u]: continue - visited[u] = True - for v in neighbors[u]: - if visited[v]: continue - ce = edge_clear_fn(u, v) - cand = min(best[u], ce) - if cand > best[v]: - best[v] = cand - parent[v] = u - heapq.heappush(heap, (-best[v], v)) - if not target_idx_set: - return -1, parent, {"best": best} - tgt_list = list(target_idx_set) - scores = best[tgt_list] - if np.all(scores == -np.inf): - return -1, parent, {"best": best} - best_tgt = tgt_list[int(np.argmax(scores))] - return best_tgt, parent, {"best": best} - -def _reconstruct_path(parent: np.ndarray, tgt: int) -> np.ndarray: - if tgt < 0: return np.zeros((0,2), float) - path = [] - cur = tgt - while cur != -1: - path.append(cur) - cur = parent[cur] - path = path[::-1] - return np.array(path, int) - -def _union_bounds(b1, b2): - x1a,x1b,y1a,y1b = b1; x2a,x2b,y2a,y2b = b2 - return (min(x1a,x2a), max(x1b,x2b), min(y1a,y2a), max(y1b,y2b)) - -def _pad_bounds(b, px, py): - xmin,xmax,ymin,ymax = b - return (xmin - px, xmax + px, ymin - py, ymax + py) - -def pick_maximin_boundary_exit_decoupled( - entry: np.ndarray, - next_bounds: Tuple[float,float,float,float], # boundary where exit must lie - search_bounds: Tuple[float,float,float,float], # larger ROI for graph - C_inf: np.ndarray, - R_inf: np.ndarray, - grid_res: int = 12, - k: int = 7, - min_node_clear: float = 1e-3, - nsamp_edge: int = 20, - min_exit_clearance: float = 0.0, # extra beyond tube - min_edge_clearance: float = 0.0, # extra beyond tube -) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], float]: - """Search in a larger ROI but choose exit on next_bounds; enforce hard clearances.""" - B_next = _sample_rect_boundary(next_bounds, n_per_edge=max(14, grid_res)) - I_roi = _sample_interior(search_bounds, grid_res=grid_res) - P = np.vstack([entry[None, :], B_next, I_roi]) - - clr = _clearance_points(P, C_inf, R_inf) - keep = clr > max(min_node_clear, 1e-4) - keep[0] = True - P, clr = P[keep], clr[keep] - if P.shape[0] < 3: - return None, None, -np.inf - - entry_idx = int(np.argmin(np.linalg.norm(P - entry[None, :], axis=1))) - - xmin,xmax,ymin,ymax = next_bounds - epsb = 1e-6 - on_next = (np.isclose(P[:,0], xmin, atol=epsb) | np.isclose(P[:,0], xmax, atol=epsb) | - np.isclose(P[:,1], ymin, atol=epsb) | np.isclose(P[:,1], ymax, atol=epsb)) - ok_exit = on_next & (clr >= (min_exit_clearance + 1e-6)) - targets = set(np.nonzero(ok_exit)[0].tolist()) - targets.discard(entry_idx) - if not targets: - return None, None, -np.inf - - nbrs = _knn_indices(P, k=k) - for u in range(P.shape[0]): - nbrs[u] = [v for v in nbrs[u] if _los_segment(P[u], P[v], C_inf, R_inf)] - - req_edge = float(min_edge_clearance) - edge_cache: Dict[Tuple[int,int], float] = {} - def edge_clear(u, v): - key = (u, v) if u < v else (v, u) - if key in edge_cache: return edge_cache[key] - ce = _segment_clearance(P[u], P[v], C_inf, R_inf, nsamp=nsamp_edge) - edge_cache[key] = (ce if ce >= req_edge - 1e-6 else -np.inf) - return edge_cache[key] - - tgt_idx, parent, aux = _maximin_widest_path(entry_idx, targets, P, nbrs, edge_clear) - if tgt_idx < 0 or not np.isfinite(aux["best"][tgt_idx]) or aux["best"][tgt_idx] <= 0.0: - return None, None, -np.inf - - path_idx = _reconstruct_path(parent, tgt_idx) - return P[tgt_idx].copy(), P[path_idx], float(aux["best"][tgt_idx]) - - -# ============================== -# Stage manager with widest-path exits (decoupled search) -# ============================== - -@dataclass -class Stage: - center: np.ndarray - size: Tuple[float, float] - entry_point: np.ndarray - exit_point: np.ndarray - stage_id: int - def __post_init__(self): - W, H = self.size; cx, cy = self.center - self.bounds = (cx - W/2, cx + W/2, cy - H/2, cy + H/2) - -class StageManagerMaximin: - """Build stages using max–min exits; search over padded union ROI; optionally corridor-follow centers.""" - def __init__( - self, - start_xy: np.ndarray, - goal_xy: np.ndarray, - stage_size=(2.6, 2.0), - overlap_ratio=0.3, - obstacles: Optional[Tuple[np.ndarray, np.ndarray]] = None, # (C, R) - inflate: float = 0.35, - graph_cfg: Optional[GenCfg] = None, - ): - self.start = np.array(start_xy, float) - self.goal = np.array(goal_xy, float) - self.stage_size = tuple(stage_size) - self.overlap_ratio = float(overlap_ratio) - - # store inflated obstacles for planning/search - self.C = np.zeros((0, 2), float); self.R = np.zeros((0,), float) - if obstacles is not None: - C, R = obstacles - self.C = np.array(C, float) - self.R = np.array(R, float) + float(inflate) - - self.gcfg = graph_cfg or GenCfg() - self.inflate = float(inflate) - - self.stages: List[Stage] = [] - self.current_stage_idx = 0 - self._build() - - def _build(self): - d = self.goal - self.start - L = np.linalg.norm(d) - if L < 1e-6: - self.stages = [Stage((self.start + self.goal) / 2, self.stage_size, - self.start.copy(), self.goal.copy(), 0)] - return - - step = self.stage_size[0] * (1.0 - self.overlap_ratio) - n = max(1, int(np.ceil(L / step))) - - # initial straight-line centers - centers = [] - for i in range(n): - t = 0.0 if n == 1 else i / (n - 1) - centers.append(self.start + t * d) - centers = [np.array(c, float) for c in centers] - - W, H = self.stage_size - pad_x = self.gcfg.roi_pad_scale_x * W - pad_y = self.gcfg.roi_pad_scale_y * H - - for i in range(n): - c_i = centers[i] - entry = self.start if i == 0 else self.stages[-1].exit_point - - if i == n - 1: - exitp = self.goal - else: - nxt_center = centers[i + 1] - next_bounds = (nxt_center[0] - W/2, nxt_center[0] + W/2, - nxt_center[1] - H/2, nxt_center[1] + H/2) - cur_bounds = (c_i[0] - W/2, c_i[0] + W/2, - c_i[1] - H/2, c_i[1] + H/2) - - # bigger search ROI: union(current, next) padded - if self.gcfg.search_union_current_next: - roi = _union_bounds(cur_bounds, next_bounds) - else: - roi = next_bounds - search_bounds = _pad_bounds(roi, pad_x, pad_y) - - cand, poly, bottle = pick_maximin_boundary_exit_decoupled( - entry=np.array(entry, float), - next_bounds=next_bounds, - search_bounds=search_bounds, - C_inf=self.C, R_inf=self.R, - grid_res=self.gcfg.grid_res, - k=self.gcfg.knn_k, - min_node_clear=self.gcfg.min_node_clear, - nsamp_edge=self.gcfg.nsamp_edge, - # require exactly the tube plus tiny extra margins - min_exit_clearance=self.gcfg.min_exit_margin, - min_edge_clearance=self.gcfg.min_edge_margin, - ) - - if cand is not None: - exitp = cand - # corridor-follow next center toward the discovered path - if self.gcfg.follow_corridor and poly is not None and poly.shape[0] >= 2: - seg_len = np.linalg.norm(np.diff(poly, axis=0), axis=1) - if seg_len.size > 0 and np.sum(seg_len) > 1e-9: - s = np.cumsum(seg_len) - j = int(np.searchsorted(s, s[-1]*0.33)) # ~first third - j = min(max(j, 0), poly.shape[0]-1) - desired_center = poly[j] - centers[i + 1] = 0.7 * centers[i + 1] + 0.3 * desired_center - else: - # fallback to straight-line fraction - exitp = self.start + (i + 1) / n * d - - self.stages.append(Stage(c_i, self.stage_size, np.array(entry, float), np.array(exitp, float), i)) - - def current(self) -> Stage: - return self.stages[self.current_stage_idx] - - def advance_if_needed(self, robot_xy: np.ndarray, thresh=0.45) -> bool: - st = self.current() - if (np.linalg.norm(robot_xy - st.exit_point) < thresh) and (self.current_stage_idx < len(self.stages) - 1): - self.current_stage_idx += 1 - return True - if self.current_stage_idx < len(self.stages) - 1: - nxt = self.stages[self.current_stage_idx + 1] - path_dir = self.goal - self.start - if np.dot(robot_xy - st.center, path_dir) > np.dot(nxt.center - st.center, path_dir): - self.current_stage_idx += 1 - return True - return False - - -# ============================== -# Planner wrapper using StageManagerMaximin -# ============================== - -class PlannerMaximin: - def __init__(self, start_xy, goal_xy, stage_size=(2.5,2.0), overlap=0.3, - n_ctrl=20, K=240, d_hat=1.0, lam_reg=3.0, mu=0.25, seed=7, - obstacles=None, inflate=0.35, graph_cfg: Optional[GenCfg]=None): - - self.sm = StageManagerMaximin( - np.array(start_xy), np.array(goal_xy), - stage_size=stage_size, overlap_ratio=overlap, - obstacles=obstacles, inflate=inflate, graph_cfg=graph_cfg - ) - - self.sys = ssi.DeformableSplineSystem(n_ctrl=n_ctrl, K=K, d_hat=d_hat, lam_reg=lam_reg, mu=mu, seed=seed) - # initialize pose - self.sys.o = torch.tensor(start_xy, dtype=self.sys.dtype) - d = np.array(goal_xy)-np.array(start_xy) - self.sys.theta = torch.tensor(math.atan2(d[1],d[0]), dtype=self.sys.dtype) - - self.stage_field = ssi.StageForceFieldTorch(r_safe=0.35, r_contact=0.18, - w_goal=1.0, w_radial=2.0, w_tangential=1.2, w_flow=0.8, w_boundary=0.6) - - def stage_slice(self, C,R,W) -> Tuple[np.ndarray,np.ndarray,np.ndarray]: - st = self.sm.current(); xmin,xmax,ymin,ymax = st.bounds - m = (C[:,0] >= xmin-0.5) & (C[:,0] <= xmax+0.5) & (C[:,1] >= ymin-0.5) & (C[:,1] <= ymax+0.5) - if np.any(m): return C[m],R[m],W[m] - return C,R,W - - def step(self, dt, world_obs: ssi.WorldObstacles, obs_weighs=None): - st = self.sm.current() - C,R,W = self.stage_slice(world_obs.C_np, world_obs.R_np, world_obs.W_np) - obs_t = ssi.ObstacleProviderTorch(C,R,W, dtype=self.sys.dtype) - barrier = ssi.IPCBarrier(obs_t, d_hat=world_obs.d_hat) - if obs_weighs is not None: - barrier.set_weights(obs_weighs) - info = self.sys.step(dt, obs_t, barrier, self.stage_field, st.bounds, tuple(st.exit_point)) - self.sm.advance_if_needed(np.array(info["center"],float)) - return info - - -# ============================== -# Episode / dataset I/O -# ============================== - -def perturb_robot_params(cfg: GenCfg) -> Dict[str, float]: - return { - "lam_reg": cfg.lam_reg * rand_uniform(*cfg.lam_reg_mult_range), - "mu": cfg.mu * rand_uniform(*cfg.mu_mult_range), - "d_hat": cfg.d_hat * rand_uniform(*cfg.d_hat_mult_range), - "radius": rand_uniform(*cfg.radius_abs_range), - } - -def stage_term_weights(stage_field) -> Dict[str, float]: - return { - "w_goal": float(getattr(stage_field, "w_goal", 1.0)), - "w_radial": float(getattr(stage_field, "w_radial", 2.0)), - "w_tangential": float(getattr(stage_field, "w_tangential", 1.0)), - "w_flow": float(getattr(stage_field, "w_flow", 0.6)), - "w_boundary": float(getattr(stage_field, "w_boundary", 0.6)), - } - -def episode_success(cfg: GenCfg, final_center: np.ndarray, min_d_history: List[float]) -> bool: - goal_dist = np.linalg.norm(np.asarray(final_center, float) - cfg.goal) - reached = (goal_dist <= cfg.goal_tol) - collision_free = (len(min_d_history) > 0) and (np.min(min_d_history) > cfg.min_clear_tol) - return bool(reached and collision_free) - -def save_episode_snapshot( - snapshot_path: str, - planner: PlannerMaximin, - frames: list, - world_obs: ssi.WorldObstacles, - start_xy: np.ndarray, - goal_xy: np.ndarray, - cfg: GenCfg, -): - os.makedirs(os.path.dirname(snapshot_path), exist_ok=True) - if not frames: - return None - start_Pw = frames[0]["Pw"].numpy() - end_Pw = frames[-1]["Pw"].numpy() - mid_Pw = frames[len(frames)//2]["Pw"].numpy() if len(frames) > 2 else None - - fig = plt.figure(figsize=(7.6, 3.8), dpi=140) - ax = fig.add_subplot(111) - ax.set_aspect("equal", "box") - - # raw obstacles - ssi.draw_obstacles(ax, world_obs.C_np, world_obs.R_np, ec="k") - # overlay inflated keep-out (what the planner respects) - if cfg.overlay_inflated_in_snapshot: - ssi.draw_obstacles(ax, world_obs.C_np, world_obs.R_np + cfg.inflate, ec="#00cfdc", alpha=0.35) - - ssi.render_curve(ax, planner.sys.B, start_Pw, "start", "#1f77b4") - if mid_Pw is not None: - ssi.render_curve(ax, planner.sys.B, mid_Pw, "mid", "#ff7f0e") - ssi.render_curve(ax, planner.sys.B, end_Pw, "end", "#2ca02c") - - for i, st in enumerate(planner.sm.stages): - xmin, xmax, ymin, ymax = st.bounds - ax.add_patch(plt.Rectangle((xmin,ymin),xmax-xmin,ymax-ymin, fill=False, - lw=1.6, ec="#00cfdc" if i==planner.sm.current_stage_idx else (0.4,0.6,0.6,0.6))) - ax.plot([st.entry_point[0]],[st.entry_point[1]],"o",ms=4,color="#85ff85") - ax.plot([st.exit_point[0]],[st.exit_point[1]],"o",ms=4,color="#ff8585") - - ax.plot([start_xy[0]],[start_xy[1]],"o",color="green",ms=6) - ax.plot([goal_xy[0]],[goal_xy[1]],"*",color="gold",ms=10) - ax.legend(loc="upper left", fontsize=8) - ax.set_title("Maximin (decoupled ROI) stage exits: start/mid/end") - ax.set_xlim(-2.5, 10.0); ax.set_ylim(-2.4, 2.4) - - fig.savefig(snapshot_path, bbox_inches="tight") - plt.close(fig) - return snapshot_path - -def planner_from_cfg( - cfg: GenCfg, - world_obs: ssi.WorldObstacles, - lam_reg: float, - mu: float, - d_hat: float, - radius: float -) -> PlannerMaximin: - planner = PlannerMaximin( - start_xy=cfg.start, - goal_xy=cfg.goal, - stage_size=cfg.stage_size, - overlap=cfg.overlap, - n_ctrl=cfg.n_ctrl, - K=cfg.K, - d_hat=d_hat, - lam_reg=lam_reg, - mu=mu, - seed=cfg.seed_robot, - obstacles=(world_obs.C_np, world_obs.R_np), - inflate=cfg.inflate, - graph_cfg=cfg, - ) - with torch.no_grad(): - sys = planner.sys - P0 = sys.P0_loc - cur_norm = torch.mean(torch.linalg.norm(P0, dim=-1)).item() + 1e-8 - scale = float(radius / cur_norm) - sys.P0_loc *= scale - sys.Ploc *= scale - return planner - -def run_episode( - cfg: GenCfg, - world_obs: ssi.WorldObstacles, - lam_reg: float, - mu: float, - d_hat: float, - radius: float, - snapshot_path: str, -) -> Dict[str, Any]: - planner = planner_from_cfg(cfg, world_obs, lam_reg, mu, d_hat, radius) - - frames = [] - energies = [] - frame_state = [] - min_d_history = [] - - meta = { - "start": cfg.start.tolist(), - "goal": cfg.goal.tolist(), - "stage_size": list(cfg.stage_size), - "overlap": float(cfg.overlap), - "inflate": float(cfg.inflate), - "dt": float(cfg.dt), - "total_steps": int(cfg.total_steps), - "n_ctrl": int(cfg.n_ctrl), - "K": int(cfg.K), - "graph": { - "grid_res": cfg.grid_res, - "knn_k": cfg.knn_k, - "nsamp_edge": cfg.nsamp_edge, - "min_node_clear": cfg.min_node_clear, - "follow_corridor": bool(cfg.follow_corridor), - "roi_pad_scale_x": cfg.roi_pad_scale_x, - "roi_pad_scale_y": cfg.roi_pad_scale_y, - "search_union_current_next": bool(cfg.search_union_current_next), - "min_edge_margin": cfg.min_edge_margin, - "min_exit_margin": cfg.min_exit_margin, - } - } - - tw = stage_term_weights(planner.stage_field) - - for t in range(cfg.total_steps): - info = planner.step(cfg.dt, world_obs) - min_d_history.append(float(info["min_d"])) - - if (t % cfg.snapshot_every) == 0 or t in (0, cfg.total_steps // 2, cfg.total_steps - 1): - sys = planner.sys - with torch.no_grad(): - Pw = sys.world_points().detach().cpu() - X = (sys.B @ Pw).detach().cpu() - frames.append({ - "t": t, - "Pw": Pw, - "X": X, - "o": sys.o.detach().cpu(), - "theta": float(sys.theta.item()), - "v_o": sys.v_o.detach().cpu(), - "omega": float(sys.omega.item()), - "Ploc": sys.Ploc.detach().cpu(), - "Vloc": sys.Vloc.detach().cpu(), - "center": sys.o.detach().cpu(), - }) - - if (t % cfg.snapshot_every) == 0: - st = planner.sm.current() - frame_state.append({ - "t": t, - "stage_idx": int(planner.sm.current_stage_idx), - "bounds": list(st.bounds), - "entry_point": st.entry_point.tolist(), - "exit_point": st.exit_point.tolist(), - "center": st.center.tolist(), - }) - - energies.append({ - "t": t, - "U_barrier": float(info["U_barrier"]), - "U_reg": float(info["U_reg"]), - "U_area": float(info["U_area"]), - "min_d": float(info["min_d"]), - }) - - final_center = planner.sys.o.detach().cpu().numpy() - final_theta = float(planner.sys.theta.item()) - success = episode_success(cfg, final_center, min_d_history) - - snap = save_episode_snapshot(snapshot_path, planner, frames, world_obs, cfg.start, cfg.goal, cfg) - - episode = { - "meta": meta, - "params": { - "lam_reg": float(lam_reg), - "mu": float(mu), - "d_hat": float(d_hat), - "radius": float(radius), - "gamma": float(planner.sys.gamma), - "gamma_th": float(planner.sys.gamma_th), - "gamma_o": float(planner.sys.gamma_o), - **tw, - }, - "obstacles": { - "centers": world_obs.C_np.astype(np.float32), - "radii": world_obs.R_np.astype(np.float32), - "weights": world_obs.W_np.astype(np.float32), - "d_hat_env": float(world_obs.d_hat), - }, - "frames": frames, - "energies": energies, - "frame_state": frame_state, - "success": bool(success), - "final_center": final_center.tolist(), - "final_theta": final_theta, - "snapshot_path": snap, - } - return episode - -def save_episode(root: str, case_tag: str, idx: int, episode: Dict[str, Any]) -> str: - os.makedirs(os.path.join(root, "episodes"), exist_ok=True) - path = os.path.join(root, "episodes", f"episode_{case_tag}_{idx:02d}.pt") - ep = dict(episode) - obs = ep["obstacles"] - obs_t = { - "centers": torch.from_numpy(obs["centers"]), - "radii": torch.from_numpy(obs["radii"]), - "weights": torch.from_numpy(obs["weights"]), - "d_hat_env": obs["d_hat_env"], - } - ep["obstacles"] = obs_t - torch.save(ep, path) - return path - -def write_manifest(root: str, records: List[Dict[str, Any]]): - man_path = os.path.join(root, "manifest.json") - with open(man_path, "w") as f: - json.dump(records, f, indent=2) - return man_path - - -# ============================== -# Dataset loader -# ============================== - -class NavigatorEpisodes(Dataset): - def __init__(self, root: str, filter_case: Optional[str] = None, success_only: Optional[bool] = None): - super().__init__() - manifest = os.path.join(root, "manifest.json") - with open(manifest, "r") as f: - self.records = json.load(f) - if filter_case is not None: - self.records = [r for r in self.records if r["case"] == filter_case] - if success_only is not None: - self.records = [r for r in self.records if r["success"] == bool(success_only)] - self.root = root - def __len__(self): return len(self.records) - def __getitem__(self, idx: int): - rec = self.records[idx] - ep = torch.load(rec["path"], map_location="cpu") - return ep - - -# ============================== -# Main -# ============================== - -def main(): - cfg = GenCfg() - set_all_seeds(cfg.seed) - - os.makedirs(cfg.out_root, exist_ok=True) - os.makedirs(os.path.join(cfg.out_root, "episodes"), exist_ok=True) - records: List[Dict[str, Any]] = [] - - # ---------- Case 1: environment variations (TIGHT) ---------- - for i in range(cfg.num_env_variations): - lam_reg, mu, d_hat, radius = cfg.lam_reg, cfg.mu, cfg.d_hat, cfg.radius - C, R, W = sample_obstacles_case1_tight(cfg) - world = ssi.WorldObstacles(C, R, W, d_hat=cfg.d_hat) - - snap_path = os.path.join(cfg.out_root, "episodes", f"episode_env_{i:02d}_snapshot.png") - ep = run_episode(cfg, world, lam_reg, mu, d_hat, radius, snapshot_path=snap_path) - path = save_episode(cfg.out_root, "env", i, ep) - records.append({ - "case": "env", - "idx": i, - "path": path, - "success": ep["success"], - "n_obstacles": int(world.C_np.shape[0]), - "snapshot": ep["snapshot_path"], - }) - print(f"[Case1 env #{i:02d}] saved: {path} | success={ep['success']} | n_obs={world.C_np.shape[0]}") - - # ---------- Case 2: robot variations (fixed obstacle field) ---------- - # Use one moderately tight field from the same generator but fixed RNG seed for reproducibility - rng_state = np.random.get_state() - np.random.seed(cfg.seed + 777) - C0, R0, W0 = sample_obstacles_case1_tight(cfg) - np.random.set_state(rng_state) - - world_fixed = ssi.WorldObstacles(C0, R0, W0, d_hat=cfg.d_hat) - for i in range(cfg.num_robot_variations): - p = perturb_robot_params(cfg) - snap_path = os.path.join(cfg.out_root, "episodes", f"episode_robot_{i:02d}_snapshot.png") - ep = run_episode(cfg, world_fixed, p["lam_reg"], p["mu"], p["d_hat"], p["radius"], snapshot_path=snap_path) - path = save_episode(cfg.out_root, "robot", i, ep) - records.append({ - "case": "robot", - "idx": i, - "path": path, - "success": ep["success"], - "lam_reg": p["lam_reg"], - "mu": p["mu"], - "d_hat": p["d_hat"], - "radius": p["radius"], - "snapshot": ep["snapshot_path"], - }) - print(f"[Case2 robot #{i:02d}] saved: {path} | success={ep['success']} | " - f"lam={p['lam_reg']:.3f} mu={p['mu']:.3f} d_hat={p['d_hat']:.3f} r={p['radius']:.3f}") - - man_path = write_manifest(cfg.out_root, records) - print(f"\nWrote manifest: {man_path}") - print("Done.") - - -if __name__ == "__main__": - main() diff --git a/scripts/train_sdf.py b/scripts/train_sdf.py deleted file mode 100644 index d9be890..0000000 --- a/scripts/train_sdf.py +++ /dev/null @@ -1,109 +0,0 @@ -"""Self-supervised training of the SDF navigation coefficients on a real scene. - -Reads a ``nav_sdf.npz`` (from ``build_sdf.py``) and trains ``sdf_nav.CoefMLP`` to -predict ``(alpha, beta, gamma)`` for the differentiable SDF surrogate by -backpropagating a reach-goal + no-collision objective through the rollout — no -labels. The coefficient net is biased toward the known-good navigating regime, so -it starts near-optimal and converges in a few hundred–thousand steps. - -Runs on CPU (the SDF removes the per-step obstacle-search cost of the circle -surrogate, so CPU is fast) or on GPU automatically if a CUDA torch is present. - -Usage: - python scripts/train_sdf.py /nav_sdf.npz -o coef_sdf.pt --steps 1500 -""" -from __future__ import annotations - -import argparse -import json -import os - -import numpy as np -import torch -import torch.nn.functional as F - -import sdf_nav - - -def main() -> None: - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument("sdf_npz") - ap.add_argument("-o", "--out", default="coef_sdf.pt") - ap.add_argument("--steps", type=int, default=1500) - ap.add_argument("--batch", type=int, default=128) - ap.add_argument("--horizon", type=int, default=28) - ap.add_argument("--dt", type=float, default=0.06) - ap.add_argument("--d-hat", type=float, default=0.35, help="wall barrier reach (normalized)") - ap.add_argument("--robot-radius-world", type=float, default=3.0) - ap.add_argument("--vmax", type=float, default=0.9) - ap.add_argument("--nsub-infer", type=int, default=3) - ap.add_argument("--lr", type=float, default=3e-4) - ap.add_argument("--threads", type=int, default=6) - ap.add_argument("--bundle", default=None, help="scene dir (for the eval penetration check); optional") - ap.add_argument("--seed", type=int, default=0) - args = ap.parse_args() - - dev = torch.device("cuda" if torch.cuda.is_available() else "cpu") - torch.set_num_threads(args.threads) - torch.manual_seed(args.seed) - np.random.seed(args.seed) - - d = np.load(args.sdf_npz) - phi, nxg, nyg = d["phi"], d["normal_x"], d["normal_y"] - bounds = [float(v) for v in d["bounds"]] - center = [float(v) for v in d["center"]] - S = float(d["scale"]); region = float(d["region"]) - rr = args.robot_radius_world * S - field = sdf_nav.SDFField(phi, nxg, nyg, bounds, center, S, device=dev) - mnx, mny, mxx, mxy = bounds - print("device=%s field=%s phi[%.2f,%.2f] rr=%.3f S=%.5f" % (dev, phi.shape, phi.min(), phi.max(), rr, S)) - - # free-point pool (normalized, centered) within the working region - ny_, nx_ = phi.shape - gxs = np.linspace(mnx, mxx, nx_); gys = np.linspace(mny, mxy, ny_) - GX, GY = np.meshgrid(gxs, gys) - free = phi > (rr + 0.02) # drivable = clearance beyond the robot - reg = (np.abs(GX - center[0]) < region) & (np.abs(GY - center[1]) < region) - sel = free & reg - pool = np.stack([(GX[sel] - center[0]) * S, (GY[sel] - center[1]) * S], 1).astype(np.float32) - poolt = torch.from_numpy(pool).to(dev) - print("pool=%d" % len(pool)) - - model = sdf_nav.CoefMLP().to(dev) - opt = torch.optim.Adam(model.parameters(), lr=args.lr) - kw = dict(rr=rr, d_hat=args.d_hat, dt=args.dt, vmax=args.vmax) - import time - t0 = time.time() - for it in range(args.steps): - si = torch.randint(0, len(poolt), (args.batch,), device=dev) - o0 = poolt[si] - ang = torch.rand(args.batch, device=dev) * 6.2832 - dd = 1.0 + torch.rand(args.batch, device=dev) # 1..2 local goal - goal = o0 + torch.stack([torch.cos(ang), torch.sin(ang)], -1) * dd.unsqueeze(-1) - al, be, ga = model(sdf_nav.coef_feats(field, o0, goal)) - oT, vT, clr = sdf_nav.sdf_rollout(field, o0, torch.zeros_like(o0), goal, al, be, ga, - args.horizon, nsub=1, **kw) - L_goal = ((oT - goal) ** 2).sum(-1).mean() - L_col = F.softplus((0.02 - clr) / 0.02).mean() - L_reg = ((be - 3.0) ** 2).mean() + ((ga - 4.0) ** 2).mean() + ((al - 1.0) ** 2).mean() - loss = L_goal + 3.0 * L_col + 0.1 * L_reg - opt.zero_grad(); loss.backward() - torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step() - if it % 300 == 0: - print("it=%d/%d L_goal=%.3f L_col=%.3f al=%.2f be=%.2f ga=%.2f (%.0fs)" - % (it, args.steps, float(L_goal), float(L_col), float(al.mean()), - float(be.mean()), float(ga.mean()), time.time() - t0), flush=True) - tt = time.time() - t0 - print("TIME %d steps %.1fs = %.1f ms/step on %s" % (args.steps, tt, 1000 * tt / args.steps, dev)) - - meta = {"scale": S, "center": center, "region": region, "rr": rr, "d_hat": args.d_hat, - "dt": args.dt, "horizon": args.horizon, "nsub": args.nsub_infer, "vmax": args.vmax, - "steps": args.steps, "bounds": bounds, "kind": "sdf", "sdf_npz": os.path.abspath(args.sdf_npz)} - torch.save({"model_state_dict": model.state_dict(), "meta": meta}, args.out) - print("SAVED %s" % args.out) - json.dump({"steps": args.steps, "ms_per_step": 1000 * tt / args.steps, "device": str(dev)}, - open(os.path.splitext(args.out)[0] + "_result.json", "w")) - - -if __name__ == "__main__": - main() diff --git a/scripts/working_spline_pathwise.py b/scripts/working_spline_pathwise.py deleted file mode 100644 index 8032cca..0000000 --- a/scripts/working_spline_pathwise.py +++ /dev/null @@ -1,501 +0,0 @@ -# stagewise_spline_icp_with_stage_field.py -# Deformable periodic B-spline with ICP barrier + stage force field (goal+radial+tangent+flow+boundary) -# Produces a snapshot (and optional GIF) showing start/mid/end. - -from __future__ import annotations -import os, math, argparse -from dataclasses import dataclass -from typing import List, Tuple, Optional -import numpy as np -import torch - -import matplotlib -matplotlib.use("Agg") -import matplotlib.pyplot as plt - -# --------------------------- utils --------------------------- - -def rot2(theta: torch.Tensor): - c = torch.cos(theta); s = torch.sin(theta) - return torch.stack([torch.stack([c,-s],-1), torch.stack([s,c],-1)], -2) - -# --------------------------- stages (numpy) --------------------------- - -@dataclass -class Stage: - center: np.ndarray - size: Tuple[float,float] - entry_point: np.ndarray - exit_point: np.ndarray - stage_id: int - def __post_init__(self): - W,H = self.size; cx,cy = self.center - self.bounds = (cx-W/2, cx+W/2, cy-H/2, cy+H/2) - -class StageManager: - def __init__(self, start_xy: np.ndarray, goal_xy: np.ndarray, stage_size=(2.6,2.0), overlap_ratio=0.3): - self.start = np.array(start_xy,float); self.goal = np.array(goal_xy,float) - self.stage_size = tuple(stage_size); self.overlap_ratio = float(overlap_ratio) - self.stages: List[Stage] = []; self.current_stage_idx = 0; self._build() - def _build(self): - d = self.goal - self.start; L = np.linalg.norm(d) - if L < 1e-6: - self.stages = [Stage((self.start+self.goal)/2,self.stage_size,self.start.copy(),self.goal.copy(),0)] - return - u = d / L; step = self.stage_size[0]*(1-self.overlap_ratio); n=max(1,int(np.ceil(L/step))) - for i in range(n): - t = 0.0 if n==1 else i/(n-1) - c = self.start + t*d - entry = self.start if i==0 else self.stages[-1].exit_point - exitp = self.goal if i==n-1 else (self.start + (i+1)/n*d) - self.stages.append(Stage(c,self.stage_size,entry,exitp,i)) - def current(self)->Stage: return self.stages[self.current_stage_idx] - def advance_if_needed(self, robot_xy: np.ndarray, thresh=0.45)->bool: - st = self.current() - if np.linalg.norm(robot_xy - st.exit_point) < thresh and self.current_stage_idx < len(self.stages)-1: - self.current_stage_idx += 1; return True - if self.current_stage_idx < len(self.stages)-1: - nxt = self.stages[self.current_stage_idx+1]; path_dir = self.goal - self.start - if np.dot(robot_xy - st.center, path_dir) > np.dot(nxt.center - st.center, path_dir): - self.current_stage_idx += 1; return True - return False - -# --------------------------- obstacles & ICP (torch) --------------------------- - -class ObstacleProviderTorch: - def __init__(self, centers, radii, weights=None, device="cpu", dtype=torch.float32, eps=1e-12): - self.C = torch.as_tensor(centers, dtype=dtype, device=device) # (C,2) - self.R = torch.as_tensor(radii, dtype=dtype, device=device) # (C,) - if weights is None: weights = torch.ones_like(self.R) - self.W = torch.as_tensor(weights, dtype=dtype, device=device) - self.eps = float(eps) - def compute_all(self, q: torch.Tensor): - diff = q.unsqueeze(-2) - self.C # (...,C,2) - r = torch.linalg.norm(diff, dim=-1, keepdim=True).clamp_min(1e-12) - d = r.squeeze(-1) - self.R # (...,C) - g = diff / r # (...,C,2) - return d, g - -class IPCBarrier: - def __init__(self, obstacles: ObstacleProviderTorch, d_hat=1.0, violation_penalty=5e2, eps=1e-12): - self.obs = obstacles; self.d_hat=float(d_hat); self.vp=float(violation_penalty); self.eps=float(eps) - def _piecewise(self, d: torch.Tensor): - dh = d.new_tensor(self.d_hat); vp = d.new_tensor(self.vp); safe = torch.clamp(d, min=self.eps) - b_in = -(d-dh)**2 * torch.log(safe/dh) - F_in = (dh-d)*(2.0*torch.log(safe/dh) - dh/safe) + 1.0 - b = torch.where(d<=self.eps, vp, torch.where(d (K,2) - if C.numel() == 0: return torch.zeros_like(X) - K, Cn = X.shape[0], C.shape[0] - diff = X.unsqueeze(1) - C.unsqueeze(0) # (K,C,2) - dist = torch.linalg.norm(diff, dim=-1).clamp_min(1e-8) # (K,C) - d = dist - R.unsqueeze(0) # (K,C) - # very strong contact push - contact = (d <= self.r_contact).unsqueeze(-1) # (K,C,1) - strong = (diff / dist.unsqueeze(-1)) * contact * 1e4 # (K,C,2) - # ICP-like gradient in [r_contact, r_safe] - d_ref = X.new_tensor(self.r_safe) - in_band = ((d > self.r_contact) & (d < d_ref)).unsqueeze(-1) # (K,C,1) - db_dd = -(2*(d - d_ref)*torch.log(d.clamp_min(1e-8)/d_ref) - + (d - d_ref)**2 / d.clamp_min(1e-8)) # (K,C) - grad = db_dd.unsqueeze(-1) * (diff / dist.unsqueeze(-1)) # (K,C,2) - grad = torch.where(in_band, grad, torch.zeros_like(grad)) # (K,C,2) - return (strong + grad).sum(dim=1) # (K,2) - - def _tangent(self, X, target, C, R): - if C.numel() == 0: return torch.zeros_like(X) - K, Cn = X.shape[0], C.shape[0] - diff = X.unsqueeze(1) - C.unsqueeze(0) # (K,C,2) - dist = torch.linalg.norm(diff, dim=-1).clamp_min(1e-8) # (K,C) - d = dist - R.unsqueeze(0) # (K,C) - n_hat = diff / dist.unsqueeze(-1) # (K,C,2) - # rotate by +90° - J = self.J.to(device=X.device, dtype=X.dtype) - t_hat = (J @ n_hat.transpose(-1, -2)).transpose(-1, -2) # (K,C,2) - - # goal alignment sign - gdir = (target.unsqueeze(0) - X) # (K,2) - gdir = gdir / torch.linalg.norm(gdir, dim=1, keepdim=True).clamp_min(1e-8) - gdir = gdir.unsqueeze(1).expand(-1, Cn, -1) # (K,C,2) - sign = torch.sign((t_hat * gdir).sum(dim=-1, keepdim=True)) # (K,C,1) - - # strength - inner = torch.clamp(self.r_safe - d, min=0.0) / self.r_safe # (K,C) - strength = torch.where(d < self.r_safe, inner, X.new_tensor(0.1))# (K,C) - strength = strength.unsqueeze(-1) # (K,C,1) - influence = (d < (self.r_safe + 0.5)).unsqueeze(-1).float() # (K,C,1) - - F = sign * strength * t_hat * influence # (K,C,2) - return F.sum(dim=1) # (K,2) - - def _flow(self, X, target, C, R): - if C.numel() == 0: return torch.zeros_like(X) - K, Cn = X.shape[0], C.shape[0] - diff = X.unsqueeze(1) - C.unsqueeze(0) # (K,C,2) - dist = torch.linalg.norm(diff, dim=-1).clamp_min(1e-8) # (K,C) - d = dist - R.unsqueeze(0) # (K,C) - n_hat = diff / dist.unsqueeze(-1) # (K,C,2) - J = self.J.to(device=X.device, dtype=X.dtype) - t_hat = (J @ n_hat.transpose(-1, -2)).transpose(-1, -2) # (K,C,2) - - obs_to_goal = (target - C) # (C,2) - goal_unit = obs_to_goal / torch.linalg.norm(obs_to_goal, dim=-1, keepdim=True).clamp_min(1e-8) # (C,2) - goal_unit = goal_unit.unsqueeze(0).expand(K, -1, -1) # (K,C,2) - sign = torch.sign((t_hat * goal_unit).sum(dim=-1, keepdim=True)) # (K,C,1) - flow_dir = sign * t_hat # (K,C,2) - - # blend a small radial component when far - far = (d > self.r_safe).float().unsqueeze(-1) # (K,C,1) - radial_comp = goal_unit # (K,C,2) - flow_dir = 0.7*flow_dir + 0.3*far*radial_comp # (K,C,2) - - flow_strength = torch.where(d < self.r_contact, X.new_tensor(1.0), - torch.where(d < self.r_safe, X.new_tensor(0.8), - torch.clamp(0.3*(1.0-(d-self.r_safe)), min=0.0))) # (K,C) - F = flow_dir * flow_strength.unsqueeze(-1) # (K,C,2) - return F.sum(dim=1) # (K,2) - - def _boundary(self, X, bounds): - xmin,xmax,ymin,ymax = [X.new_tensor(v) for v in bounds] - m = X.new_tensor(0.5) - F = torch.zeros_like(X) # (K,2) - left = X[:,0] < (xmin+m); right = X[:,0] > (xmax-m) - bot = X[:,1] < (ymin+m); top = X[:,1] > (ymax-m) - F[left,0] += ((xmin+m) - X[left,0]) / m * 2.0 - F[right,0] += (-(X[right,0] - (xmax-m)) / m) * 2.0 - F[bot,1] += ((ymin+m) - X[bot,1]) / m * 2.0 - F[top,1] += (-(X[top,1] - (ymax-m)) / m) * 2.0 - return F - - # ---- full field ---- - def field(self, X: torch.Tensor, target: torch.Tensor, bounds, C: torch.Tensor, R: torch.Tensor) -> torch.Tensor: - # Everything reduced to (K,2) before the sum - Fg = self._goal(X, target) # (K,2) - Fr = self._radial(X, C, R) # (K,2) - Ft = self._tangent(X, target, C, R) # (K,2) - Ff = self._flow(X, target, C, R) # (K,2) - Fb = self._boundary(X, bounds) # (K,2) - return Fg - # return (self.w_goal*Fg + self.w_radial*Fr + self.w_tangential*Ft - # + self.w_flow*Ff + self.w_boundary*Fb) # (K,2) - -# --------------------------- deformable system --------------------------- - -class DeformableSplineSystem: - def __init__(self, n_ctrl=20, K=220, device="cpu", dtype=torch.float32, - lam_reg=20.0, mu=0.25, d_hat=1.0, - M=0.1, gamma=4.0, I=0.6, gamma_th=1.6, M_o=1.5, gamma_o=4.0, - seed=7, radius=0.25): - torch.manual_seed(seed); self.device=device; self.dtype=dtype - self.B,self.D = periodic_cubic_uniform_matrices(n_ctrl,K,device,dtype) - self.w = torch.full((K,), 1.0/float(K), dtype=dtype, device=device) - ang = torch.linspace(0,2*math.pi,n_ctrl+1,device=device,dtype=dtype)[:-1] - self.P0_loc = radius * torch.stack([torch.cos(ang), torch.sin(ang)],-1) - self.Ploc = self.P0_loc.clone(); self.Vloc = torch.zeros_like(self.Ploc) - self.theta = torch.tensor(0.0,dtype=dtype,device=device); self.omega = torch.tensor(0.0,dtype=dtype,device=device) - self.o = torch.tensor([0.0,0.0],dtype=dtype,device=device); self.v_o = torch.tensor([0.0,0.0],dtype=dtype,device=device) - self.lam_reg=float(lam_reg); self.mu=float(mu); self.d_hat=float(d_hat) - self.M=float(M); self.gamma=float(gamma); self.I=float(I); self.gamma_th=float(gamma_th) - self.M_o=float(M_o); self.gamma_o=float(gamma_o) - self.J = torch.tensor([[0.,-1.],[1.,0.]], dtype=dtype, device=device) - self._prev_o = self.o.clone(); self._no_prog = 0; self._stuck=False; self.A0=None - - def world_points(self): - R = rot2(self.theta); return (self.Ploc @ R.T) + self.o - - def sample_curve(self): - Pw = self.world_points(); X = self.B @ Pw; Xp = self.D @ Pw - ell = Xp.norm(dim=1).clamp_min(1e-8); T_hat = Xp/ell.unsqueeze(-1) - return Pw,X,Xp,ell,T_hat - - def step( - self, - dt: float, - obstacle_provider: ObstacleProviderTorch, - barrier: IPCBarrier, - stage_force, # StageForceFieldTorch - stage_bounds, # (xmin, xmax, ymin, ymax) - target_xy: Tuple[float, float] - ): - # --- samples and tangents --- - Pw, X, Xp, ell, T_hat = self.sample_curve() - w = self.w - R = rot2(self.theta) - J = self.J.to(device=X.device, dtype=X.dtype) - target = torch.as_tensor(target_xy, dtype=self.dtype, device=self.device) - - # --- ICP barrier (conservative) --- - bsum, g_sum = barrier.barrier_and_grad(X) # (K,), (K,2) - term1 = (w * ell).unsqueeze(-1) * g_sum # (K,2) - term2 = (w * bsum).unsqueeze(-1) * T_hat # (K,2) - gradP_world = self.B.T @ term1 + self.D.T @ term2 # (n_ctrl,2) - - # --- Adaptive area control (conservative) --- - A, dA_dX = polygon_area_and_grad(X) # scalar, (K,2) - d_all, _ = obstacle_provider.compute_all(X) # (K,C) - min_d_each = d_all.min(dim=1).values - min_clear = min_d_each.min() - - # Robust A0 init (handles missing or None or non-finite) - A0_val = getattr(self, "A0", None) - if (A0_val is None) or (not np.isfinite(A0_val)): - self.A0 = float(torch.abs(A.detach())) + 1e-6 - A0_t = X.new_tensor(self.A0) - - alpha, beta = 0.25, 2.5 - squeeze_factor = alpha + (1 - alpha) * torch.tanh(beta * torch.clamp(min_clear, min=0.0)) - A_target = squeeze_factor * A0_t - - k_area = 1.5 - gradX_area = 2.0 * k_area * (A - A_target) * dA_dX / X.shape[0] # (K,2) - gradP_world = gradP_world + self.B.T @ (w.unsqueeze(-1) * gradX_area) - - # --- chain to local + shape regularizer --- - gradP_loc = (gradP_world @ R) + self.lam_reg * (self.Ploc - self.P0_loc) - F_loc_cons = -gradP_loc - - # --- conservative rigid forces (barrier + area) --- - X_minus_o = X - self.o.unsqueeze(0) - dU_do = torch.sum(term1, dim=0) + torch.sum(w.unsqueeze(-1) * gradX_area, dim=0) # (2,) - F_o_cons = -dU_do - gJx = torch.einsum("kd,dd,kd->k", g_sum, J, X_minus_o) - dU_dtheta = torch.sum(w * ell * gJx) - g_area_J = torch.einsum("kd,dd,kd->k", gradX_area, J, X_minus_o) - dU_dtheta = dU_dtheta + torch.sum(w * g_area_J) - tau_cons = -dU_dtheta - - # --- contact friction (dissipative) --- - dists, _ = obstacle_provider.compute_all(X) # (K,C) - min_d = dists.min(dim=1).values # (K,) - contact_mask = (min_d < barrier.d_hat).to(X.dtype) - Pdot_world = (self.Vloc @ R.T) + self.omega * ((self.Ploc @ R.T) @ J.T) + self.v_o - Xdot = self.B @ Pdot_world # (K,2) - v_t = torch.sum(Xdot * T_hat, dim=1) # (K,) - p_k = g_sum.norm(dim=1) # (K,) - f_fric = -self.mu * (p_k * v_t * contact_mask).unsqueeze(-1) * T_hat # (K,2) - F_world_from_fric = self.B.T @ (w.unsqueeze(-1) * f_fric) # (n_ctrl,2) - F_loc_from_fric = F_world_from_fric @ R.T # (n_ctrl,2) - F_o_from_fric = torch.sum(w.unsqueeze(-1) * f_fric, dim=0) # (2,) - tau_fric = torch.sum(w * torch.sum(f_fric * (J @ X_minus_o.T).T, dim=1)) - - # --- stage force field (non-conservative) in sample space --- - F_stage = stage_force.field(X, target, stage_bounds, obstacle_provider.C, obstacle_provider.R) # (K,2) - - F_stage = 1.5 * F_stage - - # Map stage forces to generalized coords - F_world_from_stage = self.B.T @ (w.unsqueeze(-1) * F_stage) # (n_ctrl,2) - F_loc_from_stage = F_world_from_stage @ R.T # (n_ctrl,2) - F_o_from_stage = torch.sum(w.unsqueeze(-1) * F_stage, dim=0) # (2,) - tau_stage = torch.sum(w * torch.sum(F_stage * (J @ X_minus_o.T).T, dim=1)) - - # --- integrate dynamics --- - F_loc_total = F_loc_cons + F_loc_from_fric + F_loc_from_stage - self.Vloc = self.Vloc + dt * (F_loc_total / self.M - self.gamma * self.Vloc) - self.Ploc = self.Ploc + dt * self.Vloc - - tau_total = tau_cons + tau_fric + tau_stage - self.gamma_th * self.omega - self.omega = self.omega + dt * (tau_total / self.I) - self.theta = self.theta + dt * self.omega - - F_o_total = F_o_cons + F_o_from_fric + F_o_from_stage - self.gamma_o * self.v_o - self.v_o = self.v_o + dt * (F_o_total / self.M_o) - self.o = self.o + dt * self.v_o - # --- diagnostics --- - - U_barrier = torch.sum(w * bsum * ell) - U_reg = 0.5 * self.lam_reg * ((self.Ploc - self.P0_loc) ** 2).sum() - U_area = float((k_area * (A - A_target) ** 2).item()) - - return { - "U_barrier": float(U_barrier.item()), - "U_reg": float(U_reg.item()), - "U_area": U_area, - "center": self.o.detach().cpu().numpy().tolist(), - "theta": float(self.theta.item()), - "min_d": float(min_d.min().item()), - } - - -# --------------------------- world obstacles holder --------------------------- - -class WorldObstacles: - def __init__(self, centers, radii, weights=None, d_hat=0.28): - self.C_np=np.asarray(centers,float); self.R_np=np.asarray(radii,float) - self.W_np=np.asarray(weights if weights is not None else np.ones_like(self.R_np),float) - self.d_hat=float(d_hat) - -# --------------------------- viz helpers --------------------------- - -def draw_obstacles(ax, C,R, **kw): - C=np.asarray(C); R=np.asarray(R) - for (cx,cy),rr in zip(C,R): - ax.add_patch(plt.Circle((cx,cy), rr, fill=False, lw=2, **kw)) - -def render_curve(ax, B, Pw, label=None, color=None): - Xv=(B @ torch.tensor(Pw)).numpy(); ax.plot(Xv[:,0],Xv[:,1],lw=2,label=label,color=color) - -# --------------------------- planner wrapper --------------------------- - -class StagewiseSplinePlanner: - def __init__(self, start_xy, goal_xy, stage_size=(2.5,2.0), overlap=0.3, - n_ctrl=20, K=240, d_hat=1.0, lam_reg=20.0, mu=0.25, seed=7): - self.sm = StageManager(np.array(start_xy), np.array(goal_xy), stage_size, overlap) - self.sys = DeformableSplineSystem(n_ctrl=n_ctrl, K=K, d_hat=d_hat, lam_reg=lam_reg, mu=mu, seed=seed) - # initialize pose - self.sys.o = torch.tensor(start_xy, dtype=self.sys.dtype) - d = np.array(goal_xy)-np.array(start_xy); self.sys.theta = torch.tensor(math.atan2(d[1],d[0]), dtype=self.sys.dtype) - self.stage_field = StageForceFieldTorch(r_safe=0.35, r_contact=0.18, - w_goal=1.0, w_radial=2.0, w_tangential=1.4, w_flow=0.9, w_boundary=0.5) - - def _stage_slice(self, C,R,W) -> Tuple[np.ndarray,np.ndarray,np.ndarray]: - st = self.sm.current(); xmin,xmax,ymin,ymax = st.bounds - m = (C[:,0] >= xmin-0.5) & (C[:,0] <= xmax+0.5) & (C[:,1] >= ymin-0.5) & (C[:,1] <= ymax+0.5) - if np.any(m): return C[m],R[m],W[m] - return C,R,W - - def step(self, dt, world_obs: WorldObstacles): - st = self.sm.current(); C,R,W = self._stage_slice(world_obs.C_np, world_obs.R_np, world_obs.W_np) - obs_t = ObstacleProviderTorch(C,R,W, dtype=self.sys.dtype) - barrier = IPCBarrier(obs_t, d_hat=world_obs.d_hat) - info = self.sys.step(dt, obs_t, barrier, self.stage_field, st.bounds, tuple(st.exit_point)) - self.sm.advance_if_needed(np.array(info["center"],float)) - return info - -# --------------------------- demo --------------------------- - -def run_demo(out_dir="./out_stage_spline", total_steps=900, dt=0.015, - stage_size=(2.6,2.0), overlap=0.3, snapshot_every=5): - os.makedirs(out_dir,exist_ok=True) - # centers=np.array([[1.0,0.5],[2.5,-0.8],[3.2,0.3],[4.0,1.0],[4.8,-0.5], - # [5.5,0.8],[6.2,-0.2],[6.8,0.9],[7.5,-0.7],[8.2,0.4], - # [2.0,1.2],[3.5,-1.3],[5.0,1.5],[6.5,-1.1],[7.8,1.2]],float) - # radii=np.array([0.40,0.50,0.30,0.60,0.40,0.50,0.30,0.40,0.50,0.30,0.30,0.40,0.30,0.40,0.30],float) - # weights=np.array([1.0,1.2,0.8,1.5,0.9,1.1,0.7,1.0,1.3,0.8,0.6,1.0,0.7,1.1,0.6],float) - centers=np.array([[2.5,-0.8],[3.2,0.3],[4.0,1.0],[4.8,-0.5], - [5.5,0.8],[6.2,-0.2],[6.8,0.9],[7.5,-0.7],[8.2,0.4], - [2.0,1.2],[3.5,-1.3],[5.0,1.5],[6.5,-1.1],[7.8,1.2]],float) - radii=np.array([0.50,0.30,0.60,0.40,0.50,0.30,0.40,0.50,0.30,0.30,0.40,0.30,0.40,0.30],float) - weights=np.array([1.2,0.8,1.5,0.9,1.1,0.7,1.0,1.3,0.8,0.6,1.0,0.7,1.1,0.6],float) - world_obs=WorldObstacles(centers,radii,weights,d_hat=1.0) - - start=np.array([-1.0,-0.5]); goal=np.array([9.0,0.2]) - planner=StagewiseSplinePlanner(start, goal, stage_size=stage_size, overlap=overlap, - n_ctrl=20, K=240, d_hat=1.0, lam_reg=20.0, mu=0.25, seed=7) - - frames=[]; mid=None - for t in range(total_steps): - planner.step(dt, world_obs) - if t%snapshot_every==0 or t in (0,total_steps//2,total_steps-1): - Pw=planner.sys.world_points().detach().cpu().numpy() - frames.append(Pw) - if t==total_steps//2: mid=Pw.copy() - - # snapshot - fig=plt.figure(figsize=(7.6,3.8),dpi=140); ax=fig.add_subplot(111); ax.set_aspect("equal","box") - draw_obstacles(ax,centers,radii, ec="k") - if frames: render_curve(ax, planner.sys.B, frames[0], "start", "#1f77b4") - if mid is not None: render_curve(ax, planner.sys.B, mid, "mid", "#ff7f0e") - render_curve(ax, planner.sys.B, frames[-1], "end", "#2ca02c") - for i,st in enumerate(planner.sm.stages): - xmin,xmax,ymin,ymax=st.bounds - ax.add_patch(plt.Rectangle((xmin,ymin),xmax-xmin,ymax-ymin, fill=False, - lw=1.6, ec="#00cfdc" if i==planner.sm.current_stage_idx else (0.4,0.6,0.6,0.6))) - ax.plot([st.entry_point[0]],[st.entry_point[1]],"o",ms=4,color="#85ff85") - ax.plot([st.exit_point[0]],[st.exit_point[1]],"o",ms=4,color="#ff8585") - ax.plot([start[0]],[start[1]],"o",color="green",ms=6); ax.plot([goal[0]],[goal[1]],"*",color="gold",ms=10) - ax.legend(loc="upper left",fontsize=8) - ax.set_title("Stagewise deformable spline planning: ICP + stage field (goal+radial+tangent+flow) + area") - ax.set_xlim(-2.5,10.0); ax.set_ylim(-2.4,2.4) - snap=os.path.join(out_dir,"stagewise_spline_snapshot.png"); fig.savefig(snap,bbox_inches="tight"); plt.close(fig) - - # gif (optional) - gif=os.path.join(out_dir,"stagewise_spline_motion.gif") - try: - import imageio - imgs=[] - for Pw in frames: - fig=plt.figure(figsize=(6.4,3.6),dpi=110); ax=fig.add_subplot(111); ax.set_aspect("equal","box") - draw_obstacles(ax,centers,radii, ec="k"); render_curve(ax, planner.sys.B, Pw) - ax.set_xlim(-2.5,10.0); ax.set_ylim(-2.4,2.4); ax.set_xticks([]); ax.set_yticks([]) - fig.canvas.draw() - buf=np.frombuffer(fig.canvas.tostring_argb(),dtype=np.uint8) - arr=buf.reshape(fig.canvas.get_width_height()[::-1]+(4,))[:,:,1:]; imgs.append(arr); plt.close(fig) - imageio.mimsave(gif, imgs, fps=max(6,int(1.0/dt))) - except Exception: - gif=None - - return {"snapshot":snap, "gif":gif, "stages":len(planner.sm.stages), - "final_center":planner.sys.o.detach().cpu().numpy().tolist(), - "final_theta":float(planner.sys.theta.item())} - -# --------------------------- CLI --------------------------- - -def parse_args(): - ap=argparse.ArgumentParser("Stagewise deformable spline with stage field") - ap.add_argument("--out",type=str,default="./out_stage_spline") - ap.add_argument("--steps",type=int,default=3000) - ap.add_argument("--dt",type=float,default=0.02) - ap.add_argument("--stage_w",type=float,default=2.6) - ap.add_argument("--stage_h",type=float,default=2.0) - ap.add_argument("--overlap",type=float,default=0.3) - return ap.parse_args() - -if __name__=="__main__": - args=parse_args() - out=run_demo(out_dir=args.out, total_steps=args.steps, dt=args.dt, - stage_size=(args.stage_w,args.stage_h), overlap=args.overlap) - print("Outputs:"); print(out) diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..e46a899 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,45 @@ +"""The grl-snam CLI wiring: commands registered, help works, demos resolvable.""" + +from click.testing import CliRunner + +from grl_snam import demos +from grl_snam.cli import main + + +def test_cli_help_lists_pipeline(): + result = CliRunner().invoke(main, ["--help"]) + assert result.exit_code == 0 + assert "pipeline" in result.output + + +def test_cli_registers_core_commands(): + names = set(main.commands) + for cmd in ( + "selftest", + "obstacles", + "build-sdf", + "train", + "capture", + "pipeline", + "demo", + "lab-demo", + "eval", + "train-coef", + ): + assert cmd in names, f"missing command: {cmd}" + + +def test_capture_subcommands(): + assert {"drive", "multigoal"} <= set(main.commands["capture"].commands) + + +def test_demo_registry_and_paths_resolve(): + reg = demos.registry() + assert "austin-freedrive" in reg and "lab" in reg + for name in reg: + path = demos.demo_path(name) + assert path and path.endswith(".py") + + +def test_demo_path_unknown_is_none(): + assert demos.demo_path("does-not-exist") is None diff --git a/tests/test_lab.py b/tests/test_lab.py index 0a4400f..45160c8 100644 --- a/tests/test_lab.py +++ b/tests/test_lab.py @@ -1,66 +1,47 @@ -"""Tests for grl_snam_lab. +"""Tests for grl_snam.demos.lab — the analytic lab demo. -The geometry helpers are pure Python and always run; the end-to-end Lab test -is skipped where the compiled pycvc / pycvc_gl bindings aren't importable. +The pure geometry (terrain height, agent loop, track) always runs; the end-to-end +scene build is skipped where the compiled pycvc / pycvc_gl bindings aren't importable. """ import pytest -from grl_snam_lab import terrain_mesh -from grl_snam_lab.lab import _flatten_points, polyline_indices +from grl_snam.demos import lab -def test_terrain_mesh_counts_and_placement(): - heights = [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]] # 2 rows x 3 cols - verts, tris = terrain_mesh(heights, bounds=(-10, -20, 10, 20)) - assert len(verts) == 2 * 3 * 3 # 6 vertices * xyz - assert len(tris) == (2 - 1) * (3 - 1) * 6 # 2 cells * 2 tris * 3 idx - # first vertex at (min_x, min_y, height[0][0]); last at (max_x, max_y, h[-1][-1]) - assert verts[0:3] == [-10.0, -20.0, 0.0] - assert verts[-3:] == [10.0, 20.0, 5.0] +def test_terrain_height_peaks_at_centre(): + assert lab.terrain_height(0.0, 0.0) > lab.terrain_height(95.0, 95.0) -def test_terrain_mesh_rejects_degenerate(): - with pytest.raises(ValueError): - terrain_mesh([[1.0]], bounds=(0, 0, 1, 1)) +def test_track_is_closed_and_draped(): + assert len(lab.TRACK) == lab._N_TRACK + 1 + assert lab.TRACK[0] == pytest.approx(lab.TRACK[-1], abs=1e-9) # closed loop + x, y, z = lab.TRACK[0] + assert abs(z - (lab.terrain_height(x, y) + lab.AGENT_LIFT)) < 1e-6 -def test_polyline_indices(): - assert polyline_indices(4) == [0, 1, 1, 2, 2, 3] - assert polyline_indices(1) == [] +def test_agent_position_wraps_one_lap(): + a = lab.agent_position(0.0) + b = lab.agent_position(lab.LOOP_SECONDS) # exactly one lap -> back to start + assert all(abs(u - v) < 1e-6 for u, v in zip(a, b)) -def test_flatten_points(): - buf, n = _flatten_points([(1, 2, 3), (4, 5, 6)]) - assert n == 2 and buf == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] +def test_terrain_heights_grid_shape(): + g = lab._terrain_heights(8) + assert len(g) == 8 and all(len(row) == 8 for row in g) # ── end-to-end (requires the compiled bindings) ───────────────────────── - pytest.importorskip("pycvc", reason="pycvc bindings not installed") pytest.importorskip("pycvc_gl", reason="pycvc_gl bindings not installed") -def test_lab_builds_a_scene(): - from grl_snam_lab import Lab - - lab = Lab() - lab.add_terrain([[0, 0, 0], [0, 1, 0], [0, 0, 0]], bounds=(-5, -5, 5, 5)) - lab.add_mesh( - "bldg", - [0, 0, 0, 2, 0, 0, 0, 2, 0], - [0, 1, 2], - color=(0.6, 0.6, 0.6), - ) - lab.add_path("agent0", [(-4, 0, 1), (0, 0, 1), (4, 2, 1)], color=(1, 0, 0)) - lab.add_markers("agents", [(-4, 0, 1), (4, 2, 1)]) - lab.add_field( - "risk", - [float(i % 5) for i in range(3 * 3 * 3)], - dims=(3, 3, 3), - bounds=(-5, -5, 0, 5, 5, 4), - ) - lab.pump() - for name in ("terrain", "bldg", "agent0", "agents", "risk"): - assert lab.has(name), f"missing node {name}" - assert lab.num_nodes() == 5 +def test_lab_demo_builds_a_scene(): + # Build the scene graph and assert its nodes — but do NOT render_png here: offscreen + # GL rendering segfaults on headless CI runners (there is no GL context). + from pycvc_gl.lab import Lab + + scene = lab.demo_scene(Lab()) + assert scene.num_nodes() >= 3 + for name in ("terrain", "agent0_track", "agent0"): + assert scene.has(name) diff --git a/tests/test_nav.py b/tests/test_nav.py new file mode 100644 index 0000000..a822cb9 --- /dev/null +++ b/tests/test_nav.py @@ -0,0 +1,63 @@ +"""SdfNavigator + live metrics on a synthetic OPEN field (no compiled bindings needed).""" + +import numpy as np + +import sdf_nav +from grl_snam.metrics import NavMetrics, NavStats, hud_lines +from grl_snam.nav import SdfNavigator + + +def _open_field(): + """A flat, obstacle-free SDF (clearance everywhere) + matching physics meta.""" + n = 64 + phi = np.full((n, n), 5.0, np.float32) # 5 normalized units of clearance everywhere + zeros = np.zeros((n, n), np.float32) + bounds = (-100.0, -100.0, 100.0, 100.0) + scale = 0.05 # 200 world units -> 10 normalized + field = sdf_nav.SDFField(phi, zeros, zeros, bounds, (0.0, 0.0), scale) + meta = dict( + scale=scale, + center=(0.0, 0.0), + region=100.0, + rr=0.15, + d_hat=0.35, + dt=0.06, + nsub=1, + vmax=0.9, + bounds=list(bounds), + ) + return field, meta + + +def test_navigator_makes_progress_and_emits_metrics(): + field, meta = _open_field() + model = sdf_nav.CoefMLP() + model.eval() + nav = SdfNavigator(field, model, meta, reach_tol=0.8) + nav.start((-50.0, 0.0), (50.0, 0.0)) + metrics, best_i, _o, _v = nav.drive_to_goal(max_steps=2000) + assert metrics, "navigator produced no steps" + assert isinstance(metrics[best_i], NavMetrics) + assert metrics[best_i].goal_dist_m < metrics[0].goal_dist_m # closed distance to the goal + last = metrics[-1] + assert last.alpha > 0 and last.beta > 0 and last.gamma > 0 # network emitted coefficients + assert last.clearance_m > 0 # open field + + +def test_hud_lines_and_stats(): + m = NavMetrics( + step=3, + goal_dist_m=42.0, + clearance_m=8.0, + alpha=1.0, + beta=3.0, + gamma=4.0, + speed_mps=6.5, + mode="seek", + ) + stats = NavStats() + stats.update(m) + lines = hud_lines(m, stats) + assert any("coeffs" in ln for ln in lines) + assert any("reached" in ln for ln in lines) + assert stats.steps == 1