From ef6356ff82c94bda7085b39ae703abdd2b0e0f90 Mon Sep 17 00:00:00 2001 From: Joe Rivera Date: Sun, 27 Sep 2026 01:17:05 -0500 Subject: [PATCH 1/3] refactor: move the flat-layout research core under grl_snam/ MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit sdf_nav / material_nav / train_coef_energy / eval_coef_energy / surrogate_robust were top-level modules at the repo root, imported absolutely (import sdf_nav) by the package + tests + scripts + experiments and shipped via explicit pyproject py-module includes — a flat-layout legacy. Move all five under grl_snam/ and rewrite every import site (78 files) to grl_snam. (aliased imports and a docstring reference preserved; the surrogate_robust<->train_coef_energy try/except cycle survives). pyproject: drop the 5 top-level includes (now under the grl_snam package); add them to the coverage omit list to keep the 80% gate scope unchanged. Verified: all modules import, every .py compiles, and the full test suite passes (378 passed, 26 skipped — the native-lib-gated ones). --- experiments/__pycache__/e.cpython-312.pyc | Bin 33927 -> 33914 bytes experiments/__pycache__/e1.cpython-312.pyc | Bin 75349 -> 75345 bytes experiments/__pycache__/e10.cpython-312.pyc | Bin 65188 -> 65175 bytes experiments/__pycache__/e10_1.cpython-312.pyc | Bin 54374 -> 54370 bytes experiments/__pycache__/e3.cpython-312.pyc | Bin 48646 -> 48633 bytes experiments/__pycache__/e4.cpython-312.pyc | Bin 43176 -> 43163 bytes experiments/__pycache__/e5.cpython-312.pyc | Bin 75135 -> 75131 bytes experiments/__pycache__/e7.cpython-312.pyc | Bin 45964 -> 45951 bytes experiments/__pycache__/e9.cpython-312.pyc | Bin 45630 -> 45617 bytes experiments/e.py | 2 +- experiments/e0.py | 4 +-- experiments/e1.py | 4 +-- experiments/e10.py | 2 +- experiments/e10_1.py | 4 +-- experiments/e3.py | 2 +- experiments/e4.py | 2 +- experiments/e5.py | 4 +-- experiments/e5_old.py | 2 +- experiments/e7.py | 2 +- experiments/e9.py | 2 +- grl_snam/__init__.py | 2 +- grl_snam/adaptation.py | 2 +- grl_snam/demos/austin_freedrive.py | 2 +- grl_snam/demos/austin_learned.py | 2 +- grl_snam/dynamics.py | 2 +- .../eval_coef_energy.py | 4 +-- grl_snam/fog_stories.py | 2 +- grl_snam/material.py | 2 +- material_nav.py => grl_snam/material_nav.py | 2 +- grl_snam/nav.py | 2 +- grl_snam/network.py | 2 +- grl_snam/route.py | 2 +- grl_snam/scenario.py | 2 +- sdf_nav.py => grl_snam/sdf_nav.py | 0 grl_snam/squad.py | 2 +- .../surrogate_robust.py | 4 +-- grl_snam/swarm.py | 2 +- grl_snam/tools/belief_bench.py | 2 +- grl_snam/tools/capture.py | 2 +- grl_snam/tools/coef_eval.py | 2 +- grl_snam/tools/coef_export.py | 2 +- grl_snam/tools/coef_train.py | 2 +- grl_snam/tools/drive_bench.py | 2 +- grl_snam/tools/material_demo.py | 2 +- grl_snam/tools/scorecard_eval.py | 2 +- grl_snam/tools/sdf.py | 2 +- grl_snam/tools/squad_bench.py | 2 +- grl_snam/tools/swarm_bench.py | 2 +- grl_snam/tools/train.py | 2 +- .../train_coef_energy.py | 2 +- pyproject.toml | 23 +++++++++++------- scripts/train_on_geometry.py | 6 ++--- tests/test_belief.py | 2 +- tests/test_belief_occupancy_parity.py | 2 +- tests/test_bicycle_parity.py | 2 +- tests/test_clearance_route.py | 2 +- tests/test_clearance_routing_wiring.py | 2 +- tests/test_coef_eval.py | 2 +- tests/test_coef_export_checkpoint.py | 2 +- tests/test_coef_mlp_parity.py | 2 +- tests/test_cuda_drive_parity.py | 2 +- tests/test_drive_step_parity.py | 2 +- tests/test_ext_force.py | 2 +- tests/test_ipc_barrier_is_gradient.py | 2 +- tests/test_lam_head.py | 2 +- tests/test_material_fork_xcheck.py | 6 ++--- tests/test_material_grid.py | 4 +-- tests/test_material_nav.py | 2 +- tests/test_material_parity.py | 2 +- tests/test_material_rollout_parity.py | 2 +- tests/test_material_scenario.py | 2 +- tests/test_matnet_parity.py | 2 +- tests/test_nav.py | 2 +- tests/test_nav_cpp_parity.py | 2 +- tests/test_nav_ext_force.py | 2 +- tests/test_nav_native_drive.py | 2 +- tests/test_risk_lever.py | 2 +- tests/test_scenario.py | 2 +- tests/test_scorecard_eval.py | 2 +- tests/test_sdf_perf.py | 2 +- tests/test_sdf_sample_parity.py | 2 +- tests/test_sim_thread.py | 2 +- tests/test_sim_world_parity.py | 2 +- tests/test_squad.py | 2 +- tests/test_swarm.py | 2 +- tests/test_swarm_native_drive.py | 2 +- tests/test_swarm_risk_drive.py | 2 +- tests/test_vehicle.py | 2 +- tests/test_vehicle_refinements.py | 2 +- 89 files changed, 103 insertions(+), 98 deletions(-) rename eval_coef_energy.py => grl_snam/eval_coef_energy.py (99%) rename material_nav.py => grl_snam/material_nav.py (99%) rename sdf_nav.py => grl_snam/sdf_nav.py (100%) rename surrogate_robust.py => grl_snam/surrogate_robust.py (98%) rename train_coef_energy.py => grl_snam/train_coef_energy.py (99%) diff --git a/experiments/__pycache__/e.cpython-312.pyc b/experiments/__pycache__/e.cpython-312.pyc index f28d2ce49e35112647f6e37194686217a65ec3be..918a5326592aedaf6f3017bb4612d093580a261b 100644 GIT binary patch delta 53 zcmZqgWct;?#Cw{Tmx}=iW^dY&HFqQL5++80%_o>F8@Qy>i*n+N^AdCQHZwKzGcuY^ I7HTO40Mkbi3jhEB delta 66 zcmey>!PMT##Cw{Tmx}=int88gOxei0gh^aYKR2&LKP9mwQ9m&yv!pVySl=bUUw`u^ VCd&p!!OdTr_!${>CbPDb0sz1z6=eVb diff --git a/experiments/__pycache__/e1.cpython-312.pyc b/experiments/__pycache__/e1.cpython-312.pyc index f1aa0f8209c03336b439ebbb3ec2508c0c981ad5..a4866571212d701582623ff408261d330ce5b85f 100644 GIT binary patch delta 85 zcmcb5hUMZJ7T(jmyj%=GFniOEtUDWdPqQ!zHos@t{+@;Lxh$7ddQnb%ab9As-t_BA ijE!89aQ^m4Wkw4|F4MaVns*u0?=o0zpRK~UixB`acO9$% delta 89 zcmcb3hUMxR7T(jmyj%=G!2kA2#;J|Gr&+|+^>gz|^ivW`67>^PGD|8Gi}hUs{Pmk} tvTVP}!uVX4QE>VxCB{ZZf$gEnj24VsI(Hc~?=q<0WiZ@6U4?NMBLH-q9lQVl diff --git a/experiments/__pycache__/e10.cpython-312.pyc b/experiments/__pycache__/e10.cpython-312.pyc index 6404a54d334fa525a56517e9e53b13b98ccb7ab8..b407040865d338383453a6918cdef9396f83c862 100644 GIT binary patch delta 53 zcmZ4TmwEbMX5Q1hyj%=GFniOEtc@Fa-?K0ZZ5C!N{>~+pUX&AGoR^rZx4HdKAS0vc J^PGD|8Gi}hUs{Pj2g WWhws7D7d-mPaq?s&g3(HPXhoT4jG>S diff --git a/experiments/__pycache__/e10_1.cpython-312.pyc b/experiments/__pycache__/e10_1.cpython-312.pyc index e17cc6b61a0c22f9597b84498a8f6670f46292b3..8a4f1781d7a472e4b2dc70fc6317f9f02483dc45 100644 GIT binary patch delta 76 zcmaF1g89)3X5Q1hyj%=GFniOEtZ5s0C$lh$Y~Icia+XUfy(lNXI4?0*Z!*uNMlMMN Yx9`$VF4MaVns*u0?=o0To_P5x04QM{YybcN delta 64 zcmaE~g8A7BX5Q1hyj%=Gkn-kAM$1Ot$t;XolNU0JGU{($$P#jvQE>AAi;au|le;ec Ub8n7wI7mitEDC>BP+&2=oT+qtCDi*n+N^AdCQHnZ&UV`Ma) I9K3fk0QPhezW@LL delta 66 zcmezQo2l&&6YptWUM>b8xI6W7hTTTqC>C*b{oK3~{glL#ME%5+%#zB)VttnYfBnrF VEUnua1vh`+?Z?QdGudtLW&jZd7x@4H diff --git a/experiments/__pycache__/e4.cpython-312.pyc b/experiments/__pycache__/e4.cpython-312.pyc index 50aeef8ce8df058782b7a4956601ccd724510257..a2d0117a4586c8afef1c49c79c2314cbf1136030 100644 GIT binary patch delta 54 zcmZ2+k!kitCf?J$yj%=GFniOEtPdM`@3S%rZf0hCJfBM{y(lNXI4?0*Z}an|fsBl% Kn=O~;GXek?0TZSG delta 67 zcmbPzk!i(6Cf?J$yj%=GP@sP$f*X`nh=}`YDMeiTa5tnI)Bp#riG*{`#BW WvOb>AD7g97(m+N=oy|JS^BDmf#uyF& diff --git a/experiments/__pycache__/e5.cpython-312.pyc b/experiments/__pycache__/e5.cpython-312.pyc index 3761aa23db385383a5d161d759395321018c60f6..8c9691bcf39972af8a13defa41f3ec5d5eadb70c 100644 GIT binary patch delta 85 zcmex=iskny7T(jmyj%=GFniOEtQ#A7Z?G^5HveJS{)dGzPnJt6y(lNXI4?0*Z+f;8 iV*wZ`=%*x>Br@|P^2mrAS9h(3E diff --git a/experiments/__pycache__/e7.cpython-312.pyc b/experiments/__pycache__/e7.cpython-312.pyc index dab39dbf5f3004d26cec0aaaf831ef3dbf0dde98..8f765609c1a864bf1bae5eee5e35b30400b05648 100644 GIT binary patch delta 54 zcmeBq&h-Bo6YptWUM>b8n7wI7){>39acqo&o156oS8_?E7v;nk=OyOqZQiv$h>_8B JGw;TFCIAlJ688WA delta 66 zcmezWjH%~26YptWUM>b8s9Jd?WA;YgI5u&0{oK3~{glL#ME%5+%#zB)VttnYfBnrl VZ00K&1vjrM0vuR^J695E<7UKW_ diff --git a/experiments/__pycache__/e9.cpython-312.pyc b/experiments/__pycache__/e9.cpython-312.pyc index 3b334bd21d8f6a16e177afd567340a82ed10bd88..a1de830ca8ccf70b0dd92f28a9661ae40c8ce3fb 100644 GIT binary patch delta 54 zcmdn@glXdwCf?J$yj%=GFniOEtkWBL>)03tH!osyTfrrjUX&AGoR^rZxB2C|Bt}Nl K%|07GF#-S^S`;|| delta 67 zcmdn^glXRsCf?J$yj%=Gu&(S%#-WY8b!_75`nh=}`YDMeiTa5tnI)Bp#riG*{`#AH W*xXhy3U0o$E{Ty*XS4l=PmBO1^B5-p diff --git a/experiments/e.py b/experiments/e.py index 30d0831..6831c53 100644 --- a/experiments/e.py +++ b/experiments/e.py @@ -15,7 +15,7 @@ from collections import defaultdict # ==== project imports ==== -from train_coef_energy import CoefEnergyNet +from grl_snam.train_coef_energy import CoefEnergyNet import scripts.ring_dataset_maxmin as gen import scripts.spline_stagewise6 as ssi diff --git a/experiments/e0.py b/experiments/e0.py index 33fa27d..0bbaab2 100644 --- a/experiments/e0.py +++ b/experiments/e0.py @@ -26,11 +26,11 @@ import matplotlib.pyplot as plt # ==== project imports (unchanged) ==== -from train_coef_energy import CoefEnergyNet +from grl_snam.train_coef_energy import CoefEnergyNet from src.utils.online_stage_manager import StageManagerOnline import scripts.ring_dataset_maxmin as gen import scripts.spline_stagewise6 as ssi -from eval_coef_energy import HistSecantController +from grl_snam.eval_coef_energy import HistSecantController import re def mkdir(p): diff --git a/experiments/e1.py b/experiments/e1.py index c0aa45f..5131186 100644 --- a/experiments/e1.py +++ b/experiments/e1.py @@ -18,10 +18,10 @@ import matplotlib.pyplot as plt # ==== project imports ==== -from train_coef_energy import CoefEnergyNet +from grl_snam.train_coef_energy import CoefEnergyNet import scripts.ring_dataset_maxmin as gen import scripts.spline_stagewise6 as ssi -from eval_coef_energy import HistSecantController +from grl_snam.eval_coef_energy import HistSecantController import re def mkdir(p): os.makedirs(p, exist_ok=True); return p diff --git a/experiments/e10.py b/experiments/e10.py index 505a1a6..4e34898 100644 --- a/experiments/e10.py +++ b/experiments/e10.py @@ -15,7 +15,7 @@ from collections import defaultdict # ==== project imports ==== -from train_coef_energy import CoefEnergyNet +from grl_snam.train_coef_energy import CoefEnergyNet import scripts.ring_dataset_maxmin as gen import scripts.spline_stagewise6 as ssi diff --git a/experiments/e10_1.py b/experiments/e10_1.py index 761fb3b..ba8dd76 100644 --- a/experiments/e10_1.py +++ b/experiments/e10_1.py @@ -24,10 +24,10 @@ from collections import defaultdict # ==== project imports ==== -from train_coef_energy import CoefEnergyNet +from grl_snam.train_coef_energy import CoefEnergyNet import scripts.ring_dataset_maxmin as gen import scripts.spline_stagewise6 as ssi -from eval_coef_energy import HistSecantController +from grl_snam.eval_coef_energy import HistSecantController import re def mkdir(p): os.makedirs(p, exist_ok=True); return p diff --git a/experiments/e3.py b/experiments/e3.py index ed6a641..89c4376 100644 --- a/experiments/e3.py +++ b/experiments/e3.py @@ -22,7 +22,7 @@ import matplotlib.pyplot as plt # Project imports -from train_coef_energy import CoefEnergyNet +from grl_snam.train_coef_energy import CoefEnergyNet import scripts.ring_dataset_maxmin as gen import scripts.spline_stagewise6 as ssi diff --git a/experiments/e4.py b/experiments/e4.py index e599580..4152377 100644 --- a/experiments/e4.py +++ b/experiments/e4.py @@ -30,7 +30,7 @@ import matplotlib.pyplot as plt # --- Project imports -from train_coef_energy import CoefEnergyNet +from grl_snam.train_coef_energy import CoefEnergyNet import scripts.ring_dataset_maxmin as gen import scripts.spline_stagewise6 as ssi diff --git a/experiments/e5.py b/experiments/e5.py index c67f76c..251702e 100644 --- a/experiments/e5.py +++ b/experiments/e5.py @@ -22,10 +22,10 @@ import matplotlib.pyplot as plt # ==== project imports ==== -from train_coef_energy import CoefEnergyNet +from grl_snam.train_coef_energy import CoefEnergyNet import scripts.ring_dataset_maxmin as gen import scripts.spline_stagewise6 as ssi -from eval_coef_energy import HistSecantController +from grl_snam.eval_coef_energy import HistSecantController import re def mkdir(p): os.makedirs(p, exist_ok=True); return p diff --git a/experiments/e5_old.py b/experiments/e5_old.py index ad48c30..e8d1bd6 100644 --- a/experiments/e5_old.py +++ b/experiments/e5_old.py @@ -19,7 +19,7 @@ import matplotlib.pyplot as plt # ==== your project imports (unchanged) ==== -from train_coef_energy import CoefEnergyNet +from grl_snam.train_coef_energy import CoefEnergyNet import scripts.ring_dataset_maxmin as gen import scripts.spline_stagewise6 as ssi # ======================= DWA (fixed) ======================= diff --git a/experiments/e7.py b/experiments/e7.py index 29ea43e..b64096e 100644 --- a/experiments/e7.py +++ b/experiments/e7.py @@ -24,7 +24,7 @@ from tqdm import tqdm # Project imports (match your repo) -from train_coef_energy import CoefEnergyNet +from grl_snam.train_coef_energy import CoefEnergyNet import scripts.ring_dataset_maxmin as gen import scripts.spline_stagewise6 as ssi diff --git a/experiments/e9.py b/experiments/e9.py index 3fdc824..766877e 100644 --- a/experiments/e9.py +++ b/experiments/e9.py @@ -24,7 +24,7 @@ import imageio.v3 as iio # ---- Project imports (your codebase) -from train_coef_energy import CoefEnergyNet +from grl_snam.train_coef_energy import CoefEnergyNet import scripts.ring_dataset_maxmin as gen import scripts.spline_stagewise6 as ssi diff --git a/grl_snam/__init__.py b/grl_snam/__init__.py index 33a61c5..5cbbc90 100644 --- a/grl_snam/__init__.py +++ b/grl_snam/__init__.py @@ -35,7 +35,7 @@ def __getattr__(name: str): """Lazy attribute access so heavy imports (torch) happen on demand.""" if name in {"CoefEnergyNet", "integrate_surrogate"}: - from train_coef_energy import CoefEnergyNet, integrate_surrogate # noqa: PLC0415 + from grl_snam.train_coef_energy import CoefEnergyNet, integrate_surrogate # noqa: PLC0415 return {"CoefEnergyNet": CoefEnergyNet, "integrate_surrogate": integrate_surrogate}[name] raise AttributeError(f"module 'grl_snam' has no attribute {name!r}") diff --git a/grl_snam/adaptation.py b/grl_snam/adaptation.py index 4c41be3..54fecb7 100644 --- a/grl_snam/adaptation.py +++ b/grl_snam/adaptation.py @@ -2,6 +2,6 @@ from __future__ import annotations -from eval_coef_energy import HistSecantController, OnlineFinetuner # noqa: F401 +from grl_snam.eval_coef_energy import HistSecantController, OnlineFinetuner # noqa: F401 __all__ = ["HistSecantController", "OnlineFinetuner"] diff --git a/grl_snam/demos/austin_freedrive.py b/grl_snam/demos/austin_freedrive.py index 6dbe444..ac0c424 100644 --- a/grl_snam/demos/austin_freedrive.py +++ b/grl_snam/demos/austin_freedrive.py @@ -25,7 +25,7 @@ import numpy as np import torch -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from grl_snam.demos._common import ( SimPacer, current_host, diff --git a/grl_snam/demos/austin_learned.py b/grl_snam/demos/austin_learned.py index afb6f9f..de84e66 100644 --- a/grl_snam/demos/austin_learned.py +++ b/grl_snam/demos/austin_learned.py @@ -15,7 +15,7 @@ import numpy as np import torch -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from grl_snam.demos._common import ( SimPacer, current_host, diff --git a/grl_snam/dynamics.py b/grl_snam/dynamics.py index 9580049..6b92c54 100644 --- a/grl_snam/dynamics.py +++ b/grl_snam/dynamics.py @@ -7,7 +7,7 @@ from __future__ import annotations -from surrogate_robust import ( # noqa: F401 +from grl_snam.surrogate_robust import ( # noqa: F401 integrate_surrogate_v2, multi_start_penalty, ) diff --git a/eval_coef_energy.py b/grl_snam/eval_coef_energy.py similarity index 99% rename from eval_coef_energy.py rename to grl_snam/eval_coef_energy.py index b3865b7..4e17e12 100644 --- a/eval_coef_energy.py +++ b/grl_snam/eval_coef_energy.py @@ -34,10 +34,10 @@ import imageio.v3 as iio # --- project imports (adjust if paths differ) -from train_coef_energy import CoefEnergyNet # new model +from grl_snam.train_coef_energy import CoefEnergyNet # new model import scripts.ring_dataset_maxmin as gen import scripts.spline_stagewise6 as ssi -from surrogate_robust import integrate_surrogate_v2 +from grl_snam.surrogate_robust import integrate_surrogate_v2 # ---- paste near top of your eval file, after imports ---- class HistSecantController: diff --git a/grl_snam/fog_stories.py b/grl_snam/fog_stories.py index 8dd72bc..7667f87 100644 --- a/grl_snam/fog_stories.py +++ b/grl_snam/fog_stories.py @@ -401,7 +401,7 @@ def build_scenario( """ import torch # noqa: PLC0415 -- lazy: importing this module must stay cheap - import sdf_nav + import grl_snam.sdf_nav as sdf_nav from grl_snam.scenario import FogScenario if model is None: diff --git a/grl_snam/material.py b/grl_snam/material.py index 7021f63..77d7931 100644 --- a/grl_snam/material.py +++ b/grl_snam/material.py @@ -49,7 +49,7 @@ import torch import torch.nn.functional as F -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from . import nav_native as _native from .material_palette import MATERIAL_ID, OPEN_AIR_ID, RISK_MATERIAL_IDS diff --git a/material_nav.py b/grl_snam/material_nav.py similarity index 99% rename from material_nav.py rename to grl_snam/material_nav.py index 4bdff0a..3c991c7 100644 --- a/material_nav.py +++ b/grl_snam/material_nav.py @@ -48,7 +48,7 @@ import torch.nn as nn import torch.nn.functional as F -from train_coef_energy import ipc_piecewise +from grl_snam.train_coef_energy import ipc_piecewise # --------------------------------------------------------------------------- # Model diff --git a/grl_snam/nav.py b/grl_snam/nav.py index 96e4896..18ceb67 100644 --- a/grl_snam/nav.py +++ b/grl_snam/nav.py @@ -21,7 +21,7 @@ import numpy as np import torch -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from . import nav_native as _native from .metrics import NavMetrics diff --git a/grl_snam/network.py b/grl_snam/network.py index e684a18..8538781 100644 --- a/grl_snam/network.py +++ b/grl_snam/network.py @@ -7,6 +7,6 @@ from __future__ import annotations -from train_coef_energy import CoefEnergyNet # noqa: F401 +from grl_snam.train_coef_energy import CoefEnergyNet # noqa: F401 __all__ = ["CoefEnergyNet"] diff --git a/grl_snam/route.py b/grl_snam/route.py index e3c8982..0364838 100644 --- a/grl_snam/route.py +++ b/grl_snam/route.py @@ -21,7 +21,7 @@ import numpy as np -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from grl_snam import planner diff --git a/grl_snam/scenario.py b/grl_snam/scenario.py index bb74cd8..a87e8b8 100644 --- a/grl_snam/scenario.py +++ b/grl_snam/scenario.py @@ -33,7 +33,7 @@ import numpy as np -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from .belief import BeliefGrid, DynamicLayer, composite_occupancy from .nav import SdfNavigator diff --git a/sdf_nav.py b/grl_snam/sdf_nav.py similarity index 100% rename from sdf_nav.py rename to grl_snam/sdf_nav.py diff --git a/grl_snam/squad.py b/grl_snam/squad.py index 2caeb23..277b83f 100644 --- a/grl_snam/squad.py +++ b/grl_snam/squad.py @@ -26,7 +26,7 @@ import numpy as np import torch -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from . import nav_native as _native from .fog_stories import Story, build_scenario diff --git a/surrogate_robust.py b/grl_snam/surrogate_robust.py similarity index 98% rename from surrogate_robust.py rename to grl_snam/surrogate_robust.py index 4bf3aed..69d174c 100644 --- a/surrogate_robust.py +++ b/grl_snam/surrogate_robust.py @@ -5,7 +5,7 @@ Drop-in usage in train_coef_energy.py (pseudo): ------------------------------------------------ -from surrogate_robust import integrate_surrogate_v2, multi_start_penalty +from grl_snam.surrogate_robust import integrate_surrogate_v2, multi_start_penalty # in Trainer.step_batch(...): oT, vT, clr = integrate_surrogate_v2(o0, v0, goal, C, R, mask, @@ -31,7 +31,7 @@ # You can import from train_coef_energy if available; otherwise include a local copy try: - from train_coef_energy import ipc_piecewise + from grl_snam.train_coef_energy import ipc_piecewise except Exception: def ipc_piecewise(d: torch.Tensor, d_hat: torch.Tensor | float, vp: float = -5e2, eps: float = 1e-9, max_grad: float = 200.0, max_b: float = 200.0): diff --git a/grl_snam/swarm.py b/grl_snam/swarm.py index 42c6a98..0516b4f 100644 --- a/grl_snam/swarm.py +++ b/grl_snam/swarm.py @@ -41,7 +41,7 @@ import numpy as np import torch -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from . import nav_native as _native from .belief import BeliefGrid, DynamicLayer, composite_occupancy diff --git a/grl_snam/tools/belief_bench.py b/grl_snam/tools/belief_bench.py index a1d21e5..769aa94 100644 --- a/grl_snam/tools/belief_bench.py +++ b/grl_snam/tools/belief_bench.py @@ -22,7 +22,7 @@ except ImportError: # pragma: no cover raise SystemExit("belief_bench needs torch") -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from .. import nav_native, planner from ..fog_stories import STORIES, shrunk diff --git a/grl_snam/tools/capture.py b/grl_snam/tools/capture.py index c015cf1..41636e8 100644 --- a/grl_snam/tools/capture.py +++ b/grl_snam/tools/capture.py @@ -23,7 +23,7 @@ import numpy as np import torch -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from ..metrics import NavStats, hud_lines from ..nav import SdfNavigator, select_reachable_goals diff --git a/grl_snam/tools/coef_eval.py b/grl_snam/tools/coef_eval.py index 6096ea8..dbcec44 100644 --- a/grl_snam/tools/coef_eval.py +++ b/grl_snam/tools/coef_eval.py @@ -51,7 +51,7 @@ import numpy as np import torch -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from grl_snam import planner from grl_snam.fog_stories import STORIES, shrunk from grl_snam.material import FrictionField diff --git a/grl_snam/tools/coef_export.py b/grl_snam/tools/coef_export.py index ab47091..fc85aaf 100644 --- a/grl_snam/tools/coef_export.py +++ b/grl_snam/tools/coef_export.py @@ -111,7 +111,7 @@ def write_coef_mlp(model, path, meta: bytes = b""): def main(argv=None): import torch - import sdf_nav + import grl_snam.sdf_nav as sdf_nav argv = list(sys.argv[1:] if argv is None else argv) if len(argv) != 2: diff --git a/grl_snam/tools/coef_train.py b/grl_snam/tools/coef_train.py index 9940464..d09d88d 100644 --- a/grl_snam/tools/coef_train.py +++ b/grl_snam/tools/coef_train.py @@ -28,7 +28,7 @@ except ImportError: # pragma: no cover raise SystemExit("coef_train needs torch") -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from .. import planner from ..fog_stories import STORIES, shrunk diff --git a/grl_snam/tools/drive_bench.py b/grl_snam/tools/drive_bench.py index 80562d9..ebaea8f 100644 --- a/grl_snam/tools/drive_bench.py +++ b/grl_snam/tools/drive_bench.py @@ -24,7 +24,7 @@ def _run_once(n: int, ticks: int): import numpy as np import torch - import sdf_nav + import grl_snam.sdf_nav as sdf_nav from grl_snam import planner from grl_snam.fog_stories import STORIES, shrunk from grl_snam.squad import AgentSpec diff --git a/grl_snam/tools/material_demo.py b/grl_snam/tools/material_demo.py index 2b0ee74..31e9086 100644 --- a/grl_snam/tools/material_demo.py +++ b/grl_snam/tools/material_demo.py @@ -46,7 +46,7 @@ def _meta(): def _model(seed=0): import torch - import sdf_nav + import grl_snam.sdf_nav as sdf_nav torch.manual_seed(seed) m = sdf_nav.CoefMLP() diff --git a/grl_snam/tools/scorecard_eval.py b/grl_snam/tools/scorecard_eval.py index c4e099e..cdc9325 100644 --- a/grl_snam/tools/scorecard_eval.py +++ b/grl_snam/tools/scorecard_eval.py @@ -23,7 +23,7 @@ import numpy as np import torch -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from grl_snam import planner from grl_snam.fog_stories import STORIES, shrunk from grl_snam.material_palette import terrain_risk_share diff --git a/grl_snam/tools/sdf.py b/grl_snam/tools/sdf.py index a9e49bc..573404a 100644 --- a/grl_snam/tools/sdf.py +++ b/grl_snam/tools/sdf.py @@ -15,7 +15,7 @@ import numpy as np -import sdf_nav +import grl_snam.sdf_nav as sdf_nav TARGET_EXTENT = 10.0 # the surrogate's normalized working scale (~10 units across the region) diff --git a/grl_snam/tools/squad_bench.py b/grl_snam/tools/squad_bench.py index 61ae5e4..de68da4 100644 --- a/grl_snam/tools/squad_bench.py +++ b/grl_snam/tools/squad_bench.py @@ -26,7 +26,7 @@ def _run_once(grid: int, n: int, ticks: int, seed: int, route_clearance=None): import numpy as np import torch - import sdf_nav + import grl_snam.sdf_nav as sdf_nav from grl_snam import nav_native, planner from grl_snam.fog_stories import STORIES, shrunk from grl_snam.squad import AgentSpec, Squad diff --git a/grl_snam/tools/swarm_bench.py b/grl_snam/tools/swarm_bench.py index 386603d..df11b81 100644 --- a/grl_snam/tools/swarm_bench.py +++ b/grl_snam/tools/swarm_bench.py @@ -32,7 +32,7 @@ except ImportError: # pragma: no cover - torch is a hard dep of the swarm raise SystemExit("swarm_bench needs torch") -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from .. import planner from ..fog_stories import STORIES, shrunk diff --git a/grl_snam/tools/train.py b/grl_snam/tools/train.py index d524183..c0319ad 100644 --- a/grl_snam/tools/train.py +++ b/grl_snam/tools/train.py @@ -18,7 +18,7 @@ import torch import torch.nn.functional as F -import sdf_nav +import grl_snam.sdf_nav as sdf_nav def train_sdf( diff --git a/train_coef_energy.py b/grl_snam/train_coef_energy.py similarity index 99% rename from train_coef_energy.py rename to grl_snam/train_coef_energy.py index 2f5b5e8..4c34747 100644 --- a/train_coef_energy.py +++ b/grl_snam/train_coef_energy.py @@ -57,7 +57,7 @@ import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader -from surrogate_robust import integrate_surrogate_v2, multi_start_penalty +from grl_snam.surrogate_robust import integrate_surrogate_v2, multi_start_penalty import numpy as np diff --git a/pyproject.toml b/pyproject.toml index d5fa234..e7246d9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,13 +43,9 @@ packages = [ { 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" }, - { include = "material_nav.py" }, + # The research core (sdf_nav / material_nav / train_coef_energy / eval_coef_energy / + # surrogate_robust) now lives UNDER grl_snam/ and is imported as grl_snam., so it + # ships with the package — no more top-level module includes (they were a flat-layout legacy). ] [tool.poetry.group.dev.dependencies] @@ -98,8 +94,9 @@ ignore = [ # (the maintained `grl_snam` package only). Omitted below is code with no # headless-CI surface — runnable GL/matplotlib demos, video-capture / # benchmark / one-off visualization scripts — plus trivial stubs and the thin -# CLI entry point. The flat-layout research modules (sdf_nav.py, material_nav.py, -# ...) sit outside `grl_snam` and, like lint, outside this gate. `pytest --cov` +# CLI entry point. The research core (grl_snam/sdf_nav.py, material_nav.py, +# ...) is now under `grl_snam` but explicitly OMITted below to keep it outside this +# gate, as it was pre-refactor when it lived at the top level. `pytest --cov` # and CI's `coverage run` both read this config, so the local and gated numbers # match. Current: ~85% of the in-scope surface (gate floor is 80). # ----------------------------------------------------------------------------- @@ -124,6 +121,14 @@ omit = [ "grl_snam/adaptation.py", # stub "grl_snam/network.py", # stub "grl_snam/utils/*", # stub + # The research core, moved under grl_snam/ (was top-level, outside this gate). Kept OUT of the + # coverage gate to preserve the pre-refactor scope — these are large torch research modules + # exercised by the parity/native tests, not unit-covered line-by-line. + "grl_snam/sdf_nav.py", + "grl_snam/material_nav.py", + "grl_snam/train_coef_energy.py", + "grl_snam/eval_coef_energy.py", + "grl_snam/surrogate_robust.py", ] [tool.coverage.report] diff --git a/scripts/train_on_geometry.py b/scripts/train_on_geometry.py index 6493836..7c77810 100644 --- a/scripts/train_on_geometry.py +++ b/scripts/train_on_geometry.py @@ -41,9 +41,9 @@ for _m in ["imageio", "imageio.v3", "scripts.ring_dataset_maxmin", "scripts.spline_stagewise6"]: sys.modules.setdefault(_m, types.ModuleType(_m)) -from eval_coef_energy import build_local_feats # noqa: E402 -from surrogate_robust import integrate_surrogate_v2 # noqa: E402 -from train_coef_energy import CoefEnergyNet # noqa: E402 +from grl_snam.eval_coef_energy import build_local_feats # noqa: E402 +from grl_snam.surrogate_robust import integrate_surrogate_v2 # noqa: E402 +from grl_snam.train_coef_energy import CoefEnergyNet # noqa: E402 TARGET_EXTENT = 10.0 # normalize a working region to ~10 units (the tuned regime) diff --git a/tests/test_belief.py b/tests/test_belief.py index f2a5d9c..3e586cc 100644 --- a/tests/test_belief.py +++ b/tests/test_belief.py @@ -167,7 +167,7 @@ def test_composite_merges_belief_and_dynamic(): def test_belief_occupancy_feeds_build_sdf(): torch = __import__("pytest").importorskip("torch") # noqa: F841 - import sdf_nav + import grl_snam.sdf_nav as sdf_nav b = BeliefGrid((N, N), BOUNDS) truth = _wall_truth(col=40) diff --git a/tests/test_belief_occupancy_parity.py b/tests/test_belief_occupancy_parity.py index 1e6fb34..356ee51 100644 --- a/tests/test_belief_occupancy_parity.py +++ b/tests/test_belief_occupancy_parity.py @@ -64,7 +64,7 @@ def test_composite_occupancy_bit_identical(unknown): def test_composite_feeds_bit_identical_field(): """The whole point: the C++ occupancy must produce the SAME SDF field as the numpy occupancy (the fidelity boundary is the built field).""" - import sdf_nav + import grl_snam.sdf_nav as sdf_nav b, dyn = _evolved_belief() story = shrunk(STORIES["city"], n=192, max_steps=100) diff --git a/tests/test_bicycle_parity.py b/tests/test_bicycle_parity.py index c755a87..1b72a54 100644 --- a/tests/test_bicycle_parity.py +++ b/tests/test_bicycle_parity.py @@ -13,7 +13,7 @@ torch = pytest.importorskip("torch") pytest.importorskip("pycvc") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam import nav_native # noqa: E402 from grl_snam.fog_stories import STORIES, shrunk # noqa: E402 diff --git a/tests/test_clearance_route.py b/tests/test_clearance_route.py index 7692ab5..247f649 100644 --- a/tests/test_clearance_route.py +++ b/tests/test_clearance_route.py @@ -5,7 +5,7 @@ import numpy as np -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from grl_snam import planner diff --git a/tests/test_clearance_routing_wiring.py b/tests/test_clearance_routing_wiring.py index 7be005e..56700fb 100644 --- a/tests/test_clearance_routing_wiring.py +++ b/tests/test_clearance_routing_wiring.py @@ -2,7 +2,7 @@ import numpy as np -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from grl_snam import planner from grl_snam.route import cells_for_metres, plan_clearance_route diff --git a/tests/test_coef_eval.py b/tests/test_coef_eval.py index 0139b52..9b734a1 100644 --- a/tests/test_coef_eval.py +++ b/tests/test_coef_eval.py @@ -3,7 +3,7 @@ import numpy as np import torch -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from grl_snam.tools import coef_eval diff --git a/tests/test_coef_export_checkpoint.py b/tests/test_coef_export_checkpoint.py index b418238..59494c2 100644 --- a/tests/test_coef_export_checkpoint.py +++ b/tests/test_coef_export_checkpoint.py @@ -11,7 +11,7 @@ import torch -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from grl_snam.tools import coef_export diff --git a/tests/test_coef_mlp_parity.py b/tests/test_coef_mlp_parity.py index fb031d0..ac668a2 100644 --- a/tests/test_coef_mlp_parity.py +++ b/tests/test_coef_mlp_parity.py @@ -15,7 +15,7 @@ torch = pytest.importorskip("torch") pytest.importorskip("pycvc") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam import nav_native # noqa: E402 from grl_snam.tools import coef_export # noqa: E402 diff --git a/tests/test_cuda_drive_parity.py b/tests/test_cuda_drive_parity.py index f54dfba..65eb6b9 100644 --- a/tests/test_cuda_drive_parity.py +++ b/tests/test_cuda_drive_parity.py @@ -14,7 +14,7 @@ torch = pytest.importorskip("torch") pytest.importorskip("pycvc") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam import nav_native # noqa: E402 from grl_snam.fog_stories import STORIES, shrunk # noqa: E402 from grl_snam.tools import coef_export # noqa: E402 diff --git a/tests/test_drive_step_parity.py b/tests/test_drive_step_parity.py index e0e12de..9def95c 100644 --- a/tests/test_drive_step_parity.py +++ b/tests/test_drive_step_parity.py @@ -13,7 +13,7 @@ torch = pytest.importorskip("torch") pytest.importorskip("pycvc") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam import nav_native # noqa: E402 from grl_snam.fog_stories import STORIES, shrunk # noqa: E402 from grl_snam.tools import coef_export # noqa: E402 diff --git a/tests/test_ext_force.py b/tests/test_ext_force.py index f08612b..9a5fcf7 100644 --- a/tests/test_ext_force.py +++ b/tests/test_ext_force.py @@ -12,7 +12,7 @@ torch = pytest.importorskip("torch") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 BOUNDS = (-100.0, -100.0, 100.0, 100.0) CENTER = (0.0, 0.0) diff --git a/tests/test_ipc_barrier_is_gradient.py b/tests/test_ipc_barrier_is_gradient.py index d103358..e381f3d 100644 --- a/tests/test_ipc_barrier_is_gradient.py +++ b/tests/test_ipc_barrier_is_gradient.py @@ -7,7 +7,7 @@ import torch -from train_coef_energy import ipc_piecewise +from grl_snam.train_coef_energy import ipc_piecewise def test_dbdd_is_the_autograd_derivative_of_b(): diff --git a/tests/test_lam_head.py b/tests/test_lam_head.py index 13a3b3a..1fa7db2 100644 --- a/tests/test_lam_head.py +++ b/tests/test_lam_head.py @@ -13,7 +13,7 @@ import pytest import torch -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from grl_snam.fog_stories import STORIES, shrunk from grl_snam.material import city_material_grid from grl_snam.tools import coef_train diff --git a/tests/test_material_fork_xcheck.py b/tests/test_material_fork_xcheck.py index fee34e7..4567c3c 100644 --- a/tests/test_material_fork_xcheck.py +++ b/tests/test_material_fork_xcheck.py @@ -68,7 +68,7 @@ def fork_modules(): def test_barrier_sampler_integrator_bitwise(fork_modules): - import material_nav as mnav + import grl_snam.material_nav as mnav fork_tm, _ = fork_modules tg = torch.Generator().manual_seed(7) @@ -114,7 +114,7 @@ def test_barrier_sampler_integrator_bitwise(fork_modules): def test_model_state_dict_and_forward_bitwise(fork_modules): - import material_nav as mnav + import grl_snam.material_nav as mnav fork_tm, _ = fork_modules m_fork = fork_tm.CoefEnergyNetMaterial() @@ -135,7 +135,7 @@ def test_model_state_dict_and_forward_bitwise(fork_modules): def test_gate_bitwise_random_grids(fork_modules): - import material_nav as mnav + import grl_snam.material_nav as mnav _, fork_exp1 = fork_modules rng = np.random.default_rng(3) diff --git a/tests/test_material_grid.py b/tests/test_material_grid.py index 1ac7e8c..0af2be5 100644 --- a/tests/test_material_grid.py +++ b/tests/test_material_grid.py @@ -16,7 +16,7 @@ torch = pytest.importorskip("torch") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam.material import ( # noqa: E402 MaterialGrid, MaterialParams, @@ -198,7 +198,7 @@ def test_barrier_scale_oracle_phi_in_metres(): """The k/S rescale trap: phi stays in metres so at phi == d_hat_sdf_m the barrier factor is EXACTLY -sigmoid(0) = -0.5, and one metre to either side is sigmoid(+-k_sharp).""" - from sdf_nav import _material_force + from grl_snam.sdf_nav import _material_force class _Flat: def __init__(self, phi_m): diff --git a/tests/test_material_nav.py b/tests/test_material_nav.py index d63295e..0e6f34f 100644 --- a/tests/test_material_nav.py +++ b/tests/test_material_nav.py @@ -19,7 +19,7 @@ import pytest import torch -import material_nav as mnav +import grl_snam.material_nav as mnav # --------------------------------------------------------------------------- # Witness gate diff --git a/tests/test_material_parity.py b/tests/test_material_parity.py index 06048dd..ba0c20d 100644 --- a/tests/test_material_parity.py +++ b/tests/test_material_parity.py @@ -31,7 +31,7 @@ if not nav_native.HAS_MATERIAL: pytest.skip("this pycvc build has no material kernels", allow_module_level=True) -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 BOUNDS = (-100.0, -100.0, 100.0, 100.0) CENTER = (0.0, 0.0) diff --git a/tests/test_material_rollout_parity.py b/tests/test_material_rollout_parity.py index 45138d0..1044228 100644 --- a/tests/test_material_rollout_parity.py +++ b/tests/test_material_rollout_parity.py @@ -20,7 +20,7 @@ torch = pytest.importorskip("torch") pytest.importorskip("pycvc") -import material_nav # noqa: E402 +import grl_snam.material_nav as material_nav # noqa: E402 from grl_snam import nav_native # noqa: E402 pytestmark = pytest.mark.skipif( diff --git a/tests/test_material_scenario.py b/tests/test_material_scenario.py index b5869ab..d9e19d8 100644 --- a/tests/test_material_scenario.py +++ b/tests/test_material_scenario.py @@ -12,7 +12,7 @@ torch = pytest.importorskip("torch") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam.material import MaterialGrid, MaterialParams # noqa: E402 from grl_snam.scenario import FogScenario # noqa: E402 diff --git a/tests/test_matnet_parity.py b/tests/test_matnet_parity.py index 7365e6c..5b7d004 100644 --- a/tests/test_matnet_parity.py +++ b/tests/test_matnet_parity.py @@ -24,7 +24,7 @@ torch = pytest.importorskip("torch") -import material_nav # noqa: E402 +import grl_snam.material_nav as material_nav # noqa: E402 from grl_snam import nav_native # noqa: E402 from grl_snam.tools.matnet_export import ( # noqa: E402 matnet_forward_numpy, diff --git a/tests/test_nav.py b/tests/test_nav.py index 5d2f575..aa12357 100644 --- a/tests/test_nav.py +++ b/tests/test_nav.py @@ -2,7 +2,7 @@ import numpy as np -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from grl_snam.metrics import NavMetrics, NavStats, hud_lines from grl_snam.nav import SdfNavigator diff --git a/tests/test_nav_cpp_parity.py b/tests/test_nav_cpp_parity.py index 6a8a338..c94d658 100644 --- a/tests/test_nav_cpp_parity.py +++ b/tests/test_nav_cpp_parity.py @@ -16,7 +16,7 @@ if not hasattr(pycvc, "nav_astar"): pytest.skip("pycvc build has no cvc::nav kernels", allow_module_level=True) -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam import nav_native, planner # noqa: E402 _SIZES = [3, 5, 8, 13, 16, 24, 33, 48, 64, 96, 128] diff --git a/tests/test_nav_ext_force.py b/tests/test_nav_ext_force.py index 9daf551..45b6ceb 100644 --- a/tests/test_nav_ext_force.py +++ b/tests/test_nav_ext_force.py @@ -10,7 +10,7 @@ import numpy as np import torch -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from grl_snam.nav import SdfNavigator diff --git a/tests/test_nav_native_drive.py b/tests/test_nav_native_drive.py index f46a087..85673a2 100644 --- a/tests/test_nav_native_drive.py +++ b/tests/test_nav_native_drive.py @@ -17,7 +17,7 @@ torch = pytest.importorskip("torch") pytest.importorskip("pycvc") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam import nav_native # noqa: E402 from grl_snam.fog_stories import STORIES, shrunk # noqa: E402 from grl_snam.nav import SdfNavigator # noqa: E402 diff --git a/tests/test_risk_lever.py b/tests/test_risk_lever.py index c58de8b..588568a 100644 --- a/tests/test_risk_lever.py +++ b/tests/test_risk_lever.py @@ -13,7 +13,7 @@ import pytest import torch -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from grl_snam.fog_stories import STORIES, shrunk from grl_snam.material import city_material_grid from grl_snam.tools import coef_train diff --git a/tests/test_scenario.py b/tests/test_scenario.py index 4d63c66..65fa670 100644 --- a/tests/test_scenario.py +++ b/tests/test_scenario.py @@ -10,7 +10,7 @@ torch = pytest.importorskip("torch") -import sdf_nav # noqa: E402 (imports torch; must follow the skip guard) +import grl_snam.sdf_nav as sdf_nav # noqa: E402 (imports torch; must follow the skip guard) from grl_snam.scenario import Event, FogScenario # noqa: E402 N = 96 diff --git a/tests/test_scorecard_eval.py b/tests/test_scorecard_eval.py index e4c2435..34bdc1e 100644 --- a/tests/test_scorecard_eval.py +++ b/tests/test_scorecard_eval.py @@ -7,7 +7,7 @@ torch = pytest.importorskip("torch") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam.scorecard import NavScorecard # noqa: E402 from grl_snam.tools import scorecard_eval # noqa: E402 diff --git a/tests/test_sdf_perf.py b/tests/test_sdf_perf.py index 9f97f5f..d232565 100644 --- a/tests/test_sdf_perf.py +++ b/tests/test_sdf_perf.py @@ -10,7 +10,7 @@ torch = pytest.importorskip("torch") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 def _reference_edt2(mask): diff --git a/tests/test_sdf_sample_parity.py b/tests/test_sdf_sample_parity.py index 181a804..fa68168 100644 --- a/tests/test_sdf_sample_parity.py +++ b/tests/test_sdf_sample_parity.py @@ -18,7 +18,7 @@ torch = pytest.importorskip("torch") pytest.importorskip("pycvc") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam import nav_native # noqa: E402 pytestmark = pytest.mark.skipif( diff --git a/tests/test_sim_thread.py b/tests/test_sim_thread.py index b23f212..865bf85 100644 --- a/tests/test_sim_thread.py +++ b/tests/test_sim_thread.py @@ -15,7 +15,7 @@ torch = pytest.importorskip("torch") pytest.importorskip("pycvc") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam import nav_native, planner # noqa: E402 from grl_snam.fog_stories import STORIES, shrunk # noqa: E402 from grl_snam.squad import AgentSpec # noqa: E402 diff --git a/tests/test_sim_world_parity.py b/tests/test_sim_world_parity.py index 86c73d2..01f208a 100644 --- a/tests/test_sim_world_parity.py +++ b/tests/test_sim_world_parity.py @@ -13,7 +13,7 @@ torch = pytest.importorskip("torch") pytest.importorskip("pycvc") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam import nav_native, planner # noqa: E402 from grl_snam.fog_stories import STORIES, shrunk # noqa: E402 from grl_snam.squad import AgentSpec # noqa: E402 diff --git a/tests/test_squad.py b/tests/test_squad.py index 33d7e52..94299bb 100644 --- a/tests/test_squad.py +++ b/tests/test_squad.py @@ -237,7 +237,7 @@ def test_stagger_assigns_distinct_phases(): # ── stage-2: batched vehicle rollout (bit-identical to serial) ─────────────── -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 def _shared_model(): diff --git a/tests/test_swarm.py b/tests/test_swarm.py index 5a88a75..a3dcb10 100644 --- a/tests/test_swarm.py +++ b/tests/test_swarm.py @@ -20,7 +20,7 @@ torch = pytest.importorskip("torch") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam import planner # noqa: E402 from grl_snam.fog_stories import STORIES, shrunk # noqa: E402 from grl_snam.metrics import NavStats # noqa: E402 diff --git a/tests/test_swarm_native_drive.py b/tests/test_swarm_native_drive.py index 124f09a..a0c55bf 100644 --- a/tests/test_swarm_native_drive.py +++ b/tests/test_swarm_native_drive.py @@ -20,7 +20,7 @@ torch = pytest.importorskip("torch") pytest.importorskip("pycvc") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam import nav_native, planner # noqa: E402 from grl_snam.fog_stories import STORIES, shrunk # noqa: E402 from grl_snam.squad import AgentSpec # noqa: E402 diff --git a/tests/test_swarm_risk_drive.py b/tests/test_swarm_risk_drive.py index 1f1742e..e3eb5af 100644 --- a/tests/test_swarm_risk_drive.py +++ b/tests/test_swarm_risk_drive.py @@ -9,7 +9,7 @@ import numpy as np import pytest -import sdf_nav +import grl_snam.sdf_nav as sdf_nav from grl_snam.fog_stories import STORIES, shrunk from grl_snam.material import city_material_grid from grl_snam.scorecard import NavScorecard diff --git a/tests/test_vehicle.py b/tests/test_vehicle.py index 3b19084..a57c03f 100644 --- a/tests/test_vehicle.py +++ b/tests/test_vehicle.py @@ -17,7 +17,7 @@ # These import torch at module level, so they must come AFTER the skip guard # or collection itself crashes on a torch-less environment. -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam.nav import SdfNavigator # noqa: E402 RR, DHAT, DT, VMAX = 0.15, 0.35, 0.06, 0.9 diff --git a/tests/test_vehicle_refinements.py b/tests/test_vehicle_refinements.py index 960129c..ad4bced 100644 --- a/tests/test_vehicle_refinements.py +++ b/tests/test_vehicle_refinements.py @@ -17,7 +17,7 @@ torch = pytest.importorskip("torch") -import sdf_nav # noqa: E402 +import grl_snam.sdf_nav as sdf_nav # noqa: E402 from grl_snam.material import FrictionField # noqa: E402 RR, DHAT, DT, VMAX = 0.15, 0.35, 0.06, 0.9 From 0012a214405e4828917b570a57a8c9cc0b9daca3 Mon Sep 17 00:00:00 2001 From: Joe Rivera Date: Sun, 27 Sep 2026 01:32:40 -0500 Subject: [PATCH 2/3] refactor: hold the moved research modules to the package lint standard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Moving the five modules under grl_snam/ pulled them into the black+ruff gate (they were unlinted at the repo root). Rather than carve them back out, bring them up to standard: black-format all five, ruff --fix (44 auto-fixes: import sort, pyupgrade, unused-import cleanup — noqa re-exports preserved), and drop two dead F841 unused locals (surrogate_robust dev, train_coef_energy W) by hand. Lint now clean across the whole grl_snam package with the five IN scope; full suite still 378 passed / 26 skipped. --- grl_snam/eval_coef_energy.py | 409 +++++++++++++++++++++++++--------- grl_snam/material_nav.py | 18 +- grl_snam/surrogate_robust.py | 141 ++++++++---- grl_snam/train_coef_energy.py | 292 ++++++++++++++++-------- pyproject.toml | 7 +- 5 files changed, 614 insertions(+), 253 deletions(-) diff --git a/grl_snam/eval_coef_energy.py b/grl_snam/eval_coef_energy.py index 4e17e12..4bb3dd9 100644 --- a/grl_snam/eval_coef_energy.py +++ b/grl_snam/eval_coef_energy.py @@ -26,27 +26,44 @@ • Your stagewise module is available as scripts.spline_stagewise6 (or adapt imports below). """ from __future__ import annotations -import os, glob, math, argparse + +import argparse +import glob +import math +import os from datetime import datetime + +import imageio.v3 as iio import numpy as np import torch -import torch.nn.functional as F -import imageio.v3 as iio -# --- project imports (adjust if paths differ) -from grl_snam.train_coef_energy import CoefEnergyNet # new model import scripts.ring_dataset_maxmin as gen import scripts.spline_stagewise6 as ssi from grl_snam.surrogate_robust import integrate_surrogate_v2 +# --- project imports (adjust if paths differ) +from grl_snam.train_coef_energy import CoefEnergyNet # new model + + # ---- paste near top of your eval file, after imports ---- class HistSecantController: """ One-step history-based (secant) sensitivity controller. Adjusts: beta, gamma, and top-K alphas (nearest obstacles) without extra sims. """ - def __init__(self, k_alpha=2, lr_beta=0.15, lr_gamma=0.10, lr_alpha=0.4, - safe_margin=0.08, v_min=0.25, v_max=0.5, prog_eps=0.01, ema=0.9): + + def __init__( + self, + k_alpha=2, + lr_beta=0.15, + lr_gamma=0.10, + lr_alpha=0.4, + safe_margin=0.08, + v_min=0.25, + v_max=0.5, + prog_eps=0.01, + ema=0.9, + ): self.k_alpha = k_alpha self.lr_b, self.lr_g, self.lr_a = lr_beta, lr_gamma, lr_alpha self.safe_margin = safe_margin @@ -57,32 +74,57 @@ def __init__(self, k_alpha=2, lr_beta=0.15, lr_gamma=0.10, lr_alpha=0.4, # history self.prev = None # dict with {theta, y} # running Jacobian (optional EMA) - self.J = None # shape (3, D) where D = 2 + k_alpha + self.J = None # shape (3, D) where D = 2 + k_alpha def _select_alpha_indices(self, o_w, Cw, Rw, Ww): - if Cw.shape[0] == 0: return np.array([], dtype=int) - d = np.linalg.norm(o_w[None,:] - Cw, axis=1) - Rw - j = np.argsort(d)[:self.k_alpha] + if Cw.shape[0] == 0: + return np.array([], dtype=int) + d = np.linalg.norm(o_w[None, :] - Cw, axis=1) - Rw + j = np.argsort(d)[: self.k_alpha] return j - def update(self, alphas, beta, gamma, # current params (torch, shapes [1,N], [1], [1]) - o_w, v_w, goal, Cw, Rw, Ww, # numpy (world slice you already have) - clearance_now, dist_now, speed_now): + def update( + self, + alphas, + beta, + gamma, # current params (torch, shapes [1,N], [1], [1]) + o_w, + v_w, + goal, + Cw, + Rw, + Ww, # numpy (world slice you already have) + clearance_now, + dist_now, + speed_now, + ): # choose alpha subset idx = self._select_alpha_indices(o_w, Cw, Rw, Ww) - a_sub = (alphas.squeeze(0).detach().cpu().numpy()[idx] if idx.size else np.zeros((0,), dtype=np.float32)) + a_sub = ( + alphas.squeeze(0).detach().cpu().numpy()[idx] + if idx.size + else np.zeros((0,), dtype=np.float32) + ) b = float(beta.squeeze(0).detach().cpu().item()) g = float(gamma.squeeze(0).detach().cpu().item()) # current observable vector (we use signs so 'lower is worse') - y = np.array([-float(clearance_now), float(dist_now),# float(speed_now), - -float(speed_now)], dtype=np.float32) + y = np.array( + [-float(clearance_now), float(dist_now), -float(speed_now)], # float(speed_now), + dtype=np.float32, + ) theta = np.concatenate([np.array([b, g], dtype=np.float32), a_sub], axis=0) # D = 2+k_alpha # targets: keep clearance >= safe_margin, decrease dist, keep speed >= v_min if safe - y_tgt = np.array([-self.safe_margin, dist_now - self.prog_eps, - # min(speed_now, self.v_max), - -max(speed_now, self.v_min if clearance_now >= self.safe_margin else 0.0)], dtype=np.float32) + y_tgt = np.array( + [ + -self.safe_margin, + dist_now - self.prog_eps, + # min(speed_now, self.v_max), + -max(speed_now, self.v_min if clearance_now >= self.safe_margin else 0.0), + ], + dtype=np.float32, + ) # if first frame: store and return if self.prev is None: @@ -105,7 +147,7 @@ def update(self, alphas, beta, gamma, # current params (torch, shapes [1,N], [1 self.J = (self.J * self.ema + J_new * (1.0 - self.ema)) if self.J is not None else J_new # desired Δy - dy_des = (y_tgt - y) # move current y toward target y + dy_des = y_tgt - y # move current y toward target y # small LS step on parameters: minimize ||J Δθ - dy_des||_2 # Δθ* = argmin ||J Δθ - dy_des|| -> normal eq: (J^T J) Δθ = J^T dy_des @@ -135,29 +177,39 @@ def update(self, alphas, beta, gamma, # current params (torch, shapes [1,N], [1 a_full = alphas.squeeze(0).detach().cpu().numpy() # back to torch on the right device/dtype - dev = alphas.device; dtype = alphas.dtype + dev = alphas.device + dtype = alphas.dtype al_out = torch.as_tensor(a_full, device=dev, dtype=dtype).unsqueeze(0) - b_out = torch.as_tensor([b_new], device=dev, dtype=beta.dtype) - g_out = torch.as_tensor([g_new], device=dev, dtype=gamma.dtype) + b_out = torch.as_tensor([b_new], device=dev, dtype=beta.dtype) + g_out = torch.as_tensor([g_new], device=dev, dtype=gamma.dtype) return al_out, b_out, g_out ##### Test-Time Finetuner (TTT) + class OnlineFinetuner: - def __init__(self, model, lr=1e-4, max_steps=2, weight_decay=0.0, prox_lambda=1e-3, head_name_filters=("head","out","proj")): + def __init__( + self, + model, + lr=1e-4, + max_steps=2, + weight_decay=0.0, + prox_lambda=1e-3, + head_name_filters=("head", "out", "proj"), + ): self.model = model self.max_steps = max_steps self.prox_lambda = torch.tensor(prox_lambda, dtype=torch.float64) # select a small, stable subset of parameters (final/near-final layers) trainable = [] - for n,p in model.named_parameters(): + for n, p in model.named_parameters(): if any(tag in n for tag in head_name_filters): p.requires_grad_(True) - trainable.append((n,p)) + trainable.append((n, p)) else: p.requires_grad_(False) - self.params = [p for _,p in trainable] + self.params = [p for _, p in trainable] self.opt = torch.optim.Adam(self.params, lr=lr, weight_decay=weight_decay) # store anchors for proximal regularization self._anchors = {id(p): p.detach().clone() for p in self.params} @@ -171,7 +223,11 @@ def step(self, obs_feats, obs_mask, goal_feats, targets): for _ in range(self.max_steps): self.opt.zero_grad() a_pred, b_pred, g_pred = self.model(obs_feats, obs_mask, goal_feats) - L_a = torch.nn.functional.mse_loss(a_pred, targets["alphas"]) if a_pred.numel() else a_pred.sum()*0 + L_a = ( + torch.nn.functional.mse_loss(a_pred, targets["alphas"]) + if a_pred.numel() + else a_pred.sum() * 0 + ) L_b = torch.nn.functional.mse_loss(b_pred, targets["beta"]) L_g = torch.nn.functional.mse_loss(g_pred, targets["gamma"]) # proximal anchor: keep heads near checkpoint @@ -182,19 +238,41 @@ def step(self, obs_feats, obs_mask, goal_feats, targets): L.backward() torch.nn.utils.clip_grad_norm_(self.params, 1.0) self.opt.step() - losses = {"L": float(L.item()), "La": float(L_a.item() if a_pred.numel() else 0.0), - "Lb": float(L_b.item()), "Lg": float(L_g.item()), "Lprox": float((self.prox_lambda*L_prox).item())} + losses = { + "L": float(L.item()), + "La": float(L_a.item() if a_pred.numel() else 0.0), + "Lb": float(L_b.item()), + "Lg": float(L_g.item()), + "Lprox": float((self.prox_lambda * L_prox).item()), + } self.model.eval() return losses + # ----------------------------- # Helpers for local features (match training features) # ----------------------------- -def ttt_rollout_loss(model, obs_feats_t, obs_mask_t, goal_feats_t, - o_t, v_t, goal_t, C_t, R_t, mask_t, d_hat_t, dt_t, - o_tp1_obs, v_tp1_obs, - integrate_surrogate, robot_radius=None, margin_factor=0.5): + +def ttt_rollout_loss( + model, + obs_feats_t, + obs_mask_t, + goal_feats_t, + o_t, + v_t, + goal_t, + C_t, + R_t, + mask_t, + d_hat_t, + dt_t, + o_tp1_obs, + v_tp1_obs, + integrate_surrogate, + robot_radius=None, + margin_factor=0.5, +): """ Builds loss: ||ô - o_{t+1}||^2 + 0.25||v̂ - v_{t+1}||^2 using a single surrogate step. All tensors are shaped for B=1 (batch size 1). @@ -202,23 +280,63 @@ def ttt_rollout_loss(model, obs_feats_t, obs_mask_t, goal_feats_t, model.train() a, b, g = model(obs_feats_t, obs_mask_t, goal_feats_t) o_hat, v_hat, _ = integrate_surrogate( - o_t, v_t, goal_t, C_t, R_t, mask_t, a, b, g, d_hat_t, dt_t, H=torch.tensor([1], device=o_t.device), + o_t, + v_t, + goal_t, + C_t, + R_t, + mask_t, + a, + b, + g, + d_hat_t, + dt_t, + H=torch.tensor([1], device=o_t.device), robot_radius=(robot_radius if robot_radius is not None else 0.0), - margin_factor=margin_factor + margin_factor=margin_factor, + ) + L = torch.nn.functional.mse_loss(o_hat, o_tp1_obs) + 0.25 * torch.nn.functional.mse_loss( + v_hat, v_tp1_obs ) - L = torch.nn.functional.mse_loss(o_hat, o_tp1_obs) + 0.25 * torch.nn.functional.mse_loss(v_hat, v_tp1_obs) return L, {"La": float(a.mean().item() if a.numel() else 0.0)} -def ttt_constraint_loss(model, obs_feats_t, obs_mask_t, goal_feats_t, - o_t, v_t, goal_t, C_t, R_t, mask_t, d_hat_t, dt_t, - integrate_surrogate, robot_radius=None, margin_factor=0.5, - v_max=1.0, lambda_v=None,): + +def ttt_constraint_loss( + model, + obs_feats_t, + obs_mask_t, + goal_feats_t, + o_t, + v_t, + goal_t, + C_t, + R_t, + mask_t, + d_hat_t, + dt_t, + integrate_surrogate, + robot_radius=None, + margin_factor=0.5, + v_max=1.0, + lambda_v=None, +): model.train() a, b, g = model(obs_feats_t, obs_mask_t, goal_feats_t) o_hat, v_hat, clr = integrate_surrogate( - o_t, v_t, goal_t, C_t, R_t, mask_t, a, b, g, d_hat_t, dt_t, H=torch.tensor([1], device=o_t.device), + o_t, + v_t, + goal_t, + C_t, + R_t, + mask_t, + a, + b, + g, + d_hat_t, + dt_t, + H=torch.tensor([1], device=o_t.device), robot_radius=(robot_radius if robot_radius is not None else 0.0), - margin_factor=margin_factor + margin_factor=margin_factor, ) speed_next = torch.linalg.norm(v_hat) g = torch.nn.functional.relu(speed_next - v_max) @@ -230,35 +348,46 @@ def ttt_constraint_loss(model, obs_feats_t, obs_mask_t, goal_feats_t, # augmented Lagrangian term L_speed = 0.5 * g.pow(2) # proximal-to-checkpoint on head weights - return L_speed + return L_speed -def build_local_feats(o_w: np.ndarray, goal_w: np.ndarray, C_w: np.ndarray, R_w: np.ndarray, W_w: np.ndarray): + +def build_local_feats( + o_w: np.ndarray, goal_w: np.ndarray, C_w: np.ndarray, R_w: np.ndarray, W_w: np.ndarray +): """Return (obs_feats[B=1,N,6], goal_feats[B=1,4]) with training-compatible semantics: obs_feats: [cx, cy, r, w, dx_goal, dy_goal] goal_feats: [dgx, dgy, ||dg||, 1] """ o = torch.as_tensor(o_w, dtype=torch.float32) g = torch.as_tensor(goal_w, dtype=torch.float32) - C = torch.as_tensor(C_w, dtype=torch.float32) if C_w.size else torch.zeros(0,2) + C = torch.as_tensor(C_w, dtype=torch.float32) if C_w.size else torch.zeros(0, 2) R = torch.as_tensor(R_w, dtype=torch.float32) if R_w.size else torch.zeros(0) W = torch.as_tensor(W_w, dtype=torch.float32) if W_w.size else torch.zeros(0) if C.ndim == 1: - C = C.reshape(0,2) - dg = (g - o) + C = C.reshape(0, 2) + dg = g - o gdist = torch.linalg.norm(dg).unsqueeze(0) - goal_feats = torch.stack([dg[0], dg[1], gdist[0], torch.tensor(1.0)], dim=0).unsqueeze(0) # [1,4] + goal_feats = torch.stack([dg[0], dg[1], gdist[0], torch.tensor(1.0)], dim=0).unsqueeze( + 0 + ) # [1,4] if C.shape[0] == 0: - obs_feats = torch.zeros(1,0,6) + obs_feats = torch.zeros(1, 0, 6) else: - dxdy = (g.unsqueeze(0) - C) # (N,2) - obs_feats = torch.cat([C, R.unsqueeze(-1), W.unsqueeze(-1), dxdy], dim=-1).unsqueeze(0) # [1,N,6] + dxdy = g.unsqueeze(0) - C # (N,2) + obs_feats = torch.cat([C, R.unsqueeze(-1), W.unsqueeze(-1), dxdy], dim=-1).unsqueeze( + 0 + ) # [1,N,6] return obs_feats, goal_feats + # ----------------------------- # α-mapping utilities # ----------------------------- -def map_alpha_to_world(W_in: np.ndarray, R_in: np.ndarray, alphas: np.ndarray, mode: str, k_rad: float = 0.05): + +def map_alpha_to_world( + W_in: np.ndarray, R_in: np.ndarray, alphas: np.ndarray, mode: str, k_rad: float = 0.05 +): al = np.maximum(alphas, 0.0) W_out = W_in.copy() R_out = R_in.copy() @@ -268,20 +397,24 @@ def map_alpha_to_world(W_in: np.ndarray, R_in: np.ndarray, alphas: np.ndarray, m R_out = R_out + k_rad * (al if al.size else 0.0) return W_out, R_out + # ----------------------------- # Main rollout # ----------------------------- + def main(): ap = argparse.ArgumentParser("Visual evaluation for CoefEnergyNet") ap.add_argument("--ckpt", type=str, required=True) ap.add_argument("--case", type=str, default="case1-tight") ap.add_argument("--steps", type=int, default=800) ap.add_argument("--fps", type=int, default=20) - ap.add_argument("--alpha_mode", type=str, default="weight", choices=["weight","radius","both","none"]) + ap.add_argument( + "--alpha_mode", type=str, default="weight", choices=["weight", "radius", "both", "none"] + ) ap.add_argument("--seed", type=int, default=2312) - ap.add_argument("--correction", action='store_true') - ap.add_argument("--online_finetune", action='store_true') + ap.add_argument("--correction", action="store_true") + ap.add_argument("--online_finetune", action="store_true") args = ap.parse_args() device = "cuda" if torch.cuda.is_available() else "cpu" @@ -293,17 +426,18 @@ def main(): model.load_state_dict(state) gen.set_all_seeds(args.seed) # World & planner - cfg = gen.GenCfg(); cfg.seed = args.seed; + cfg = gen.GenCfg() + cfg.seed = args.seed cfg.start = np.array([-1.0, 2.0], float) - cfg.goal = np.array([ 9.0, -0.9], float) + cfg.goal = np.array([9.0, -0.9], float) cfg.d_hat = getattr(cfg, "d_hat", 0.5) if args.case.startswith("case1"): - C,R,W = gen.sample_obstacles_case1_tight(cfg) + C, R, W = gen.sample_obstacles_case1_tight(cfg) else: # Fallback to a default sampler - C,R,W = gen.sample_obstacles_case2_harder(cfg) + C, R, W = gen.sample_obstacles_case2_harder(cfg) world = ssi.WorldObstacles(C, R, W, d_hat=cfg.d_hat) # planner = gen.planner_from_cfg(cfg, world, cfg.k_bulk, cfg.gamma_s, cfg.d_hat, cfg.radius) planner = gen.planner_from_cfg(cfg, world, cfg.k_bulk, cfg.gamma_s, cfg.d_hat, cfg.radius) @@ -311,11 +445,12 @@ def main(): # Snapshots folder snap_parent_dir = f"snaps_coef/{datetime.now().strftime('%Y%m%d_%H%M%S')}" os.makedirs(snap_parent_dir, exist_ok=True) - snap_dir = os.path.join(snap_parent_dir, os.path.basename(args.ckpt).split('.')[0]) + snap_dir = os.path.join(snap_parent_dir, os.path.basename(args.ckpt).split(".")[0]) os.makedirs(snap_dir, exist_ok=True) # --- rollout - dt = cfg.dt; T = int(args.steps) + dt = cfg.dt + T = int(args.steps) frames_png = [] frames_meta = [] @@ -323,14 +458,18 @@ def main(): def capture_frame(sys): entry = {} entry["center"] = sys.o.detach().cpu().to(torch.float32) - entry["theta"] = (sys.theta.detach().cpu().to(torch.float32) - if hasattr(sys, "theta") else torch.tensor(0.0)) + entry["theta"] = ( + sys.theta.detach().cpu().to(torch.float32) + if hasattr(sys, "theta") + else torch.tensor(0.0) + ) if hasattr(sys, "Pw") and sys.Pw is not None: entry["Pw"] = sys.Pw.detach().cpu().to(torch.float32) elif hasattr(sys, "Ploc") and sys.Ploc is not None: - o = entry["center"]; theta = float(entry["theta"]) if torch.is_tensor(entry["theta"]) else entry["theta"] - c,s = math.cos(theta), math.sin(theta) - Rm = torch.tensor([[c,-s],[s,c]], dtype=torch.float32) + o = entry["center"] + theta = float(entry["theta"]) if torch.is_tensor(entry["theta"]) else entry["theta"] + c, s = math.cos(theta), math.sin(theta) + Rm = torch.tensor([[c, -s], [s, c]], dtype=torch.float32) entry["Pw"] = (sys.Ploc.detach().cpu().to(torch.float32) @ Rm.T) + o else: entry["Pw"] = sys.world_points() @@ -350,8 +489,11 @@ def capture_frame(sys): Cw, Rw, Ww = planner.stage_slice(world.C_np, world.R_np, world.W_np) # build features obs_feats, goal_feats = build_local_feats(o_w, cfg.goal, Cw, Rw, Ww) - obs_mask = torch.ones(1, obs_feats.shape[1], dtype=torch.bool, device=device) if obs_feats.shape[1] else torch.zeros(1,0, dtype=torch.bool, device=device) - + obs_mask = ( + torch.ones(1, obs_feats.shape[1], dtype=torch.bool, device=device) + if obs_feats.shape[1] + else torch.zeros(1, 0, dtype=torch.bool, device=device) + ) # predict coefficients with torch.no_grad(): @@ -360,32 +502,65 @@ def capture_frame(sys): print(t, alphas, beta, gamma) # correction # compute current metrics (cheap) - + if args.correction and new_stage_timer >= 5: if reinitialize: - controller = HistSecantController(k_alpha=1, lr_beta=0.25, lr_gamma=0.05, lr_alpha=0.4, safe_margin=0.3*getattr(cfg,"radius",0.16), prog_eps=0.02, v_min=0.05, v_max=1.0, ema=0.99) + controller = HistSecantController( + k_alpha=1, + lr_beta=0.25, + lr_gamma=0.05, + lr_alpha=0.4, + safe_margin=0.3 * getattr(cfg, "radius", 0.16), + prog_eps=0.02, + v_min=0.05, + v_max=1.0, + ema=0.99, + ) reinitialize = False - + dist_now = np.linalg.norm(o_w - cfg.goal) - speed_now = float(np.linalg.norm(planner.sys.v_o.detach().cpu().numpy())) if hasattr(planner.sys, "v_o") else 0.0 + speed_now = ( + float(np.linalg.norm(planner.sys.v_o.detach().cpu().numpy())) + if hasattr(planner.sys, "v_o") + else 0.0 + ) # quick min clearance from current slice clr_now = np.inf if Cw.shape[0] > 0: - clr_now = float(np.min(np.linalg.norm(o_w[None,:] - Cw, axis=1) - Rw)) + clr_now = float(np.min(np.linalg.norm(o_w[None, :] - Cw, axis=1) - Rw)) # create once (outside loop) # controller = HistSecantController(k_alpha=2, safe_margin=0.08, v_min=0.25) # update params without extra sims alphas_use, beta_use, gamma_use = controller.update( - alphas, beta, gamma, o_w, speed_now, cfg.goal, Cw, Rw, Ww, clr_now, dist_now, speed_now + alphas, + beta, + gamma, + o_w, + speed_now, + cfg.goal, + Cw, + Rw, + Ww, + clr_now, + dist_now, + speed_now, ) print(t, alphas_use, beta_use, gamma_use) - al_np = alphas_use.squeeze(0).detach().cpu().numpy() if obs_feats.shape[1] else np.zeros_like(Rw) + al_np = ( + alphas_use.squeeze(0).detach().cpu().numpy() + if obs_feats.shape[1] + else np.zeros_like(Rw) + ) beta_f = float(beta_use.squeeze(0).item()) gamma_f = float(gamma_use.squeeze(0).item()) else: - al_np = alphas.squeeze(0).detach().cpu().numpy() if obs_feats.shape[1] else np.zeros_like(Rw) + al_np = ( + alphas.squeeze(0).detach().cpu().numpy() + if obs_feats.shape[1] + else np.zeros_like(Rw) + ) beta_f = float(beta.squeeze(0).item()) gamma_f = float(gamma.squeeze(0).item()) # map α @@ -397,43 +572,65 @@ def capture_frame(sys): planner.sys.gamma_o = max(0.0, gamma_f) # rebuild world slice with adjusted W/R world_step = ssi.WorldObstacles(Cw, R_adj, W_adj, d_hat=cfg.d_hat) - + # advance stage_idx = planner.sm.current_stage_idx new_stage_timer += 1 - o_t = torch.as_tensor(frames_meta[-1]["center"], device=device).unsqueeze(0) - v_t = (planner.sys.v_o.detach().to(device).unsqueeze(0) - if hasattr(planner.sys, "v_o") else torch.zeros(1,2, device=device)) + o_t = torch.as_tensor(frames_meta[-1]["center"], device=device).unsqueeze(0) + v_t = ( + planner.sys.v_o.detach().to(device).unsqueeze(0) + if hasattr(planner.sys, "v_o") + else torch.zeros(1, 2, device=device) + ) info = planner.step(dt, world_step) if args.online_finetune: - + # Build single-frame tensors from your recorded previous frame (frames_meta[-1]) and the current observed next state # o_prev, v_prev from planner.sys (stored in frames_meta[-1]); o_next from current info["center"]; v_next from planner.sys.v_o - param_dtype = next(model.parameters()).dtype + param_dtype = next(model.parameters()).dtype # AFTER the step - #o_tp1 = torch.as_tensor(info["center"], device=device, dtype=param_dtype).unsqueeze(0) - #v_tp1 = planner.sys.v_o.detach().to(device, dtype=param_dtype).unsqueeze(0) if hasattr(planner.sys, "v_o") else torch.zeros(1,2, device=device, dtype=param_dtype) + # o_tp1 = torch.as_tensor(info["center"], device=device, dtype=param_dtype).unsqueeze(0) + # v_tp1 = planner.sys.v_o.detach().to(device, dtype=param_dtype).unsqueeze(0) if hasattr(planner.sys, "v_o") else torch.zeros(1,2, device=device, dtype=param_dtype) # Pack world slice as 1xN tensors - C_t = torch.as_tensor(Cw, device=device, dtype=param_dtype).unsqueeze(0) - R_t = torch.as_tensor(Rw, device=device, dtype=param_dtype).unsqueeze(0) - mask_t = torch.ones(1, Cw.shape[0], dtype=torch.bool, device=device) if Cw.shape[0] else torch.zeros(1,0,dtype=torch.bool,device=device) - goal_t = torch.as_tensor(cfg.goal, device=device, dtype=param_dtype).unsqueeze(0) + C_t = torch.as_tensor(Cw, device=device, dtype=param_dtype).unsqueeze(0) + R_t = torch.as_tensor(Rw, device=device, dtype=param_dtype).unsqueeze(0) + mask_t = ( + torch.ones(1, Cw.shape[0], dtype=torch.bool, device=device) + if Cw.shape[0] + else torch.zeros(1, 0, dtype=torch.bool, device=device) + ) + goal_t = torch.as_tensor(cfg.goal, device=device, dtype=param_dtype).unsqueeze(0) d_hat_t = torch.as_tensor([cfg.d_hat], device=device, dtype=param_dtype) - dt_t = torch.as_tensor([cfg.dt], device=device, dtype=param_dtype) + dt_t = torch.as_tensor([cfg.dt], device=device, dtype=param_dtype) # One tiny update step (prox-regularized) finetuner.model.train() finetuner.opt.zero_grad() - L_speed = ttt_constraint_loss(finetuner.model, obs_feats.to(device), obs_mask, goal_feats.to(device), - o_t, v_t, goal_t, C_t, R_t, mask_t, d_hat_t, dt_t, - integrate_surrogate=integrate_surrogate_v2, # from surrogate_robust - robot_radius=torch.as_tensor([getattr(cfg,"radius",0.0)],device=device), - margin_factor=0.5) + L_speed = ttt_constraint_loss( + finetuner.model, + obs_feats.to(device), + obs_mask, + goal_feats.to(device), + o_t, + v_t, + goal_t, + C_t, + R_t, + mask_t, + d_hat_t, + dt_t, + integrate_surrogate=integrate_surrogate_v2, # from surrogate_robust + robot_radius=torch.as_tensor([getattr(cfg, "radius", 0.0)], device=device), + margin_factor=0.5, + ) # add proximal anchor L_proxW = torch.zeros((), device=device, dtype=param_dtype) for p in finetuner.params: - L_proxW = L_proxW + (p - finetuner._anchors[id(p)].to(p.device, dtype=p.dtype)).pow(2).mean() + L_proxW = ( + L_proxW + + (p - finetuner._anchors[id(p)].to(p.device, dtype=p.dtype)).pow(2).mean() + ) # total loss: ONLY speed constraint + prox L_total = L_speed + finetuner.prox_lambda * L_proxW @@ -441,16 +638,17 @@ def capture_frame(sys): torch.nn.utils.clip_grad_norm_(finetuner.params, 1.0) finetuner.opt.step() finetuner.model.eval() - + if stage_idx != planner.sm.current_stage_idx: reinitialize = True new_stage_timer = 0 - # snapshot every k steps - if t % 3 == 0 or t == T-1: + if t % 3 == 0 or t == T - 1: png_path = os.path.join(snap_dir, f"frame_{t:04d}.png") - gen.save_episode_snapshot(png_path, planner, frames_meta, world_step, cfg.start, cfg.goal, cfg) + gen.save_episode_snapshot( + png_path, planner, frames_meta, world_step, cfg.start, cfg.goal, cfg + ) frames_png.append(png_path) frames_meta.append(capture_frame(planner.sys)) @@ -466,13 +664,18 @@ def capture_frame(sys): iio.imwrite(gif_path, frames, loop=0, fps=args.fps) try: import imageio - with imageio.get_writer(mp4_path, format="FFMPEG", mode="I", fps=args.fps, codec="libx264", quality=7) as w: - for fr in frames: w.append_data(fr) + + with imageio.get_writer( + mp4_path, format="FFMPEG", mode="I", fps=args.fps, codec="libx264", quality=7 + ) as w: + for fr in frames: + w.append_data(fr) except Exception as e: print("MP4 writer not available:", e) print({"gif": gif_path, "mp4": mp4_path}) else: print("No frames captured.") + if __name__ == "__main__": main() diff --git a/grl_snam/material_nav.py b/grl_snam/material_nav.py index 3c991c7..366b17a 100644 --- a/grl_snam/material_nav.py +++ b/grl_snam/material_nav.py @@ -40,8 +40,8 @@ from __future__ import annotations import math +from collections.abc import Mapping from dataclasses import dataclass -from typing import List, Mapping, Tuple import numpy as np import torch @@ -145,7 +145,7 @@ def forward( obs_mask: torch.Tensor, goal_feats: torch.Tensor, risk_patch: torch.Tensor, - ) -> Tuple[torch.Tensor, ...]: + ) -> tuple[torch.Tensor, ...]: B, N = obs_feats.shape[:2] z_goal = self.goal_enc(goal_feats).unsqueeze(1) @@ -213,7 +213,7 @@ def load_geometry_weights( def sdf_barrier_grad( sdf_val: torch.Tensor, d_hat_sdf: float = 3.0, k_sharp: float = 5.0 -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor]: """Softplus hazard barrier b(phi) and its derivative db/dphi. b(phi) = (1/k) * log(1 + exp(k * (d_hat - phi))) active when phi < d_hat @@ -266,7 +266,7 @@ def integrate_surrogate_material( margin_factor: float = 0.5, mass: float = 1.0, d_hat_sdf: float = 3.0, -) -> Tuple[torch.Tensor, ...]: +) -> tuple[torch.Tensor, ...]: """Surrogate integrator with material forces (faithful port). Shapes: o0/v0/goal (B,2) in (x=col, y=row) global pixels; C (B,N,2); @@ -362,8 +362,8 @@ class GateDecision: nominal_risk: float best_risk: float feasible_count: int - selected_direction_rc: Tuple[float, float] - selected_endpoint_rc: Tuple[float, float] + selected_direction_rc: tuple[float, float] + selected_endpoint_rc: tuple[float, float] selected_min_clearance_m: float @@ -374,7 +374,7 @@ def _unit(v: np.ndarray) -> np.ndarray: return (v / norm).astype(np.float32) -def _clip_rc(point_rc: np.ndarray, shape: Tuple[int, int]) -> Tuple[int, int]: +def _clip_rc(point_rc: np.ndarray, shape: tuple[int, int]) -> tuple[int, int]: # round() is round-half-to-even; the C++ twin uses std::rint to match. return ( int(np.clip(round(float(point_rc[0])), 0, shape[0] - 1)), @@ -389,7 +389,7 @@ def _ray_cost( *, horizon_cells: int, hard_margin_m: float, -) -> Tuple[float, bool, float]: +) -> tuple[float, bool, float]: """Mean sampled risk, feasibility, and min hazard clearance along one ray. Samples at integer distances 1..horizon; the float point is bounds-checked @@ -399,7 +399,7 @@ def _ray_cost( risk = maps["risk_map"] hard = maps["hard_mask"].astype(bool) sdf = maps["sdf_hard"] - values: List[float] = [] + values: list[float] = [] min_clearance = float("inf") feasible = True for distance in range(1, horizon_cells + 1): diff --git a/grl_snam/surrogate_robust.py b/grl_snam/surrogate_robust.py index 69d174c..cf0a289 100644 --- a/grl_snam/surrogate_robust.py +++ b/grl_snam/surrogate_robust.py @@ -25,16 +25,22 @@ # add to loss: L += self.cfg.w_multi * L_multi """ from __future__ import annotations -from typing import Tuple + import torch -import torch.nn.functional as F # You can import from train_coef_energy if available; otherwise include a local copy try: from grl_snam.train_coef_energy import ipc_piecewise except Exception: - def ipc_piecewise(d: torch.Tensor, d_hat: torch.Tensor | float, vp: float = -5e2, eps: float = 1e-9, - max_grad: float = 200.0, max_b: float = 200.0): + + def ipc_piecewise( + d: torch.Tensor, + d_hat: torch.Tensor | float, + vp: float = -5e2, + eps: float = 1e-9, + max_grad: float = 200.0, + max_b: float = 200.0, + ): # Minimal compatible implementation (broadcasts d_hat) if not torch.is_tensor(d_hat): dh = d.new_tensor(float(d_hat)) @@ -44,52 +50,69 @@ def ipc_piecewise(d: torch.Tensor, d_hat: torch.Tensor | float, vp: float = -5e2 dh = dh.unsqueeze(-1) dh = torch.broadcast_to(dh, d.shape) safe = torch.clamp(d, min=float(eps)) - b_in = -(d - dh) ** 2 * torch.log(safe / dh) - dbdd_in = -(2.0 * (d - dh) * torch.log(safe / dh) + (d - dh) ** 2 / safe) # M10: analytic derivative of b (was +（d-dh)+1, an attraction band); verified vs autograd + b_in = -((d - dh) ** 2) * torch.log(safe / dh) + dbdd_in = -( + 2.0 * (d - dh) * torch.log(safe / dh) + (d - dh) ** 2 / safe + ) # M10: analytic derivative of b (was +（d-dh)+1, an attraction band); verified vs autograd b = torch.where(d <= eps, d.new_tensor(vp), torch.where(d < dh, b_in, torch.zeros_like(d))) - dbdd = torch.where(d <= eps, d.new_tensor(vp), torch.where(d < dh, dbdd_in, torch.zeros_like(d))) + dbdd = torch.where( + d <= eps, d.new_tensor(vp), torch.where(d < dh, dbdd_in, torch.zeros_like(d)) + ) b = torch.clamp(b, 0.0, float(max_b)) dbdd = torch.clamp(dbdd, -float(max_grad), float(max_grad)) return b, dbdd + @torch.no_grad() def _nearest_obstacle( - o: torch.Tensor, # (B,2) - C: torch.Tensor, # (B,N,2) - R_eff: torch.Tensor, # (B,N) - mask: torch.Tensor, # (B,N) bool -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + o: torch.Tensor, # (B,2) + C: torch.Tensor, # (B,N,2) + R_eff: torch.Tensor, # (B,N) + mask: torch.Tensor, # (B,N) bool +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Return (dmin[B], nmin[B,2], jmin[B]) at o. Uses safe advanced indexing. """ B, N = C.shape[:2] if N == 0: return ( - torch.full((B,), float('inf'), device=o.device, dtype=o.dtype), + torch.full((B,), float("inf"), device=o.device, dtype=o.dtype), torch.zeros(B, 2, device=o.device, dtype=o.dtype), torch.zeros(B, dtype=torch.long, device=o.device), ) - diff = o.unsqueeze(1) - C # (B,N,2) + diff = o.unsqueeze(1) - C # (B,N,2) r = torch.linalg.norm(diff, dim=-1).clamp_min(1e-9) # (B,N) - d = r - R_eff # (B,N) + d = r - R_eff # (B,N) d = torch.where(mask, d, torch.full_like(d, 1e6)) - jmin = d.argmin(dim=1) # (B,) - idx = torch.arange(B, device=o.device) # (B,) + jmin = d.argmin(dim=1) # (B,) + idx = torch.arange(B, device=o.device) # (B,) - rmin = r[idx, jmin] # (B,) - nmin = diff[idx, jmin, :] / rmin.unsqueeze(-1) # (B,2) - dmin = d[idx, jmin] # (B,) + rmin = r[idx, jmin] # (B,) + nmin = diff[idx, jmin, :] / rmin.unsqueeze(-1) # (B,2) + dmin = d[idx, jmin] # (B,) return dmin, nmin, jmin -def integrate_surrogate_v2(o0: torch.Tensor, v0: torch.Tensor, goal: torch.Tensor, - C: torch.Tensor, R: torch.Tensor, mask: torch.Tensor, - alphas: torch.Tensor, beta: torch.Tensor, gamma: torch.Tensor, - d_hat: torch.Tensor, dt: torch.Tensor, H: torch.Tensor, - robot_radius: torch.Tensor | float = 0.0, margin_factor: float = 0.5, - mass: float = 1.0): + +def integrate_surrogate_v2( + o0: torch.Tensor, + v0: torch.Tensor, + goal: torch.Tensor, + C: torch.Tensor, + R: torch.Tensor, + mask: torch.Tensor, + alphas: torch.Tensor, + beta: torch.Tensor, + gamma: torch.Tensor, + d_hat: torch.Tensor, + dt: torch.Tensor, + H: torch.Tensor, + robot_radius: torch.Tensor | float = 0.0, + margin_factor: float = 0.5, + mass: float = 1.0, +): """Radius-aware version of integrate_surrogate. R_eff = R + margin_factor * robot_radius.""" B, N = C.shape[:2] if not torch.is_tensor(robot_radius): @@ -98,39 +121,55 @@ def integrate_surrogate_v2(o0: torch.Tensor, v0: torch.Tensor, goal: torch.Tenso rr = robot_radius.to(device=o0.device, dtype=o0.dtype) R_eff = R + margin_factor * rr[:, None] - o = o0.clone(); v = v0.clone() + o = o0.clone() + v = v0.clone() min_clear = torch.full((B,), float("inf"), dtype=o.dtype, device=o.device) for s in range(int(H.max().item())): active = (s < H).to(o.dtype).unsqueeze(-1) F_goal = -beta.unsqueeze(-1) * (o - goal) if N == 0: - F_bar = torch.zeros_like(o); dmin = torch.full_like(min_clear, float('inf')) + F_bar = torch.zeros_like(o) + dmin = torch.full_like(min_clear, float("inf")) else: diff = o.unsqueeze(1) - C r = torch.linalg.norm(diff, dim=-1).clamp_min(1e-9) n_hat = diff / r.unsqueeze(-1) d = r - R_eff d = torch.where(mask, d, torch.full_like(d, 1e6)) - _, dbdd = ipc_piecewise(d, d_hat.view(-1,1)) + _, dbdd = ipc_piecewise(d, d_hat.view(-1, 1)) F_bar = -(alphas * dbdd).unsqueeze(-1) * n_hat F_bar = F_bar.sum(dim=1) - dmin = torch.where(mask, d, torch.full_like(d, float('inf'))).min(dim=1).values + dmin = torch.where(mask, d, torch.full_like(d, float("inf"))).min(dim=1).values min_clear = torch.minimum(min_clear, dmin) a = (F_bar + F_goal - gamma.unsqueeze(-1) * v) / float(mass) v = v + active * dt.unsqueeze(-1) * a # v_{n+1} o = o + active * dt.unsqueeze(-1) * v # uses v_{n+1} - + return o, v, min_clear -def multi_start_penalty(o0: torch.Tensor, v0: torch.Tensor, goal: torch.Tensor, - C: torch.Tensor, R: torch.Tensor, mask: torch.Tensor, - alphas: torch.Tensor, beta: torch.Tensor, gamma: torch.Tensor, - d_hat: torch.Tensor, dt_prime: torch.Tensor, H: torch.Tensor, - robot_radius: torch.Tensor | float = 0.0, margin_factor: float = 0.5, - ms_count: int = 3, ms_h: int = 2, ms_dt_mult: float = 1.5, - tau: float = 0.05) -> torch.Tensor: + +def multi_start_penalty( + o0: torch.Tensor, + v0: torch.Tensor, + goal: torch.Tensor, + C: torch.Tensor, + R: torch.Tensor, + mask: torch.Tensor, + alphas: torch.Tensor, + beta: torch.Tensor, + gamma: torch.Tensor, + d_hat: torch.Tensor, + dt_prime: torch.Tensor, + H: torch.Tensor, + robot_radius: torch.Tensor | float = 0.0, + margin_factor: float = 0.5, + ms_count: int = 3, + ms_h: int = 2, + ms_dt_mult: float = 1.5, + tau: float = 0.05, +) -> torch.Tensor: """Sample feasible starts near nearest obstacle and penalize penetrations after short rollouts.""" - dev = o0.device; B, N = C.shape[:2] + B, N = C.shape[:2] if ms_count <= 0 or N == 0: return o0.new_tensor(0.0) # effective radius @@ -153,13 +192,27 @@ def multi_start_penalty(o0: torch.Tensor, v0: torch.Tensor, goal: torch.Tensor, diff_ms = o_ms.unsqueeze(1) - C r_ms = torch.linalg.norm(diff_ms, dim=-1).clamp_min(1e-9) d_ms = torch.where(mask, r_ms - R_eff, torch.full_like(r_ms, 1e6)) - ok = (d_ms.min(dim=1).values >= 0) + ok = d_ms.min(dim=1).values >= 0 o_ms = torch.where(ok.unsqueeze(-1), o_ms, o0 + 0.5 * step) # fallback if infeasible # short rollout H_ms = torch.full_like(H, ms_h) dt_ms = ms_dt_mult * dt_prime - _, _, clr_ms = integrate_surrogate_v2(o_ms, v0, goal, C, R, mask, alphas, beta, gamma, d_hat, - dt_ms, H_ms, robot_radius=rr.squeeze(-1), margin_factor=margin_factor) - L_acc = L_acc + torch.nn.functional.softplus(((- clr_ms) / tau)).mean() - + _, _, clr_ms = integrate_surrogate_v2( + o_ms, + v0, + goal, + C, + R, + mask, + alphas, + beta, + gamma, + d_hat, + dt_ms, + H_ms, + robot_radius=rr.squeeze(-1), + margin_factor=margin_factor, + ) + L_acc = L_acc + torch.nn.functional.softplus((-clr_ms) / tau).mean() + return L_acc / float(ms_count) diff --git a/grl_snam/train_coef_energy.py b/grl_snam/train_coef_energy.py index 4c34747..0924545 100644 --- a/grl_snam/train_coef_energy.py +++ b/grl_snam/train_coef_energy.py @@ -49,19 +49,22 @@ flip the flag --gamma_rel True, which makes gamma = gamma_rel * gamma_o(ep). """ from __future__ import annotations -import os, json, math, random + +import builtins +import json +import os +import random from dataclasses import dataclass -from typing import Dict, Any, List, Optional, Tuple +from typing import Any +import numpy as np import torch import torch.nn as nn import torch.nn.functional as F -from torch.utils.data import Dataset, DataLoader -from grl_snam.surrogate_robust import integrate_surrogate_v2, multi_start_penalty +from torch.utils.data import DataLoader, Dataset +from grl_snam.surrogate_robust import integrate_surrogate_v2, multi_start_penalty -import numpy as np -import builtins try: from torch.serialization import add_safe_globals, safe_globals except Exception: @@ -90,13 +93,20 @@ def safe_torch_load(path, map_location="cpu"): with safe_globals(allow): return torch.load(path, map_location=map_location, weights_only=True) + # ----------------------------- # Small math helpers (IPC piecewise) # ----------------------------- -def ipc_piecewise(d: torch.Tensor, d_hat: torch.Tensor | float, vp: float = -5e2, eps: float = 1e-9, - max_grad: float = 200.0, max_b: float = 200.0) -> Tuple[torch.Tensor, torch.Tensor]: +def ipc_piecewise( + d: torch.Tensor, + d_hat: torch.Tensor | float, + vp: float = -5e2, + eps: float = 1e-9, + max_grad: float = 200.0, + max_b: float = 200.0, +) -> tuple[torch.Tensor, torch.Tensor]: """Return (b(d), db/dd) for IPC-like barrier. Accepts scalar or batched d_hat and broadcasts.""" # Normalize d_hat and broadcast to d if not torch.is_tensor(d_hat): @@ -108,24 +118,32 @@ def ipc_piecewise(d: torch.Tensor, d_hat: torch.Tensor | float, vp: float = -5e2 dh = torch.broadcast_to(dh, d.shape) safe = torch.clamp(d, min=float(eps)) - b_in = -(d - dh) ** 2 * torch.log(safe / dh) - dbdd_in = -(2.0 * (d - dh) * torch.log(safe / dh) + (d - dh) ** 2 / safe) # M10: analytic derivative of b (was +（d-dh)+1, an attraction band); verified vs autograd + b_in = -((d - dh) ** 2) * torch.log(safe / dh) + dbdd_in = -( + 2.0 * (d - dh) * torch.log(safe / dh) + (d - dh) ** 2 / safe + ) # M10: analytic derivative of b (was +（d-dh)+1, an attraction band); verified vs autograd b = torch.where(d <= eps, d.new_tensor(vp), torch.where(d < dh, b_in, torch.zeros_like(d))) - dbdd = torch.where(d <= eps, d.new_tensor(vp), torch.where(d < dh, dbdd_in, torch.zeros_like(d))) + dbdd = torch.where( + d <= eps, d.new_tensor(vp), torch.where(d < dh, dbdd_in, torch.zeros_like(d)) + ) b = torch.clamp(b, 0.0, float(max_b)) dbdd = torch.clamp(dbdd, -float(max_grad), float(max_grad)) return b, dbdd + + # ----------------------------- # Dataset: short random rollouts from stagewise episodes # ----------------------------- + @dataclass class ShortRolloutCfg: min_h: int = 2 max_h: int = 6 - dt_mult_range: Tuple[float, float] = (1.0, 3.0) # enlarge step for robustness + dt_mult_range: tuple[float, float] = (1.0, 3.0) # enlarge step for robustness max_skip_to_end: int = 6 # if t1+1 would exceed episode, clamp within this + class ShortRollouts(Dataset): """Builds random short rollouts from the stagewise checkpoints JSONL. @@ -138,20 +156,21 @@ class ShortRollouts(Dataset): - dt_base: float32, dt' (sampled), H (int) - o_tgt, v_tgt: target state after K steps of base integrator """ - def __init__(self, root: str, cfg: Optional[ShortRolloutCfg] = None): + + def __init__(self, root: str, cfg: ShortRolloutCfg | None = None): super().__init__() self.root = root self.cfg = cfg or ShortRolloutCfg() man = os.path.join(root, "manifest.json") - with open(man, "r") as f: + with open(man) as f: self.records = json.load(f) # Preload checkpoint paths and run header (for base dt & gamma_o) - self.items: List[Dict[str, Any]] = [] + self.items: list[dict[str, Any]] = [] for rec in self.records: ep = safe_torch_load(rec["path"], map_location="cpu") logs = ep["logs"] ck_path = logs["checkpoints_jsonl"] - with open(ck_path, "r") as f: + with open(ck_path) as f: cks = [json.loads(line) for line in f] # Drop first/last if we need v estimates if len(cks) < 3: @@ -172,12 +191,12 @@ def __len__(self) -> int: return max(1, sum(len(epi["cks"]) for epi in self.items) // 6) @staticmethod - def _vel_fd(c0: List[float], c1: List[float], dt: float) -> torch.Tensor: + def _vel_fd(c0: list[float], c1: list[float], dt: float) -> torch.Tensor: c0t = torch.tensor(c0, dtype=torch.float32) c1t = torch.tensor(c1, dtype=torch.float32) return (c1t - c0t) / float(dt) - def __getitem__(self, _idx: int) -> Dict[str, Any]: + def __getitem__(self, _idx: int) -> dict[str, Any]: # pick a random episode, then a valid t0 epi = random.choice(self.items) cks = epi["cks"] @@ -213,19 +232,26 @@ def __getitem__(self, _idx: int) -> Dict[str, Any]: d_hat = float(c0["barrier"]["barrier_d_hat"]) if "barrier" in c0 else 1.0 return { - "o0": o0, "v0": v0, "goal": goal, - "C": C, "R": R, "W": W, + "o0": o0, + "v0": v0, + "goal": goal, + "C": C, + "R": R, + "W": W, "d_hat": torch.tensor(d_hat, dtype=torch.float32), "dt_base": torch.tensor(dt, dtype=torch.float32), "dt_prime": torch.tensor(dt_mult * dt, dtype=torch.float32), "H": torch.tensor(H, dtype=torch.int64), - "o_tgt": o_tgt, "v_tgt": v_tgt, + "o_tgt": o_tgt, + "v_tgt": v_tgt, "gamma_o": torch.tensor(epi["gamma_o"], dtype=torch.float32), } + # -------- collate: pad variable-N obstacles safely (N may be 0) -------- -def collate_short(batch: List[Dict[str, Any]]) -> Dict[str, Any]: + +def collate_short(batch: list[dict[str, Any]]) -> dict[str, Any]: B = len(batch) maxN = max((item["C"].shape[0] for item in batch), default=0) # tensors (B, ...) @@ -252,35 +278,50 @@ def collate_short(batch: List[Dict[str, Any]]) -> Dict[str, Any]: R[i, :n] = b["R"] W[i, :n] = b["W"] mask[i, :n] = True - + # features per obstacle: [cx, cy, r, w, dx_goal, dy_goal] - dxdy = (goal.unsqueeze(1) - C) # (B,N,2) - obs_feats = torch.cat([C, R.unsqueeze(-1), W.unsqueeze(-1), dxdy], dim=-1) if maxN>0 else torch.zeros(B,0,6) + dxdy = goal.unsqueeze(1) - C # (B,N,2) + obs_feats = ( + torch.cat([C, R.unsqueeze(-1), W.unsqueeze(-1), dxdy], dim=-1) + if maxN > 0 + else torch.zeros(B, 0, 6) + ) # goal feats (center -> goal offset and norms) - dg = (goal - o0) + dg = goal - o0 gdist = torch.linalg.norm(dg, dim=-1, keepdim=True) goal_feats = torch.cat([dg, gdist, torch.ones_like(gdist)], dim=-1) # [B,4] return { - "o0": o0, "v0": v0, "goal": goal, - "C": C, "R": R, "W": W, "obs_mask": mask, - "obs_feats": obs_feats, "goal_feats": goal_feats, - "d_hat": d_hat, "dt_base": dt_base, "dt_prime": dt_prime, "H": H, - "o_tgt": o_tgt, "v_tgt": v_tgt, "gamma_o": gamma_o, + "o0": o0, + "v0": v0, + "goal": goal, + "C": C, + "R": R, + "W": W, + "obs_mask": mask, + "obs_feats": obs_feats, + "goal_feats": goal_feats, + "d_hat": d_hat, + "dt_base": dt_base, + "dt_prime": dt_prime, + "H": H, + "o_tgt": o_tgt, + "v_tgt": v_tgt, + "gamma_o": gamma_o, } + # ----------------------------- # Model: coefficient heads over tokens # ----------------------------- + class ObstacleEncoder(nn.Module): def __init__(self, d_in=6, d_tok=64): super().__init__() - self.mlp = nn.Sequential( - nn.Linear(d_in, 128), nn.ReLU(), - nn.Linear(128, d_tok) - ) + self.mlp = nn.Sequential(nn.Linear(d_in, 128), nn.ReLU(), nn.Linear(128, d_tok)) + def forward(self, feats: torch.Tensor) -> torch.Tensor: # [B,N,d_in] B, N = feats.shape[0], feats.shape[1] if N == 0: @@ -289,21 +330,25 @@ def forward(self, feats: torch.Tensor) -> torch.Tensor: # [B,N,d_in] z = self.mlp(x).reshape(B, N, -1) return z + class CoefEnergyNet(nn.Module): """Predicts {α_j}_j, β_goal, γ from local obstacle & goal context. α_j ≥ 0 via softplus. β_goal ≥ 0 via softplus. γ ≥ 0 via softplus. If --gamma_rel, interpret γ as a *multiplier* on episode gamma_o. """ + def __init__(self, d_obs=6, d_goal=4, d_tok=64, d_ctx=64): super().__init__() self.obs_enc = ObstacleEncoder(d_in=d_obs, d_tok=d_tok) self.goal_enc = nn.Sequential(nn.Linear(d_goal, 64), nn.ReLU(), nn.Linear(64, d_tok)) - enc = nn.TransformerEncoderLayer(d_model=d_tok, nhead=4, dim_feedforward=128, batch_first=True) + enc = nn.TransformerEncoderLayer( + d_model=d_tok, nhead=4, dim_feedforward=128, batch_first=True + ) self.fuser = nn.TransformerEncoder(enc, num_layers=2) # heads self.alpha_head = nn.Sequential(nn.Linear(d_tok, 64), nn.ReLU(), nn.Linear(64, 1)) - self.beta_head = nn.Sequential(nn.Linear(d_tok, 64), nn.ReLU(), nn.Linear(64, 1)) + self.beta_head = nn.Sequential(nn.Linear(d_tok, 64), nn.ReLU(), nn.Linear(64, 1)) self.gamma_head = nn.Sequential(nn.Linear(d_tok, 64), nn.ReLU(), nn.Linear(64, 1)) def forward(self, obs_feats: torch.Tensor, obs_mask: torch.Tensor, goal_feats: torch.Tensor): @@ -318,26 +363,40 @@ def forward(self, obs_feats: torch.Tensor, obs_mask: torch.Tensor, goal_feats: t else: z_obs = self.obs_enc(obs_feats) tokens = torch.cat([z_goal, z_obs], dim=1) # [B,1+N,d] - pad = torch.cat([torch.zeros(B, 1, dtype=torch.bool, device=obs_mask.device), ~obs_mask], dim=1) + pad = torch.cat( + [torch.zeros(B, 1, dtype=torch.bool, device=obs_mask.device), ~obs_mask], dim=1 + ) z_all = self.fuser(tokens, src_key_padding_mask=pad) ctx = z_all[:, 0] # α per-obstacle from its token (z_all[:,1:]) a = self.alpha_head(z_all[:, 1:]).squeeze(-1) a = F.softplus(a) # ≥ 0 alphas = torch.where(obs_mask, a, torch.zeros_like(a)) - beta = F.softplus(self.beta_head(ctx)).squeeze(-1) # [B] + beta = F.softplus(self.beta_head(ctx)).squeeze(-1) # [B] gamma = F.softplus(self.gamma_head(ctx)).squeeze(-1) # [B] return alphas, beta, gamma + # ----------------------------- # Surrogate integrator (rigid center only) # ----------------------------- -def integrate_surrogate(o0: torch.Tensor, v0: torch.Tensor, goal: torch.Tensor, - C: torch.Tensor, R: torch.Tensor, mask: torch.Tensor, - alphas: torch.Tensor, beta: torch.Tensor, gamma: torch.Tensor, - d_hat: torch.Tensor, dt: torch.Tensor, H: torch.Tensor, - mass: float = 1.0) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + +def integrate_surrogate( + o0: torch.Tensor, + v0: torch.Tensor, + goal: torch.Tensor, + C: torch.Tensor, + R: torch.Tensor, + mask: torch.Tensor, + alphas: torch.Tensor, + beta: torch.Tensor, + gamma: torch.Tensor, + d_hat: torch.Tensor, + dt: torch.Tensor, + H: torch.Tensor, + mass: float = 1.0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Integrate B independent samples for H_i steps with step dt_i. All inputs are batched (B, ...). Returns (oT, vT, min_clear_along) for penalty. @@ -358,7 +417,7 @@ def integrate_surrogate(o0: torch.Tensor, v0: torch.Tensor, goal: torch.Tensor, F_bar = torch.zeros_like(o) dmin = torch.full_like(min_clear, float("inf")) else: - diff = o.unsqueeze(1) - C # (B,N,2) + diff = o.unsqueeze(1) - C # (B,N,2) r = torch.linalg.norm(diff, dim=-1).clamp_min(1e-9) # (B,N) n_hat = diff / r.unsqueeze(-1) d = r - R @@ -377,10 +436,12 @@ def integrate_surrogate(o0: torch.Tensor, v0: torch.Tensor, goal: torch.Tensor, v = v + active * dt.unsqueeze(-1) * a return o, v, min_clear + # ----------------------------- # Trainer # ----------------------------- + @dataclass class TrainCfg: epochs: int = 50 @@ -392,13 +453,14 @@ class TrainCfg: w_friction: float = 0.1 w_clear: float = 5e-3 # penalty if predicted path penetrates gamma_rel: bool = False - margin_factor: float = 0.5 # minimal squeeze margin = 0.5 * radius - w_multi: float = 0.5 # weight for multi-start penalty - ms_count: int = 20 # # of aux starts per sample - ms_h: int = 3 # short horizon for each aux rollout - ms_dt_mult: float = 4.0 # enlarge dt for robustness + margin_factor: float = 0.5 # minimal squeeze margin = 0.5 * radius + w_multi: float = 0.5 # weight for multi-start penalty + ms_count: int = 20 # # of aux starts per sample + ms_h: int = 3 # short horizon for each aux rollout + ms_dt_mult: float = 4.0 # enlarge dt for robustness device: str = "cuda" if torch.cuda.is_available() else "cpu" + class Trainer: def __init__(self, model: CoefEnergyNet, cfg: TrainCfg): self.model = model.to(cfg.device) @@ -406,15 +468,13 @@ def __init__(self, model: CoefEnergyNet, cfg: TrainCfg): self.cfg = cfg self.version = 2 - - def step_batch(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: + def step_batch(self, batch: dict[str, torch.Tensor]) -> dict[str, float]: dev = self.cfg.device o0 = batch["o0"].to(dev) v0 = batch["v0"].to(dev) goal = batch["goal"].to(dev) C = batch["C"].to(dev) R = batch["R"].to(dev) - W = batch["W"].to(dev) mask = batch["obs_mask"].to(dev) obs_feats = batch["obs_feats"].to(dev) goal_feats = batch["goal_feats"].to(dev) @@ -430,64 +490,81 @@ def step_batch(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: if self.cfg.gamma_rel: gamma = gamma * gamma_o # interpret as multiplier on ep gamma_o - # integrate surrogate if self.version == 2: - + # ... after you’ve pulled batch tensors, and predicted alphas, beta, gamma: # main rollout (radius-aware) oT, vT, clr = integrate_surrogate_v2( - o0, v0, goal, C, R, mask, - alphas, beta, gamma, - d_hat, dt_prime, H, + o0, + v0, + goal, + C, + R, + mask, + alphas, + beta, + gamma, + d_hat, + dt_prime, + H, robot_radius=batch.get("radius", torch.zeros(o0.shape[0], device=o0.device)), - margin_factor=self.cfg.margin_factor # set to 0.5 in cfg + margin_factor=self.cfg.margin_factor, # set to 0.5 in cfg ) # losses L_traj = F.mse_loss(oT, o_tgt) - L_vel = F.mse_loss(vT, v_tgt) - L_alpha = (alphas.mean() if alphas.numel() else torch.tensor(0.0, device=dev)) - + L_vel = F.mse_loss(vT, v_tgt) + L_alpha = alphas.mean() if alphas.numel() else torch.tensor(0.0, device=dev) # multi-start robustness penalty near obstacles L_multi = multi_start_penalty( - o0, v0, goal, C, R, mask, - alphas, beta, gamma, - d_hat, dt_prime, H, + o0, + v0, + goal, + C, + R, + mask, + alphas, + beta, + gamma, + d_hat, + dt_prime, + H, robot_radius=batch.get("radius", 0.0), margin_factor=self.cfg.margin_factor, - ms_count=self.cfg.ms_count, # e.g., 3 - ms_h=self.cfg.ms_h, # e.g., 2 - ms_dt_mult=self.cfg.ms_dt_mult # e.g., 1.5 + ms_count=self.cfg.ms_count, # e.g., 3 + ms_h=self.cfg.ms_h, # e.g., 2 + ms_dt_mult=self.cfg.ms_dt_mult, # e.g., 1.5 ) - L_friction = F.mse_loss(gamma, gamma_o) # add to loss L = ( self.cfg.w_traj * L_traj - + self.cfg.w_vel * L_vel - + self.cfg.w_friction * L_friction - + self.cfg.w_multi * L_multi - #+ self.cfg.w_stage * L_stage # if you attached the real stagewise hook + + self.cfg.w_vel * L_vel + + self.cfg.w_friction * L_friction + + self.cfg.w_multi * L_multi + # + self.cfg.w_stage * L_stage # if you attached the real stagewise hook ) # barrier penalty if min clearance < 0 (penetration) m = self.cfg.margin_factor * batch.get("radius", 0.0) # per-sample margin tau = getattr(self.cfg, "prox_tau", 0.05) - pen = torch.nn.functional.softplus(((m - clr) / tau)).mean() + pen = torch.nn.functional.softplus((m - clr) / tau).mean() # add weight # L = L + getattr(self.cfg, "w_prox", 0.1) * pen else: - oT, vT, clr = integrate_surrogate(o0, v0, goal, C, R, mask, alphas, beta, gamma, d_hat, dt_prime, H) + oT, vT, clr = integrate_surrogate( + o0, v0, goal, C, R, mask, alphas, beta, gamma, d_hat, dt_prime, H + ) # losses L_traj = F.mse_loss(oT, o_tgt) - L_vel = F.mse_loss(vT, v_tgt) - L_alpha = (alphas.mean() if alphas.numel() else torch.tensor(0.0, device=dev)) + L_vel = F.mse_loss(vT, v_tgt) + L_alpha = alphas.mean() if alphas.numel() else torch.tensor(0.0, device=dev) # barrier penalty if min clearance < 0 (penetration) pen = F.relu(-clr).mean() L = self.cfg.w_traj * L_traj + self.cfg.w_vel * L_vel + self.cfg.w_clear * pen @@ -504,17 +581,20 @@ def step_batch(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: "friction": float(L_friction.item()), "alpha": float(L_alpha.item() if alphas.numel() else 0.0), "multi_val": float(L_multi.item()), - "pen": float(pen.item()) + "pen": float(pen.item()), } + # ----------------------------- # CLI # ----------------------------- + def main(): # python -m train_coef_energy --root ./nav_stagewise_hyperring --epochs 50 --bs 128 --lr 1e-4 --workers 4 --outdir checkpoints/coef_energy_radius --save-every 10 # default 0.1 friction + 0.5 multi-check (for alpha) import argparse + ap = argparse.ArgumentParser("Wrapper trainer with checkpointing") ap.add_argument("--root", type=str, required=True) ap.add_argument("--epochs", type=int, default=50) @@ -522,32 +602,56 @@ def main(): ap.add_argument("--lr", type=float, default=3e-4) ap.add_argument("--workers", type=int, default=4) ap.add_argument("--gamma_rel", action="store_true") - ap.add_argument("--outdir", type=str, default=os.environ.get("CKPT_DIR", "checkpoints/coef_energy")) + ap.add_argument( + "--outdir", type=str, default=os.environ.get("CKPT_DIR", "checkpoints/coef_energy") + ) ap.add_argument("--save-every", type=int, default=1) ap.add_argument("--w_friction", type=float, default=0.1) ap.add_argument("--w_multi", type=float, default=0.5) - + args = ap.parse_args() # data ds = ShortRollouts(args.root) - dl = DataLoader(ds, batch_size=args.bs, shuffle=True, num_workers=args.workers, - collate_fn=collate_short, drop_last=False) + dl = DataLoader( + ds, + batch_size=args.bs, + shuffle=True, + num_workers=args.workers, + collate_fn=collate_short, + drop_last=False, + ) # model + trainer model = CoefEnergyNet() - tcfg = TrainCfg(epochs=args.epochs, bs=args.bs, lr=args.lr, workers=args.workers, gamma_rel=args.gamma_rel, w_friction=args.w_friction, w_multi=args.w_multi) + tcfg = TrainCfg( + epochs=args.epochs, + bs=args.bs, + lr=args.lr, + workers=args.workers, + gamma_rel=args.gamma_rel, + w_friction=args.w_friction, + w_multi=args.w_multi, + ) trainer = Trainer(model, tcfg) os.makedirs(args.outdir, exist_ok=True) - best = float('inf') + best = float("inf") for ep in range(tcfg.epochs): - logs_acc = {"loss":0.0,"traj":0.0,"vel":0.0,"alpha":0.0,"friction":0.0,"multi_val":0.0, "pen":0.0} + logs_acc = { + "loss": 0.0, + "traj": 0.0, + "vel": 0.0, + "alpha": 0.0, + "friction": 0.0, + "multi_val": 0.0, + "pen": 0.0, + } n = 0 for batch in dl: logs = trainer.step_batch(batch) - for k,v in logs.items(): + for k, v in logs.items(): logs_acc[k] += v n += 1 for k in logs_acc: @@ -556,19 +660,19 @@ def main(): # --- checkpointing --- ckpt = { - 'epoch': ep, - 'model_state_dict': model.state_dict(), - 'optimizer_state_dict': trainer.opt.state_dict(), - 'avg_logs': logs_acc, - 'train_cfg': tcfg.__dict__, + "epoch": ep, + "model_state_dict": model.state_dict(), + "optimizer_state_dict": trainer.opt.state_dict(), + "avg_logs": logs_acc, + "train_cfg": tcfg.__dict__, } if ep % args.save_every == 0: torch.save(ckpt, os.path.join(args.outdir, f"epoch_{ep:03d}.pt")) torch.save(ckpt, os.path.join(args.outdir, "latest.pt")) - if logs_acc.get('traj', 1e9) < best: - best = logs_acc['traj'] + if logs_acc.get("traj", 1e9) < best: + best = logs_acc["traj"] torch.save(ckpt, os.path.join(args.outdir, "best.pt")) - + if __name__ == "__main__": main() diff --git a/pyproject.toml b/pyproject.toml index e7246d9..ec640b4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -74,9 +74,10 @@ target-version = ["py310", "py311", "py312"] [tool.ruff] line-length = 100 target-version = "py310" -# Only lint the new packaging shim, tests, and src/. The legacy flat-layout -# research modules (experiments/, scripts/, train_*.py, eval_*.py, etc.) are -# kept out of CI lint to avoid churn on imported third-party / generated code. +# Lint the whole grl_snam package + tests. The research core moved in from the repo +# root (sdf_nav / material_nav / train_coef_energy / eval_coef_energy / +# surrogate_robust) is now under grl_snam/ and held to the package lint standard +# (black + ruff), like the rest of it. experiments/ and scripts/ stay unlinted. include = ["grl_snam/**/*.py", "tests/**/*.py"] extend-exclude = ["build", "dist", "__pycache__", ".venv", "data"] From 50459c6c2f364092ac9f4eb1b5da8a4e21d18175 Mon Sep 17 00:00:00 2001 From: Joe Rivera Date: Sun, 27 Sep 2026 01:39:11 -0500 Subject: [PATCH 3/3] refactor: reformat with black 26.x (matrix black-version skew) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The cp313 CI leg uses a newer black (its cvcpkg closure is mislaid -> pip-latest) than cp312's pinned black; black 24.x formatting of the moved modules was rejected by cp313's newer black. Reformat the three with black 26.x — verified accepted by BOTH 26.5.1 and 24.10.0 (a fixed point across the matrix), ruff still clean, AST-preserving so behavior is unchanged. --- grl_snam/eval_coef_energy.py | 1 + grl_snam/surrogate_robust.py | 1 + grl_snam/train_coef_energy.py | 1 + 3 files changed, 3 insertions(+) diff --git a/grl_snam/eval_coef_energy.py b/grl_snam/eval_coef_energy.py index 4bb3dd9..4239dc3 100644 --- a/grl_snam/eval_coef_energy.py +++ b/grl_snam/eval_coef_energy.py @@ -25,6 +25,7 @@ save_episode_snapshot (same API as your old eval file). • Your stagewise module is available as scripts.spline_stagewise6 (or adapt imports below). """ + from __future__ import annotations import argparse diff --git a/grl_snam/surrogate_robust.py b/grl_snam/surrogate_robust.py index cf0a289..dba5a35 100644 --- a/grl_snam/surrogate_robust.py +++ b/grl_snam/surrogate_robust.py @@ -24,6 +24,7 @@ # add to loss: L += self.cfg.w_multi * L_multi """ + from __future__ import annotations import torch diff --git a/grl_snam/train_coef_energy.py b/grl_snam/train_coef_energy.py index 0924545..8fa4e10 100644 --- a/grl_snam/train_coef_energy.py +++ b/grl_snam/train_coef_energy.py @@ -48,6 +48,7 @@ • If you want to make damping relative to the episode's nominal gamma_o, you can flip the flag --gamma_rel True, which makes gamma = gamma_rel * gamma_o(ep). """ + from __future__ import annotations import builtins