aurasync-engine 0.4.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- aurasync_engine-0.4.1/LICENSE +21 -0
- aurasync_engine-0.4.1/PKG-INFO +132 -0
- aurasync_engine-0.4.1/README.md +116 -0
- aurasync_engine-0.4.1/aurasync/__init__.py +0 -0
- aurasync_engine-0.4.1/aurasync/conditions.py +84 -0
- aurasync_engine-0.4.1/aurasync/controller.py +36 -0
- aurasync_engine-0.4.1/aurasync/estimator.py +251 -0
- aurasync_engine-0.4.1/aurasync/event_camera.py +196 -0
- aurasync_engine-0.4.1/aurasync/frame_camera.py +163 -0
- aurasync_engine-0.4.1/aurasync/imu.py +98 -0
- aurasync_engine-0.4.1/aurasync/monitor.py +80 -0
- aurasync_engine-0.4.1/aurasync/physics.py +62 -0
- aurasync_engine-0.4.1/aurasync/safety.py +173 -0
- aurasync_engine-0.4.1/aurasync/scheduler.py +191 -0
- aurasync_engine-0.4.1/aurasync/simulate.py +168 -0
- aurasync_engine-0.4.1/aurasync/v4/__init__.py +0 -0
- aurasync_engine-0.4.1/aurasync/v4/geo_pipeline.py +61 -0
- aurasync_engine-0.4.1/aurasync/v4/geometry.py +97 -0
- aurasync_engine-0.4.1/aurasync/v4/pipeline.py +274 -0
- aurasync_engine-0.4.1/aurasync/v41/__init__.py +0 -0
- aurasync_engine-0.4.1/aurasync/v41/pipeline.py +269 -0
- aurasync_engine-0.4.1/aurasync_engine.egg-info/PKG-INFO +132 -0
- aurasync_engine-0.4.1/aurasync_engine.egg-info/SOURCES.txt +28 -0
- aurasync_engine-0.4.1/aurasync_engine.egg-info/dependency_links.txt +1 -0
- aurasync_engine-0.4.1/aurasync_engine.egg-info/requires.txt +4 -0
- aurasync_engine-0.4.1/aurasync_engine.egg-info/top_level.txt +1 -0
- aurasync_engine-0.4.1/pyproject.toml +20 -0
- aurasync_engine-0.4.1/setup.cfg +4 -0
- aurasync_engine-0.4.1/tests/test_v4.py +155 -0
- aurasync_engine-0.4.1/tests/test_v41.py +138 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Sahil Kundu
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: aurasync-engine
|
|
3
|
+
Version: 0.4.1
|
|
4
|
+
Summary: Event-camera state estimator and controller for a cart-pole (simulation prototype)
|
|
5
|
+
Author: Sahil Kundu
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Repository, https://github.com/sahil03kundu-code/AuraSync-Engine
|
|
8
|
+
Requires-Python: >=3.9
|
|
9
|
+
Description-Content-Type: text/markdown
|
|
10
|
+
License-File: LICENSE
|
|
11
|
+
Requires-Dist: numpy
|
|
12
|
+
Requires-Dist: scipy
|
|
13
|
+
Requires-Dist: torch
|
|
14
|
+
Requires-Dist: matplotlib
|
|
15
|
+
Dynamic: license-file
|
|
16
|
+
|
|
17
|
+
# AuraSync-Engine V4 — Detect → Correct → Recover (simulation prototype)
|
|
18
|
+
|
|
19
|
+
An event-camera state estimator and controller for an inverted pendulum (cart-pole) that
|
|
20
|
+
is hit by sudden impulses. V4 adds a way to **detect** when the learned observer is
|
|
21
|
+
confidently wrong, **correct** it from independent event geometry, and **recover** with
|
|
22
|
+
an uncertainty-aware safety supervisor.
|
|
23
|
+
|
|
24
|
+
**Everything here runs in simulation on a 2-core CPU.** No hardware result is claimed.
|
|
25
|
+
The V3 README is in `docs/V3_README.md`.
|
|
26
|
+
|
|
27
|
+
## V4.1 update: geometry-first, evidence-aware, selective neural assistance
|
|
28
|
+
|
|
29
|
+
**Result (held-out seeds 4000+, 272 paired stress trials):**
|
|
30
|
+
- **Adding the network did not help.** V4.1 against the same pipeline without the network
|
|
31
|
+
(D-ev): survival −1.8 points [−7, +3], safe recovery −0.4 [−6, +5].
|
|
32
|
+
- **The evidence fix is the real gain.** D-ev against Baseline D: survival
|
|
33
|
+
**+7.7 [+2, +13]**, safe recovery **+9.2 [+4, +14]**, and occlusion survival 45% → **84%**.
|
|
34
|
+
- **Recommended default:** D-ev, which needs no network and 0.73 cores (V4: 1.83).
|
|
35
|
+
|
|
36
|
+
Details: [`docs/V41_REPORT.md`](docs/V41_REPORT.md) · audit: [`docs/V41_AUDIT.md`](docs/V41_AUDIT.md) ·
|
|
37
|
+
code: `aurasync/v41/pipeline.py`, `benchmark_v41.py`, `analyze_v41.py`, `v41/` · tests: `tests/test_v41.py`.
|
|
38
|
+
|
|
39
|
+

|
|
40
|
+
|
|
41
|
+
## The idea in one analogy
|
|
42
|
+
The V3 neural observer is a ship's navigator who dead-reckons. After a big wave he is
|
|
43
|
+
sure of his position, and wrong. In V3, that was the cause of every fall.
|
|
44
|
+
|
|
45
|
+
V4 adds three things:
|
|
46
|
+
- **a lookout** who takes direct bearings on the mast: a RANSAC line fit on events less
|
|
47
|
+
than 3 ms old;
|
|
48
|
+
- **a rule for when to believe him:** a CUSUM consistency detector plus a gated Kalman
|
|
49
|
+
update, with re-initialisation only after several consistent bearings;
|
|
50
|
+
- **a way to warn the captain** (the supervisor) when nobody is sure.
|
|
51
|
+
|
|
52
|
+
## Headline results (test seeds, never used for any choice)
|
|
53
|
+
|
|
54
|
+

|
|
55
|
+
|
|
56
|
+
| Nominal, 112 trials/seed | Survival [95% CI] | Safe (≤ 12°) | Tilt error |
|
|
57
|
+
|---|---|---|---|
|
|
58
|
+
| **V4 proposed** (3 seeds) | **96.4 [93, 99]** | **62.5** | 1.22° |
|
|
59
|
+
| C: V3 + uncertainty + shield (3 seeds) | 54.8 [47, 62] | 32.1 | 6.38° |
|
|
60
|
+
| B: V3 triggered (3 seeds) | 73.8 [68, 79] | 16.7 | 4.96° |
|
|
61
|
+
| E: learned frame camera 120 fps (3 seeds) | 84.5 [79, 90] | 45.8 | 1.41° |
|
|
62
|
+
| **D: event geometry + supervisor, no network** | **99.1 [97, 100]** | **69.6** | 1.59° |
|
|
63
|
+
| F: IMU + encoder | 100 | 57.1 | 0.06° |
|
|
64
|
+
|
|
65
|
+
**What holds up:**
|
|
66
|
+
- **V4 vs V3 + shield:** survival +41.7 points [+34, +49], safe recovery +30.4
|
|
67
|
+
[+23, +37].
|
|
68
|
+
- **Faults V3 cannot survive:** V4 survives the injected sign flip and the confident
|
|
69
|
+
15° error on 16/16 trials (V3 + shield: 44% and 38%). Its alarm fires a median 4 ms
|
|
70
|
+
after the fault.
|
|
71
|
+
- **HDR:** 100% survival, where V3 and the frame camera get 0%.
|
|
72
|
+
|
|
73
|
+
**What does not:**
|
|
74
|
+
- **Baseline D**, the same detect-and-correct idea without the neural network, is level
|
|
75
|
+
with or better than V4 in 14 of 15 non-nominal conditions. In nominal conditions its
|
|
76
|
+
safe recovery is 7.1 points higher [2, 13], at a sixth of the compute.
|
|
77
|
+
- **Occlusion** is the only condition where V4 is clearly ahead (69% vs 44% survival).
|
|
78
|
+
- **Correction is essential.** Detection plus uncertainty inflation without it is worse
|
|
79
|
+
than V3: 8% nominal survival.
|
|
80
|
+
- **Uncertainty inflation** showed no benefit.
|
|
81
|
+
- **The first V4 design** (the network overwrites the state) had the lowest error and
|
|
82
|
+
27% survival.
|
|
83
|
+
- **Compute:** default V4 needs about 1.5 CPU cores of compute per second of control.
|
|
84
|
+
|
|
85
|
+
Full report with CIs, ablations, failure analyses and limitations:
|
|
86
|
+
**[`docs/V4_REPORT.md`](docs/V4_REPORT.md)**.
|
|
87
|
+
|
|
88
|
+

|
|
89
|
+
|
|
90
|
+
## Repository map
|
|
91
|
+
|
|
92
|
+
| Path | What it is |
|
|
93
|
+
|---|---|
|
|
94
|
+
| `aurasync/v4/geometry.py` | Module A: EventGeometry, the RANSAC pole-line measurement |
|
|
95
|
+
| `aurasync/v4/pipeline.py` | Modules A–C, E: `ConsistencyDetector`, `V4Pipeline` (gated fusion, confirmed re-init, inflation, position-only neural fusion, compute budget) |
|
|
96
|
+
| `aurasync/v4/geo_pipeline.py` | Baseline D: physics + gated geometry, no network |
|
|
97
|
+
| `aurasync/safety.py` | Module D: CBF supervisor, delay compensation, stale fallback (from V3) |
|
|
98
|
+
| `benchmark_v4.py` | Test benchmark: baselines A–F, V4, ablations, 9 stress conditions, 6 new fault scenarios |
|
|
99
|
+
| `analyze_v4.py` | Statistics (paired bootstrap CIs, Cohen's h) and figures → `results/v4/v4_report.json` |
|
|
100
|
+
| `v4/repro_v3.py` | Step 2: bit-exact V3 reproduction |
|
|
101
|
+
| `v4/collect_dev.py`, `v4/eval_dev.py` | Steps 3–5: validation failure catalogue and open-loop detector/correction evaluation |
|
|
102
|
+
| `v4/select_v4.py`, `v4/tune_v4.py` | Step 6: design selection and supervisor tuning (validation seeds 300–315 only) |
|
|
103
|
+
| `v4/latency_v4.py` | Module E: per-stage timings |
|
|
104
|
+
| `v4/plot_failures.py`, `v4/plot_test_cases.py` | Failure visualisations (validation and test) plus per-tick CSV logs |
|
|
105
|
+
| `tests/test_v4.py` | Unit and integration tests |
|
|
106
|
+
| `docs/` | V3 audit, V4 design, latency report, technical report, hardware-validation plan |
|
|
107
|
+
|
|
108
|
+
## Reproduce
|
|
109
|
+
```bash
|
|
110
|
+
pip install torch numpy scipy matplotlib --break-system-packages
|
|
111
|
+
python3 -m unittest discover -s tests -v # 12 tests; 4 need the V3 checkpoints + dev set
|
|
112
|
+
python3 v4/repro_v3.py # V3 reproduction check
|
|
113
|
+
python3 v4/collect_dev.py && python3 v4/eval_dev.py # validation set (about 3.4 GB in data/)
|
|
114
|
+
python3 v4/select_v4.py && python3 v4/tune_v4.py # selection and tuning (validation only)
|
|
115
|
+
WHICH=all python3 benchmark_v4.py # test benchmark, resumable (about 2 h on 2 cores)
|
|
116
|
+
python3 analyze_v4.py # tables, CIs, figures
|
|
117
|
+
python3 v4/latency_v4.py # timings (run on an idle CPU)
|
|
118
|
+
```
|
|
119
|
+
Runs are deterministic (fixed seeds, 1 thread per worker), and the V3 checkpoints in
|
|
120
|
+
`checkpoints/v3/` are reused unchanged. Rules that hold throughout:
|
|
121
|
+
- **Ground truth:** estimators never read it. Replay tests pass `s_true=None`.
|
|
122
|
+
- **Configuration:** stored in `results/v4/tuning_v4.json` (supervisor),
|
|
123
|
+
`checkpoints/v3/calibration.json` and `results/v3/tuning.json` (inherited).
|
|
124
|
+
|
|
125
|
+
## Integrity notes
|
|
126
|
+
- **Simulation only.** No real event sensor, actuator or timing hardware.
|
|
127
|
+
- **Latency.** The 1 kHz control loop is not a demonstrated 1 ms sensor-to-actuator
|
|
128
|
+
latency. See `docs/V4_LATENCY.md`.
|
|
129
|
+
- **No priority claims.** No prior-art search was done, and nothing here is claimed as
|
|
130
|
+
first of its kind, formally safe or patentable.
|
|
131
|
+
- **Sample size.** Stress and scenario cells use 16 trials and mostly one seed, so
|
|
132
|
+
differences below about 20 points there are not established.
|
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
# AuraSync-Engine V4 — Detect → Correct → Recover (simulation prototype)
|
|
2
|
+
|
|
3
|
+
An event-camera state estimator and controller for an inverted pendulum (cart-pole) that
|
|
4
|
+
is hit by sudden impulses. V4 adds a way to **detect** when the learned observer is
|
|
5
|
+
confidently wrong, **correct** it from independent event geometry, and **recover** with
|
|
6
|
+
an uncertainty-aware safety supervisor.
|
|
7
|
+
|
|
8
|
+
**Everything here runs in simulation on a 2-core CPU.** No hardware result is claimed.
|
|
9
|
+
The V3 README is in `docs/V3_README.md`.
|
|
10
|
+
|
|
11
|
+
## V4.1 update: geometry-first, evidence-aware, selective neural assistance
|
|
12
|
+
|
|
13
|
+
**Result (held-out seeds 4000+, 272 paired stress trials):**
|
|
14
|
+
- **Adding the network did not help.** V4.1 against the same pipeline without the network
|
|
15
|
+
(D-ev): survival −1.8 points [−7, +3], safe recovery −0.4 [−6, +5].
|
|
16
|
+
- **The evidence fix is the real gain.** D-ev against Baseline D: survival
|
|
17
|
+
**+7.7 [+2, +13]**, safe recovery **+9.2 [+4, +14]**, and occlusion survival 45% → **84%**.
|
|
18
|
+
- **Recommended default:** D-ev, which needs no network and 0.73 cores (V4: 1.83).
|
|
19
|
+
|
|
20
|
+
Details: [`docs/V41_REPORT.md`](docs/V41_REPORT.md) · audit: [`docs/V41_AUDIT.md`](docs/V41_AUDIT.md) ·
|
|
21
|
+
code: `aurasync/v41/pipeline.py`, `benchmark_v41.py`, `analyze_v41.py`, `v41/` · tests: `tests/test_v41.py`.
|
|
22
|
+
|
|
23
|
+

|
|
24
|
+
|
|
25
|
+
## The idea in one analogy
|
|
26
|
+
The V3 neural observer is a ship's navigator who dead-reckons. After a big wave he is
|
|
27
|
+
sure of his position, and wrong. In V3, that was the cause of every fall.
|
|
28
|
+
|
|
29
|
+
V4 adds three things:
|
|
30
|
+
- **a lookout** who takes direct bearings on the mast: a RANSAC line fit on events less
|
|
31
|
+
than 3 ms old;
|
|
32
|
+
- **a rule for when to believe him:** a CUSUM consistency detector plus a gated Kalman
|
|
33
|
+
update, with re-initialisation only after several consistent bearings;
|
|
34
|
+
- **a way to warn the captain** (the supervisor) when nobody is sure.
|
|
35
|
+
|
|
36
|
+
## Headline results (test seeds, never used for any choice)
|
|
37
|
+
|
|
38
|
+

|
|
39
|
+
|
|
40
|
+
| Nominal, 112 trials/seed | Survival [95% CI] | Safe (≤ 12°) | Tilt error |
|
|
41
|
+
|---|---|---|---|
|
|
42
|
+
| **V4 proposed** (3 seeds) | **96.4 [93, 99]** | **62.5** | 1.22° |
|
|
43
|
+
| C: V3 + uncertainty + shield (3 seeds) | 54.8 [47, 62] | 32.1 | 6.38° |
|
|
44
|
+
| B: V3 triggered (3 seeds) | 73.8 [68, 79] | 16.7 | 4.96° |
|
|
45
|
+
| E: learned frame camera 120 fps (3 seeds) | 84.5 [79, 90] | 45.8 | 1.41° |
|
|
46
|
+
| **D: event geometry + supervisor, no network** | **99.1 [97, 100]** | **69.6** | 1.59° |
|
|
47
|
+
| F: IMU + encoder | 100 | 57.1 | 0.06° |
|
|
48
|
+
|
|
49
|
+
**What holds up:**
|
|
50
|
+
- **V4 vs V3 + shield:** survival +41.7 points [+34, +49], safe recovery +30.4
|
|
51
|
+
[+23, +37].
|
|
52
|
+
- **Faults V3 cannot survive:** V4 survives the injected sign flip and the confident
|
|
53
|
+
15° error on 16/16 trials (V3 + shield: 44% and 38%). Its alarm fires a median 4 ms
|
|
54
|
+
after the fault.
|
|
55
|
+
- **HDR:** 100% survival, where V3 and the frame camera get 0%.
|
|
56
|
+
|
|
57
|
+
**What does not:**
|
|
58
|
+
- **Baseline D**, the same detect-and-correct idea without the neural network, is level
|
|
59
|
+
with or better than V4 in 14 of 15 non-nominal conditions. In nominal conditions its
|
|
60
|
+
safe recovery is 7.1 points higher [2, 13], at a sixth of the compute.
|
|
61
|
+
- **Occlusion** is the only condition where V4 is clearly ahead (69% vs 44% survival).
|
|
62
|
+
- **Correction is essential.** Detection plus uncertainty inflation without it is worse
|
|
63
|
+
than V3: 8% nominal survival.
|
|
64
|
+
- **Uncertainty inflation** showed no benefit.
|
|
65
|
+
- **The first V4 design** (the network overwrites the state) had the lowest error and
|
|
66
|
+
27% survival.
|
|
67
|
+
- **Compute:** default V4 needs about 1.5 CPU cores of compute per second of control.
|
|
68
|
+
|
|
69
|
+
Full report with CIs, ablations, failure analyses and limitations:
|
|
70
|
+
**[`docs/V4_REPORT.md`](docs/V4_REPORT.md)**.
|
|
71
|
+
|
|
72
|
+

|
|
73
|
+
|
|
74
|
+
## Repository map
|
|
75
|
+
|
|
76
|
+
| Path | What it is |
|
|
77
|
+
|---|---|
|
|
78
|
+
| `aurasync/v4/geometry.py` | Module A: EventGeometry, the RANSAC pole-line measurement |
|
|
79
|
+
| `aurasync/v4/pipeline.py` | Modules A–C, E: `ConsistencyDetector`, `V4Pipeline` (gated fusion, confirmed re-init, inflation, position-only neural fusion, compute budget) |
|
|
80
|
+
| `aurasync/v4/geo_pipeline.py` | Baseline D: physics + gated geometry, no network |
|
|
81
|
+
| `aurasync/safety.py` | Module D: CBF supervisor, delay compensation, stale fallback (from V3) |
|
|
82
|
+
| `benchmark_v4.py` | Test benchmark: baselines A–F, V4, ablations, 9 stress conditions, 6 new fault scenarios |
|
|
83
|
+
| `analyze_v4.py` | Statistics (paired bootstrap CIs, Cohen's h) and figures → `results/v4/v4_report.json` |
|
|
84
|
+
| `v4/repro_v3.py` | Step 2: bit-exact V3 reproduction |
|
|
85
|
+
| `v4/collect_dev.py`, `v4/eval_dev.py` | Steps 3–5: validation failure catalogue and open-loop detector/correction evaluation |
|
|
86
|
+
| `v4/select_v4.py`, `v4/tune_v4.py` | Step 6: design selection and supervisor tuning (validation seeds 300–315 only) |
|
|
87
|
+
| `v4/latency_v4.py` | Module E: per-stage timings |
|
|
88
|
+
| `v4/plot_failures.py`, `v4/plot_test_cases.py` | Failure visualisations (validation and test) plus per-tick CSV logs |
|
|
89
|
+
| `tests/test_v4.py` | Unit and integration tests |
|
|
90
|
+
| `docs/` | V3 audit, V4 design, latency report, technical report, hardware-validation plan |
|
|
91
|
+
|
|
92
|
+
## Reproduce
|
|
93
|
+
```bash
|
|
94
|
+
pip install torch numpy scipy matplotlib --break-system-packages
|
|
95
|
+
python3 -m unittest discover -s tests -v # 12 tests; 4 need the V3 checkpoints + dev set
|
|
96
|
+
python3 v4/repro_v3.py # V3 reproduction check
|
|
97
|
+
python3 v4/collect_dev.py && python3 v4/eval_dev.py # validation set (about 3.4 GB in data/)
|
|
98
|
+
python3 v4/select_v4.py && python3 v4/tune_v4.py # selection and tuning (validation only)
|
|
99
|
+
WHICH=all python3 benchmark_v4.py # test benchmark, resumable (about 2 h on 2 cores)
|
|
100
|
+
python3 analyze_v4.py # tables, CIs, figures
|
|
101
|
+
python3 v4/latency_v4.py # timings (run on an idle CPU)
|
|
102
|
+
```
|
|
103
|
+
Runs are deterministic (fixed seeds, 1 thread per worker), and the V3 checkpoints in
|
|
104
|
+
`checkpoints/v3/` are reused unchanged. Rules that hold throughout:
|
|
105
|
+
- **Ground truth:** estimators never read it. Replay tests pass `s_true=None`.
|
|
106
|
+
- **Configuration:** stored in `results/v4/tuning_v4.json` (supervisor),
|
|
107
|
+
`checkpoints/v3/calibration.json` and `results/v3/tuning.json` (inherited).
|
|
108
|
+
|
|
109
|
+
## Integrity notes
|
|
110
|
+
- **Simulation only.** No real event sensor, actuator or timing hardware.
|
|
111
|
+
- **Latency.** The 1 kHz control loop is not a demonstrated 1 ms sensor-to-actuator
|
|
112
|
+
latency. See `docs/V4_LATENCY.md`.
|
|
113
|
+
- **No priority claims.** No prior-art search was done, and nothing here is claimed as
|
|
114
|
+
first of its kind, formally safe or patentable.
|
|
115
|
+
- **Sample size.** Stress and scenario cells use 16 trials and mostly one seed, so
|
|
116
|
+
differences below about 20 points there are not established.
|
|
File without changes
|
|
@@ -0,0 +1,84 @@
|
|
|
1
|
+
"""Operating conditions for the stress benchmark (Upgrade D).
|
|
2
|
+
|
|
3
|
+
One SensorConditions object describes the scene and both sensors, so the event
|
|
4
|
+
camera and the frame camera always see the same world. NOMINAL reproduces the
|
|
5
|
+
v2 simulation exactly.
|
|
6
|
+
|
|
7
|
+
Physical models (deliberately simple, documented so a reviewer can argue with them):
|
|
8
|
+
brightness scene illumination scale. Event pixels compute log(I + dark
|
|
9
|
+
current), so dim scenes compress contrast and fire fewer
|
|
10
|
+
events; background noise rises ~1/sqrt(b) and the photoreceptor
|
|
11
|
+
slows (time constant 50 us / b). Frame cameras auto-expose up to
|
|
12
|
+
8 ms (inside the estimator's training range), then add gain,
|
|
13
|
+
so read noise grows ~sqrt(4 ms / (b * exposure)).
|
|
14
|
+
hdr_gain a bright window behind the top of the pole (y > 0.30 m) is
|
|
15
|
+
hdr_gain times brighter than the room. Frame pixels exposed for
|
|
16
|
+
the room saturate there (clip at 1.0); event pixels are
|
|
17
|
+
logarithmic and keep the contrast.
|
|
18
|
+
ev_noise_mult multiplies event background activity (0.3 Hz/pixel nominal)
|
|
19
|
+
thresh_sigma per-pixel contrast-threshold mismatch (0.03 nominal)
|
|
20
|
+
drop_prob fraction of events lost in readout (bus saturation)
|
|
21
|
+
occlusion (t0, t1) or a tuple of such windows: a grey panel covers the upper pole
|
|
22
|
+
x_offset_m, scale camera miscalibration: the sensor images world x * scale + offset at the
|
|
23
|
+
pixel the estimators believe shows world x (V4 'miscalib' scenario)
|
|
24
|
+
"""
|
|
25
|
+
from dataclasses import dataclass, replace
|
|
26
|
+
import numpy as np
|
|
27
|
+
|
|
28
|
+
READ_NOISE = 0.02
|
|
29
|
+
BASE_EXPOSURE = 0.004
|
|
30
|
+
MAX_EXPOSURE = 0.008
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
@dataclass(frozen=True)
|
|
34
|
+
class SensorConditions:
|
|
35
|
+
name: str = "nominal"
|
|
36
|
+
brightness: float = 1.0
|
|
37
|
+
hdr_gain: float = 1.0
|
|
38
|
+
hdr_y: float = 0.30
|
|
39
|
+
ev_noise_mult: float = 1.0
|
|
40
|
+
thresh_sigma: float = 0.03
|
|
41
|
+
drop_prob: float = 0.0
|
|
42
|
+
occlusion: tuple = None
|
|
43
|
+
occ_box: tuple = (-0.30, 0.30, 0.30, 0.62) # world x0, x1, y0, y1 [m]
|
|
44
|
+
occ_value: float = 0.5
|
|
45
|
+
x_offset_m: float = 0.0 # camera mounted off-centre (estimators assume the nominal calibration)
|
|
46
|
+
scale: float = 1.0 # focal-length / pixel-scale error
|
|
47
|
+
|
|
48
|
+
@property
|
|
49
|
+
def is_nominal(self):
|
|
50
|
+
return self == NOMINAL
|
|
51
|
+
|
|
52
|
+
# ---- event sensor ---------------------------------------------------------
|
|
53
|
+
@property
|
|
54
|
+
def ev_noise_rate(self):
|
|
55
|
+
return 0.3 * self.ev_noise_mult / np.sqrt(self.brightness)
|
|
56
|
+
|
|
57
|
+
@property
|
|
58
|
+
def photoreceptor_tau(self):
|
|
59
|
+
return 50e-6 / self.brightness
|
|
60
|
+
|
|
61
|
+
# ---- frame sensor (also used for the event sensor's APS channel) ----------
|
|
62
|
+
def frame_exposure(self, base=BASE_EXPOSURE):
|
|
63
|
+
return float(min(base / self.brightness, max(base, MAX_EXPOSURE)))
|
|
64
|
+
|
|
65
|
+
def frame_noise(self, exposure):
|
|
66
|
+
return float(READ_NOISE * np.sqrt(BASE_EXPOSURE / (self.brightness * exposure)))
|
|
67
|
+
|
|
68
|
+
def frame_sensor(self, radiance_avg):
|
|
69
|
+
"""Time-averaged scene radiance (already x brightness) -> normalised pixel values.
|
|
70
|
+
Exposure + gain map a room-lit scene back to the nominal range; anything
|
|
71
|
+
brighter than full scale saturates. Noise is added by the caller."""
|
|
72
|
+
return np.clip(radiance_avg / self.brightness, 0.0, 1.0)
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
NOMINAL = SensorConditions()
|
|
76
|
+
PRESETS = {
|
|
77
|
+
"nominal": NOMINAL,
|
|
78
|
+
"low_light": SensorConditions("low_light", brightness=0.05),
|
|
79
|
+
"hdr": SensorConditions("hdr", hdr_gain=40.0),
|
|
80
|
+
"event_noise": SensorConditions("event_noise", ev_noise_mult=10.0),
|
|
81
|
+
"thresh_mismatch": SensorConditions("thresh_mismatch", thresh_sigma=0.09),
|
|
82
|
+
"dropped_events": SensorConditions("dropped_events", drop_prob=0.5),
|
|
83
|
+
"occlusion": SensorConditions("occlusion", occlusion=(0.21, 0.31)),
|
|
84
|
+
}
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
"""1 kHz state-feedback controller.
|
|
2
|
+
|
|
3
|
+
An infinite-horizon LQR is the unconstrained limit of linear MPC: it minimises
|
|
4
|
+
the same quadratic cost over an unbounded horizon and collapses to a single
|
|
5
|
+
gain matrix, so it costs a 4-element dot product per tick and runs comfortably
|
|
6
|
+
at 1 kHz even in Python. Actuator limits are enforced by saturation.
|
|
7
|
+
A constrained MPC can be dropped in through the same `act(state)` interface.
|
|
8
|
+
"""
|
|
9
|
+
import numpy as np
|
|
10
|
+
from .physics import PlantParams, linearise
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def dlqr(A, B, Q, R, iters=20000, tol=1e-10):
|
|
14
|
+
P = Q.copy()
|
|
15
|
+
for _ in range(iters):
|
|
16
|
+
BtP = B.T @ P
|
|
17
|
+
K = np.linalg.solve(R + BtP @ B, BtP @ A)
|
|
18
|
+
P_new = Q + A.T @ P @ (A - B @ K)
|
|
19
|
+
if np.max(np.abs(P_new - P)) < tol:
|
|
20
|
+
P = P_new
|
|
21
|
+
break
|
|
22
|
+
P = P_new
|
|
23
|
+
K = np.linalg.solve(R + B.T @ P @ B, B.T @ P @ A)
|
|
24
|
+
return K
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
class LQRController:
|
|
28
|
+
def __init__(self, p: PlantParams, dt=1e-3,
|
|
29
|
+
Q=(30.0, 2.0, 200.0, 5.0), R=0.02):
|
|
30
|
+
A, B = linearise(p, dt)
|
|
31
|
+
self.K = dlqr(A, B, np.diag(Q), np.array([[R]]))
|
|
32
|
+
self.u_max = p.u_max
|
|
33
|
+
|
|
34
|
+
def act(self, s):
|
|
35
|
+
u = -float((self.K @ np.asarray(s))[0])
|
|
36
|
+
return float(np.clip(u, -self.u_max, self.u_max))
|
|
@@ -0,0 +1,251 @@
|
|
|
1
|
+
"""Physics-informed event observer (PINN + selective state-space model).
|
|
2
|
+
|
|
3
|
+
Analogy: a blindfolded juggler who knows physics. Between glimpses, the juggler
|
|
4
|
+
predicts where the ball is from Newton's laws (the physics prior). Every event
|
|
5
|
+
spike is a tiny glimpse; the network only learns how much to *correct* the
|
|
6
|
+
physical prediction. When the scene is still and no events arrive, the
|
|
7
|
+
estimate simply keeps following the physics instead of going blind.
|
|
8
|
+
|
|
9
|
+
s_{k+1} = RK4(s_k, u_k) # differentiable physics
|
|
10
|
+
+ G * MLP([z_{k+1}, h_{k+1}, s_k]) # learned event correction
|
|
11
|
+
z_k = CNN(time_surface_k) # event features
|
|
12
|
+
h_{k+1} = a(z) * h_k + (1 - a(z)) * B z # selective SSM memory
|
|
13
|
+
a(z) = exp(-softplus(W_d z) * lambda) # input-dependent step (Mamba-style)
|
|
14
|
+
|
|
15
|
+
Training loss = state error + physics residuals (kinematic consistency and a
|
|
16
|
+
penalty on unexplained corrections) - i.e. the network is told that the
|
|
17
|
+
physics is mostly right and must justify any deviation with evidence.
|
|
18
|
+
"""
|
|
19
|
+
import torch
|
|
20
|
+
import torch.nn as nn
|
|
21
|
+
import torch.nn.functional as F
|
|
22
|
+
|
|
23
|
+
SCALE = torch.tensor([0.3, 1.0, 0.2, 2.0]) # typical magnitude of x, xd, th, thd
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class TorchCartPole(nn.Module):
|
|
27
|
+
def __init__(self, M=1.0, m=0.1, l=0.25, g=9.81, b_cart=0.1, b_pole=0.002):
|
|
28
|
+
super().__init__()
|
|
29
|
+
self.M, self.m, self.l, self.g = M, m, l, g
|
|
30
|
+
self.b_cart, self.b_pole = b_cart, b_pole
|
|
31
|
+
|
|
32
|
+
def deriv(self, s, u):
|
|
33
|
+
x, xd, th, thd = s.unbind(-1)
|
|
34
|
+
sin, cos = torch.sin(th), torch.cos(th)
|
|
35
|
+
total = self.M + self.m
|
|
36
|
+
temp = (u + self.m * self.l * thd ** 2 * sin - self.b_cart * xd) / total
|
|
37
|
+
thdd = (self.g * sin - cos * temp - self.b_pole * thd / (self.m * self.l)) / (
|
|
38
|
+
self.l * (4.0 / 3.0 - self.m * cos ** 2 / total))
|
|
39
|
+
xdd = temp - self.m * self.l * thdd * cos / total
|
|
40
|
+
return torch.stack([xd, xdd, thd, thdd], -1)
|
|
41
|
+
|
|
42
|
+
def step(self, s, u, dt):
|
|
43
|
+
k1 = self.deriv(s, u)
|
|
44
|
+
k2 = self.deriv(s + 0.5 * dt * k1, u)
|
|
45
|
+
k3 = self.deriv(s + 0.5 * dt * k2, u)
|
|
46
|
+
k4 = self.deriv(s + dt * k3, u)
|
|
47
|
+
return s + dt / 6.0 * (k1 + 2 * k2 + 2 * k3 + k4)
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
class SpatialSoftmaxKeypoints(nn.Module):
|
|
51
|
+
"""Full-resolution keypoint layer (sub-pixel precise).
|
|
52
|
+
|
|
53
|
+
Each of K learned heatmaps is turned into a probability map; its mean gives
|
|
54
|
+
a sub-pixel (x, y) location and its covariance gives the orientation and
|
|
55
|
+
spread of the blob - e.g. the angle of the pole. This preserves the fine
|
|
56
|
+
geometry that strided pooling throws away (1 px at the pole tip ~ 3 deg).
|
|
57
|
+
"""
|
|
58
|
+
|
|
59
|
+
def __init__(self, cin, H, W, K=12):
|
|
60
|
+
super().__init__()
|
|
61
|
+
self.c1 = nn.Conv2d(cin, 16, 3, 1, 1)
|
|
62
|
+
self.c2 = nn.Conv2d(16, K, 3, 1, 1)
|
|
63
|
+
self.log_temp = nn.Parameter(torch.zeros(K))
|
|
64
|
+
yy, xx = torch.meshgrid(torch.linspace(-1, 1, H), torch.linspace(-1, 1, W) * (W / H), indexing="ij")
|
|
65
|
+
self.register_buffer("gx", xx.flatten())
|
|
66
|
+
self.register_buffer("gy", yy.flatten())
|
|
67
|
+
self.out_dim = 5 * K
|
|
68
|
+
|
|
69
|
+
def forward(self, x):
|
|
70
|
+
logits = self.c2(F.gelu(self.c1(x))).flatten(2) * torch.exp(self.log_temp)[None, :, None]
|
|
71
|
+
p = torch.softmax(logits, -1)
|
|
72
|
+
mx, my = (p * self.gx).sum(-1), (p * self.gy).sum(-1)
|
|
73
|
+
vxx = (p * self.gx ** 2).sum(-1) - mx ** 2
|
|
74
|
+
vyy = (p * self.gy ** 2).sum(-1) - my ** 2
|
|
75
|
+
vxy = (p * self.gx * self.gy).sum(-1) - mx * my
|
|
76
|
+
return torch.cat([mx, my, 10 * vxx, 10 * vyy, 10 * vxy], -1)
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
class EventEncoder(nn.Module):
|
|
80
|
+
def __init__(self, H=32, W=96, d=128, cin=4):
|
|
81
|
+
super().__init__()
|
|
82
|
+
yy, xx = torch.meshgrid(torch.linspace(-1, 1, H), torch.linspace(-1, 1, W), indexing="ij")
|
|
83
|
+
self.register_buffer("coords", torch.stack([xx, yy])[None])
|
|
84
|
+
self.c1 = nn.Conv2d(cin + 2, 16, 5, 2, 2)
|
|
85
|
+
self.c2 = nn.Conv2d(16, 32, 3, 2, 1)
|
|
86
|
+
self.c3 = nn.Conv2d(32, 32, 3, 2, 1)
|
|
87
|
+
self.kp = SpatialSoftmaxKeypoints(cin, H, W)
|
|
88
|
+
self.fc = nn.Linear(32 * (H // 8) * (W // 8) + self.kp.out_dim, d)
|
|
89
|
+
|
|
90
|
+
def forward(self, ts): # ts: (N, 4, H, W): ON, OFF, slow-any, APS frame
|
|
91
|
+
x = torch.cat([ts, self.coords.expand(ts.shape[0], -1, -1, -1)], 1)
|
|
92
|
+
x = F.gelu(self.c1(x)); x = F.gelu(self.c2(x)); x = F.gelu(self.c3(x))
|
|
93
|
+
return F.gelu(self.fc(torch.cat([x.flatten(1), self.kp(ts)], 1)))
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
class EventObserver(nn.Module):
|
|
97
|
+
def __init__(self, dt=1e-3, d=128, n_state=64):
|
|
98
|
+
super().__init__()
|
|
99
|
+
self.dt = dt
|
|
100
|
+
self.physics = TorchCartPole()
|
|
101
|
+
self.enc = EventEncoder(d=d)
|
|
102
|
+
self.B = nn.Linear(d, n_state)
|
|
103
|
+
self.W_delta = nn.Linear(d, n_state)
|
|
104
|
+
self.log_lambda = nn.Parameter(torch.zeros(n_state))
|
|
105
|
+
self.head = nn.Sequential(nn.Linear(d + n_state + 4, 128), nn.GELU(),
|
|
106
|
+
nn.Linear(128, 64), nn.GELU(), nn.Linear(64, 4))
|
|
107
|
+
nn.init.zeros_(self.head[-1].weight); nn.init.zeros_(self.head[-1].bias)
|
|
108
|
+
self.register_buffer("scale", SCALE.clone())
|
|
109
|
+
self.n_state = n_state
|
|
110
|
+
|
|
111
|
+
def init_hidden(self, batch):
|
|
112
|
+
return torch.zeros(batch, self.n_state)
|
|
113
|
+
|
|
114
|
+
def cell(self, z, h, s, u):
|
|
115
|
+
"""One 1 ms observer tick. z: event features at t_{k+1}; s, u at t_k."""
|
|
116
|
+
a = torch.exp(-F.softplus(self.W_delta(z)) * torch.exp(self.log_lambda))
|
|
117
|
+
h = a * h + (1 - a) * self.B(z)
|
|
118
|
+
s_pred = self.physics.step(s, u, self.dt)
|
|
119
|
+
corr = self.head(torch.cat([z, h, s / self.scale], -1))
|
|
120
|
+
# correction expressed in units of "typical magnitude per 50 ms"
|
|
121
|
+
s_new = s_pred + corr * self.scale * (self.dt / 0.05)
|
|
122
|
+
return s_new, h, corr
|
|
123
|
+
|
|
124
|
+
def forward(self, ts_seq, u_seq, s0):
|
|
125
|
+
"""ts_seq: (B, T, 4, H, W) surfaces at t_1..t_T; u_seq: (B, T); s0: (B, 4)."""
|
|
126
|
+
Bsz, T = ts_seq.shape[:2]
|
|
127
|
+
z_all = self.enc(ts_seq.flatten(0, 1)).view(Bsz, T, -1)
|
|
128
|
+
h = self.init_hidden(Bsz)
|
|
129
|
+
s = s0
|
|
130
|
+
out, corrs = [], []
|
|
131
|
+
for k in range(T):
|
|
132
|
+
s, h, c = self.cell(z_all[:, k], h, s, u_seq[:, k])
|
|
133
|
+
out.append(s); corrs.append(c)
|
|
134
|
+
return torch.stack(out, 1), torch.stack(corrs, 1)
|
|
135
|
+
|
|
136
|
+
|
|
137
|
+
def physics_informed_loss(s_hat, s_true, corr, s0, dt, w_kin=0.05, w_corr=1e-3):
|
|
138
|
+
scale = SCALE.to(s_hat)
|
|
139
|
+
# tilt and tilt-rate errors matter most to the controller: weight them up
|
|
140
|
+
w = torch.tensor([1.0, 1.0, 4.0, 2.0]).to(s_hat)
|
|
141
|
+
data = ((((s_hat - s_true) / scale) ** 2) * w).mean() / w.mean()
|
|
142
|
+
# kinematic residual: position change must match estimated velocity
|
|
143
|
+
s_prev = torch.cat([s0[:, None], s_hat[:, :-1]], 1)
|
|
144
|
+
dpos = (s_hat[..., [0, 2]] - s_prev[..., [0, 2]]) / dt
|
|
145
|
+
vel = 0.5 * (s_hat[..., [1, 3]] + s_prev[..., [1, 3]])
|
|
146
|
+
kin = (((dpos - vel) / scale[[1, 3]]) ** 2).mean()
|
|
147
|
+
# parsimony: corrections must be justified by evidence
|
|
148
|
+
reg = (corr ** 2).mean()
|
|
149
|
+
return data + w_kin * kin + w_corr * reg, dict(data=data.item(), kin=kin.item(), reg=reg.item())
|
|
150
|
+
|
|
151
|
+
|
|
152
|
+
# =============================================================================
|
|
153
|
+
# v3: confidence-aware, gap-aware observer (Upgrades A + C)
|
|
154
|
+
# =============================================================================
|
|
155
|
+
class ProbEventObserver(EventObserver):
|
|
156
|
+
"""Event observer that also says how much to trust itself.
|
|
157
|
+
|
|
158
|
+
Analogy: a navigator who, besides plotting the ship's position, keeps an
|
|
159
|
+
'error circle' around it. In fog (event silence) or a storm (event burst,
|
|
160
|
+
occlusion) the circle grows; after a clear fix it shrinks. The controller
|
|
161
|
+
and safety shield read the circle, not just the dot.
|
|
162
|
+
|
|
163
|
+
Changes vs EventObserver:
|
|
164
|
+
* variance head - predicts log-variance of the state error after each
|
|
165
|
+
correction (heteroscedastic Gaussian, trained by NLL on
|
|
166
|
+
detached residuals so the mean is trained exactly as before)
|
|
167
|
+
* activity input - mean ON / OFF / slow time-surface activity, so the model
|
|
168
|
+
can tell silence and bursts apart explicitly
|
|
169
|
+
* gap-aware - a correction can be applied after `gap` ticks of
|
|
170
|
+
physics-only prediction (event-triggered scheduling);
|
|
171
|
+
trained with random tick skipping
|
|
172
|
+
* use_aps flag - event-only ablation zeroes the APS frame channel
|
|
173
|
+
"""
|
|
174
|
+
|
|
175
|
+
def __init__(self, dt=1e-3, d=128, n_state=64, use_aps=True):
|
|
176
|
+
super().__init__(dt, d, n_state)
|
|
177
|
+
self.head = nn.Sequential(nn.Linear(d + n_state + 4 + 1, 128), nn.GELU(),
|
|
178
|
+
nn.Linear(128, 64), nn.GELU(), nn.Linear(64, 4))
|
|
179
|
+
nn.init.zeros_(self.head[-1].weight); nn.init.zeros_(self.head[-1].bias)
|
|
180
|
+
self.var_head = nn.Sequential(nn.Linear(d + n_state + 4 + 1 + 3, 64), nn.GELU(),
|
|
181
|
+
nn.Linear(64, 4))
|
|
182
|
+
nn.init.zeros_(self.var_head[-1].weight)
|
|
183
|
+
nn.init.constant_(self.var_head[-1].bias, -3.0)
|
|
184
|
+
self.register_buffer("use_aps", torch.tensor(bool(use_aps)))
|
|
185
|
+
|
|
186
|
+
def encode(self, ts):
|
|
187
|
+
"""ts (N, 4, H, W) -> (features z, activity a)."""
|
|
188
|
+
if not bool(self.use_aps):
|
|
189
|
+
ts = torch.cat([ts[:, :3], torch.zeros_like(ts[:, 3:])], 1)
|
|
190
|
+
return self.enc(ts), ts[:, :3].mean((2, 3)) * 10.0
|
|
191
|
+
|
|
192
|
+
def physics_step(self, s, u):
|
|
193
|
+
return self.physics.step(s, u, self.dt)
|
|
194
|
+
|
|
195
|
+
def correct(self, z, act, h, s_pred, gap):
|
|
196
|
+
"""Apply one neural correction to the physics-predicted state s_pred.
|
|
197
|
+
gap: (N,) ticks since the previous correction (1 = dense operation)."""
|
|
198
|
+
a = torch.exp(-F.softplus(self.W_delta(z)) * torch.exp(self.log_lambda))
|
|
199
|
+
h = a * h + (1 - a) * self.B(z)
|
|
200
|
+
g = torch.log1p(gap.float())[:, None]
|
|
201
|
+
feat = torch.cat([z, h, s_pred / self.scale, g], -1)
|
|
202
|
+
corr = self.head(feat)
|
|
203
|
+
s_new = s_pred + corr * self.scale * (self.dt * gap.float()[:, None] / 0.05)
|
|
204
|
+
logvar = self.var_head(torch.cat([feat, act], -1)).clamp(-12, 4)
|
|
205
|
+
return s_new, h, corr, logvar
|
|
206
|
+
|
|
207
|
+
def forward(self, ts_seq, u_seq, s0, skip_p=None, gen=None):
|
|
208
|
+
"""Training/eval rollout. skip_p: (B,) probability of skipping the neural
|
|
209
|
+
correction at a tick (physics-only). Returns states, corrections,
|
|
210
|
+
log-variances and the correction mask."""
|
|
211
|
+
Bsz, T = ts_seq.shape[:2]
|
|
212
|
+
if skip_p is None:
|
|
213
|
+
mask = torch.ones(Bsz, T, dtype=torch.bool)
|
|
214
|
+
else:
|
|
215
|
+
mask = torch.rand(Bsz, T, generator=gen) >= skip_p[:, None]
|
|
216
|
+
mask[:, -1] = True
|
|
217
|
+
# encode only the ticks that receive a correction (identical result, less compute)
|
|
218
|
+
flat = mask.flatten()
|
|
219
|
+
sel = ts_seq.flatten(0, 1)[flat]
|
|
220
|
+
if torch.is_grad_enabled() and sel.shape[0] > 256:
|
|
221
|
+
# gradient checkpointing in chunks: activations are recomputed in the backward
|
|
222
|
+
# pass instead of stored (same maths; keeps long windows inside 8 GB RAM)
|
|
223
|
+
from torch.utils.checkpoint import checkpoint
|
|
224
|
+
parts = [checkpoint(self.encode, c, use_reentrant=False) for c in sel.split(256)]
|
|
225
|
+
z_sel = torch.cat([p[0] for p in parts]); act_sel = torch.cat([p[1] for p in parts])
|
|
226
|
+
else:
|
|
227
|
+
z_sel, act_sel = self.encode(sel)
|
|
228
|
+
z_all = z_sel.new_zeros(Bsz * T, z_sel.shape[-1]).index_put((flat.nonzero()[:, 0],), z_sel)
|
|
229
|
+
act_all = act_sel.new_zeros(Bsz * T, act_sel.shape[-1]).index_put((flat.nonzero()[:, 0],), act_sel)
|
|
230
|
+
z_all = z_all.view(Bsz, T, -1); act_all = act_all.view(Bsz, T, -1)
|
|
231
|
+
h = self.init_hidden(Bsz)
|
|
232
|
+
s = s0
|
|
233
|
+
gap = torch.zeros(Bsz)
|
|
234
|
+
out, corrs, lvs = [], [], []
|
|
235
|
+
for k in range(T):
|
|
236
|
+
s = self.physics_step(s, u_seq[:, k])
|
|
237
|
+
gap = gap + 1
|
|
238
|
+
s_c, h_c, c, lv = self.correct(z_all[:, k], act_all[:, k], h, s, gap)
|
|
239
|
+
m = mask[:, k][:, None]
|
|
240
|
+
s = torch.where(m, s_c, s)
|
|
241
|
+
h = torch.where(m, h_c, h)
|
|
242
|
+
gap = torch.where(mask[:, k], torch.zeros_like(gap), gap)
|
|
243
|
+
out.append(s); corrs.append(c * m); lvs.append(lv)
|
|
244
|
+
return torch.stack(out, 1), torch.stack(corrs, 1), torch.stack(lvs, 1), mask
|
|
245
|
+
|
|
246
|
+
|
|
247
|
+
def gaussian_nll(s_hat, s_true, logvar, mask):
|
|
248
|
+
"""Heteroscedastic Gaussian NLL in SCALE units on detached residuals (mean untouched)."""
|
|
249
|
+
e = ((s_hat.detach() - s_true) / SCALE.to(s_hat)) ** 2
|
|
250
|
+
nll = 0.5 * (e * torch.exp(-logvar) + logvar)
|
|
251
|
+
return (nll * mask[..., None]).sum() / (mask.sum() * 4).clamp(min=1)
|