vjepa-byor 0.1.0__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.
- vjepa_byor-0.1.0/LICENSE +21 -0
- vjepa_byor-0.1.0/PKG-INFO +203 -0
- vjepa_byor-0.1.0/README.md +181 -0
- vjepa_byor-0.1.0/pyproject.toml +35 -0
- vjepa_byor-0.1.0/pyproject.toml.orig +32 -0
- vjepa_byor-0.1.0/src/vjepa_byor/__init__.py +39 -0
- vjepa_byor-0.1.0/src/vjepa_byor/actions/__init__.py +19 -0
- vjepa_byor-0.1.0/src/vjepa_byor/actions/codecs/__init__.py +5 -0
- vjepa_byor-0.1.0/src/vjepa_byor/actions/codecs/json_line.py +69 -0
- vjepa_byor-0.1.0/src/vjepa_byor/actions/state.py +47 -0
- vjepa_byor-0.1.0/src/vjepa_byor/actions/vjepa_mapping.py +88 -0
- vjepa_byor-0.1.0/src/vjepa_byor/actions/vocabulary.py +84 -0
- vjepa_byor-0.1.0/src/vjepa_byor/cli.py +340 -0
- vjepa_byor-0.1.0/src/vjepa_byor/dashboard/__init__.py +11 -0
- vjepa_byor-0.1.0/src/vjepa_byor/dashboard/hub.py +211 -0
- vjepa_byor-0.1.0/src/vjepa_byor/dashboard/static/index.html +473 -0
- vjepa_byor-0.1.0/src/vjepa_byor/dashboard/web.py +196 -0
- vjepa_byor-0.1.0/src/vjepa_byor/data/__init__.py +40 -0
- vjepa_byor-0.1.0/src/vjepa_byor/data/collector.py +402 -0
- vjepa_byor-0.1.0/src/vjepa_byor/data/dataset.py +309 -0
- vjepa_byor-0.1.0/src/vjepa_byor/data/policy.py +140 -0
- vjepa_byor-0.1.0/src/vjepa_byor/data/sample.py +98 -0
- vjepa_byor-0.1.0/src/vjepa_byor/data/storage.py +291 -0
- vjepa_byor-0.1.0/src/vjepa_byor/interfaces.py +183 -0
- vjepa_byor-0.1.0/src/vjepa_byor/planning/__init__.py +46 -0
- vjepa_byor-0.1.0/src/vjepa_byor/planning/planner.py +268 -0
- vjepa_byor-0.1.0/src/vjepa_byor/planning/ranking.py +109 -0
- vjepa_byor-0.1.0/src/vjepa_byor/planning/scorer.py +145 -0
- vjepa_byor-0.1.0/src/vjepa_byor/predictor/__init__.py +58 -0
- vjepa_byor-0.1.0/src/vjepa_byor/predictor/encoding.py +63 -0
- vjepa_byor-0.1.0/src/vjepa_byor/predictor/evaluate.py +81 -0
- vjepa_byor-0.1.0/src/vjepa_byor/predictor/finetune.py +217 -0
- vjepa_byor-0.1.0/src/vjepa_byor/predictor/inference.py +122 -0
- vjepa_byor-0.1.0/src/vjepa_byor/predictor/input_layers.py +31 -0
- vjepa_byor-0.1.0/src/vjepa_byor/predictor/latent_cache.py +162 -0
- vjepa_byor-0.1.0/src/vjepa_byor/predictor/loading.py +115 -0
- vjepa_byor-0.1.0/src/vjepa_byor/predictor/lora.py +80 -0
- vjepa_byor-0.1.0/src/vjepa_byor/registry.py +120 -0
- vjepa_byor-0.1.0/src/vjepa_byor/rig.py +93 -0
- vjepa_byor-0.1.0/src/vjepa_byor/ros2/__init__.py +94 -0
- vjepa_byor-0.1.0/src/vjepa_byor/ros2/config.py +288 -0
- vjepa_byor-0.1.0/src/vjepa_byor/ros2/dashboard.py +156 -0
- vjepa_byor-0.1.0/src/vjepa_byor/ros2/link.py +327 -0
- vjepa_byor-0.1.0/src/vjepa_byor/ros2/reach_goal.py +253 -0
- vjepa_byor-0.1.0/src/vjepa_byor/sim/__init__.py +21 -0
- vjepa_byor-0.1.0/src/vjepa_byor/sim/camera.py +115 -0
- vjepa_byor-0.1.0/src/vjepa_byor/sim/rig.py +44 -0
- vjepa_byor-0.1.0/src/vjepa_byor/sim/ros2_bridge.py +169 -0
- vjepa_byor-0.1.0/src/vjepa_byor/sim/stage.py +68 -0
- vjepa_byor-0.1.0/src/vjepa_byor/testing/__init__.py +9 -0
- vjepa_byor-0.1.0/src/vjepa_byor/testing/collector_conformance.py +61 -0
- vjepa_byor-0.1.0/src/vjepa_byor/testing/conformance.py +208 -0
- vjepa_byor-0.1.0/src/vjepa_byor/testing/fake_camera.py +50 -0
- vjepa_byor-0.1.0/src/vjepa_byor/testing/fake_transport.py +91 -0
vjepa_byor-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Ampyard
|
|
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,203 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: vjepa-byor
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Bring Your Own Robot: fine-tune V-JEPA 2-AC to predict what your robot's actions will do.
|
|
5
|
+
License-Expression: MIT
|
|
6
|
+
License-File: LICENSE
|
|
7
|
+
Requires-Dist: pillow>=10
|
|
8
|
+
Requires-Dist: torch>=2.4
|
|
9
|
+
Requires-Dist: numpy ; extra == 'ros2'
|
|
10
|
+
Requires-Dist: pyyaml ; extra == 'ros2'
|
|
11
|
+
Requires-Dist: pytest>=8 ; extra == 'testing'
|
|
12
|
+
Requires-Dist: einops>=0.7 ; extra == 'vjepa'
|
|
13
|
+
Requires-Dist: timm>=1.0 ; extra == 'vjepa'
|
|
14
|
+
Requires-Python: >=3.11
|
|
15
|
+
Project-URL: Homepage, https://github.com/ampyard/vjepa-byor
|
|
16
|
+
Project-URL: Repository, https://github.com/ampyard/vjepa-byor
|
|
17
|
+
Project-URL: Issues, https://github.com/ampyard/vjepa-byor/issues
|
|
18
|
+
Provides-Extra: ros2
|
|
19
|
+
Provides-Extra: testing
|
|
20
|
+
Provides-Extra: vjepa
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
|
|
23
|
+
# vjepa-byor
|
|
24
|
+
|
|
25
|
+
**Bring Your Own Robot:** fine-tune V-JEPA 2-AC to predict what your robot's actions will do.
|
|
26
|
+
|
|
27
|
+
## What is a rig?
|
|
28
|
+
|
|
29
|
+
A **rig** is your robot setup: the actuators that move, the controller that drives them, and the
|
|
30
|
+
camera that watches the result. A desk-sized XY stage with stepper motors, an ESP32 and a webcam is a
|
|
31
|
+
rig, and so is a simulated one. "Robot" is the public-facing word; "rig" is the term used in the code
|
|
32
|
+
(`Rig`, `RigState`, `--rig`).
|
|
33
|
+
|
|
34
|
+
The library has no hardware code of its own. To use it with your robot, you describe your rig to it:
|
|
35
|
+
which actions it can take, what its position looks like, how to talk to it, and how to see it.
|
|
36
|
+
|
|
37
|
+
## Adding your rig
|
|
38
|
+
|
|
39
|
+
A rig plugin implements the protocols in `vjepa_byor.interfaces`, bundles them in a `vjepa_byor.Rig`,
|
|
40
|
+
and registers it under the `vjepa_byor.rigs` entry-point group:
|
|
41
|
+
|
|
42
|
+
```toml
|
|
43
|
+
[project.entry-points."vjepa_byor.rigs"]
|
|
44
|
+
my-rig = "my_rig_package:make_rig"
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
The rig's name must match the name it is registered under: `make_rig` above returns
|
|
48
|
+
`Rig(name="my-rig", ...)`, otherwise `load_rig("my-rig")` refuses it.
|
|
49
|
+
|
|
50
|
+
There are two ways to talk to the hardware:
|
|
51
|
+
|
|
52
|
+
- **Text commands** (serial, WiFi, Bluetooth): give the rig a `CommandCodec`, which turns an action into
|
|
53
|
+
a command line and the reply back into a state, and a `Transport` that carries the lines.
|
|
54
|
+
- **Action tensors** (ROS 2, a vendor SDK, ...): give the rig an `ActionTransport` and no codec. Its
|
|
55
|
+
`execute(action, timeout)` carries out the move and returns the state the rig reports.
|
|
56
|
+
|
|
57
|
+
A ROS 2 robot needs no plugin: the built-in `ros2` rig is described by a YAML file
|
|
58
|
+
(`--rig ros2 --rig-config robot.yaml`, see `vjepa_byor.ros2.config`). Install with
|
|
59
|
+
`pip install vjepa-byor[ros2]` in a venv created with `--system-site-packages` from a sourced ROS 2
|
|
60
|
+
environment, since `rclpy` is not on PyPI. Without `rclpy`, `available_rigs()` does not list `ros2`.
|
|
61
|
+
|
|
62
|
+
Validate it without hardware by subclassing `vjepa_byor.testing.conformance.RigConformanceTests`
|
|
63
|
+
(install with `pip install vjepa-byor[testing]`) and backing the rig with a
|
|
64
|
+
`vjepa_byor.testing.FakeTransport` (or `FakeActionTransport`) that answers like your device.
|
|
65
|
+
|
|
66
|
+
## Trying it without hardware
|
|
67
|
+
|
|
68
|
+
The library includes `sim`, a simulated XY stage with a rendered top-down camera, so collection,
|
|
69
|
+
fine-tuning and planning run end to end with nothing attached. The examples below use it; with your own
|
|
70
|
+
rig, use its registered name in place of `sim`.
|
|
71
|
+
|
|
72
|
+
```
|
|
73
|
+
pip install vjepa-byor
|
|
74
|
+
vjepa-byor collect --rig sim --out data/sim-try --num-samples 200 --settle-time 0
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
- **Actions:** 8 compass directions × 3 step sizes (1, 5 and 20 mm by default), as `[dx_mm, dy_mm]`
|
|
78
|
+
- **State:** `XYState(x_mm, y_mm)`, reported after every move; travel is 200 × 150 mm, starting in the middle
|
|
79
|
+
- **V-JEPA mapping:** millimetres to metres in the end effector's x and y slots: `[dx, dy, 0, 0, 0, 0, 0]`
|
|
80
|
+
- **Wire format:** the JSON-line codec over an in-memory transport; out-of-range moves are refused with
|
|
81
|
+
`{"error": "out of range"}` and the carriage stays put
|
|
82
|
+
- **Camera:** 256 × 256 RGB `uint8` by default: an orange puck on a gridded mat, with a red and a blue
|
|
83
|
+
landmark block. Anti-aliased, so a 1 mm move changes the photo; deterministic, so a position always
|
|
84
|
+
gives the same photo
|
|
85
|
+
|
|
86
|
+
Options (`--rig-config sim.json`): `step_sizes_mm`, `travel_mm`, `start_mm`, `resolution` (`[height, width]`),
|
|
87
|
+
`reply_timeout`.
|
|
88
|
+
|
|
89
|
+
```python
|
|
90
|
+
from vjepa_byor import load_rig
|
|
91
|
+
|
|
92
|
+
with load_rig("sim") as rig:
|
|
93
|
+
print(rig.execute(rig.action_space.by_label("NE 5mm"))) # XYState(x_mm=105.0, y_mm=80.0)
|
|
94
|
+
photo = rig.camera.capture() # uint8 (3, 256, 256)
|
|
95
|
+
```
|
|
96
|
+
|
|
97
|
+
It is built only from the library's public building blocks (`vjepa_byor.sim`), so it doubles as an
|
|
98
|
+
example of a rig.
|
|
99
|
+
|
|
100
|
+
### The sim as a ROS 2 robot
|
|
101
|
+
|
|
102
|
+
`python -m vjepa_byor.sim.ros2_bridge` runs the sim as an ordinary ROS 2 node (needs a sourced ROS 2
|
|
103
|
+
environment and `vjepa-byor[ros2]`):
|
|
104
|
+
|
|
105
|
+
```bash
|
|
106
|
+
python -m vjepa_byor.sim.ros2_bridge --ros-args -p start_mm:="[40.0, 100.0]"
|
|
107
|
+
ros2 topic pub --once /rig/cmd geometry_msgs/msg/Vector3Stamped "{vector: {x: 0.02, y: -0.02}}"
|
|
108
|
+
ros2 topic echo /rig/state
|
|
109
|
+
```
|
|
110
|
+
|
|
111
|
+
| Name | Type | |
|
|
112
|
+
|---|---|---|
|
|
113
|
+
| `/camera/image_raw` | `sensor_msgs/Image` | `rgb8`, 256 × 256 |
|
|
114
|
+
| `/rig/state` | `sensor_msgs/JointState` | joints `x`, `y`, positions in metres |
|
|
115
|
+
| `/rig/cmd` (subscribed) | `geometry_msgs/Vector3Stamped` | relative move, `vector.x` / `vector.y` in metres |
|
|
116
|
+
| `/rig/reset` (service) | `std_srvs/Trigger` | carriage back to `start_mm` |
|
|
117
|
+
|
|
118
|
+
Photo and state share one `header.stamp` and go out at `rate_hz` (10) and after every command, including
|
|
119
|
+
a refused one, which leaves the position unchanged. Parameters: `travel_mm`, `start_mm`, `resolution`,
|
|
120
|
+
`rate_hz`, `camera_frame`, `stage_frame`.
|
|
121
|
+
|
|
122
|
+
## Collecting a fine-tuning dataset
|
|
123
|
+
|
|
124
|
+
`vjepa-byor collect` drives any registered rig through short episodes of random actions and records every
|
|
125
|
+
transition as (before photo, rig state, action, after photo):
|
|
126
|
+
|
|
127
|
+
```
|
|
128
|
+
vjepa-byor collect --rig sim --out data/sim-run1 \
|
|
129
|
+
--num-samples 1000 --steps-per-episode 20 --seed 0 --settle-time 0.5
|
|
130
|
+
vjepa-byor collect --rig sim --out data/sim-run1 --num-samples 2000 --resume # adds 1000 more
|
|
131
|
+
vjepa-byor collect --verify data/sim-run1
|
|
132
|
+
```
|
|
133
|
+
|
|
134
|
+
`--num-samples` is the target size of the whole dataset, so a resumed run only collects what is missing.
|
|
135
|
+
`--rig-config rig.json` passes options to the rig's factory, and `--state-bounds 0:200,0:150` tells the
|
|
136
|
+
random walk the limits of the state tensor so it steers away from them (without bounds it learns from
|
|
137
|
+
rejections). Ctrl-C finishes the step in flight and stops cleanly; a second Ctrl-C aborts.
|
|
138
|
+
|
|
139
|
+
The rig needs a camera (`Rig.camera`). Since the interfaces have no "home" or "read position" command,
|
|
140
|
+
the collector learns where the rig is from its reply to a first, unrecorded move.
|
|
141
|
+
|
|
142
|
+
### Dataset layout
|
|
143
|
+
|
|
144
|
+
```
|
|
145
|
+
<dataset>/
|
|
146
|
+
├── meta.json # schema_version, rig, action vocabulary, resolution [height, width], ...
|
|
147
|
+
├── manifest.jsonl # one sample per line, append-only
|
|
148
|
+
├── events.jsonl # rejections, failures, syncs, resumptions
|
|
149
|
+
├── splits.json # episode -> train / val / test (created on first read)
|
|
150
|
+
└── frames/ep000000/000.png, 001.png, ... # lossless PNGs exactly as captured
|
|
151
|
+
```
|
|
152
|
+
|
|
153
|
+
Step *k* of an episode has frame *k* as its before photo and frame *k+1* as its after photo, so photos are
|
|
154
|
+
stored once. The full format contract is in the docstring of `vjepa_byor.data.storage`.
|
|
155
|
+
|
|
156
|
+
Each manifest line is a sample:
|
|
157
|
+
|
|
158
|
+
| Field | Meaning |
|
|
159
|
+
|---|---|
|
|
160
|
+
| `sample_id`, `episode_id`, `step_index` | e.g. `ep000012_s0003`, `ep000012`, `3` |
|
|
161
|
+
| `before_frame`, `after_frame` | PNG paths relative to the dataset root |
|
|
162
|
+
| `state_before`, `state_after` | the rig's reported state, as `RigState.to_dict()` |
|
|
163
|
+
| `state_before_tensor`, `state_after_tensor` | the same, as `RigState.as_tensor()` |
|
|
164
|
+
| `action_index`, `action_name`, `action` | index into the vocabulary, its label, the action tensor |
|
|
165
|
+
| `command` | the wire line sent to the rig |
|
|
166
|
+
| `status` | `ok`, or `rejected` if the rig reported an error (the state is then unchanged) |
|
|
167
|
+
| `changed` | `state_after != state_before`, from the reported states |
|
|
168
|
+
| `resolution`, `timestamp`, `error` | `[height, width]`, ISO 8601 UTC, the rig's error message if rejected |
|
|
169
|
+
|
|
170
|
+
Read it back with `vjepa_byor.data.CollectedDataset(root, split="train", drop_no_change=False)`; splits are by
|
|
171
|
+
episode and never change once written. `vjepa_byor.data.export(root, out, split="train")` copies a split into
|
|
172
|
+
a new self-contained dataset.
|
|
173
|
+
|
|
174
|
+
Plugin authors can check the collector drives their rig by adding
|
|
175
|
+
`vjepa_byor.testing.collector_conformance.CollectorConformanceTests` next to `RigConformanceTests`; it needs
|
|
176
|
+
`rig.camera`, and `vjepa_byor.testing.fake_camera.MarkerCamera` can draw the emulated position.
|
|
177
|
+
|
|
178
|
+
## Fine-tuning the predictor
|
|
179
|
+
|
|
180
|
+
Only V-JEPA 2-AC's action-conditioned predictor is fine-tuned; the ViT-g encoder stays frozen. Install the
|
|
181
|
+
model code's dependencies with `pip install vjepa-byor[vjepa]`. The weights are not bundled (they have their
|
|
182
|
+
own license): download `vjepa2-ac-vitg.pt` (about 11 GB) from `https://dl.fbaipublicfiles.com/vjepa2/` to
|
|
183
|
+
`~/.cache/torch/hub/checkpoints/`, or pass `--checkpoint PATH`.
|
|
184
|
+
|
|
185
|
+
```
|
|
186
|
+
vjepa-byor encode --data data/sim-run1 # frozen encoder, once per frame
|
|
187
|
+
vjepa-byor finetune --data data/sim-run1 --out checkpoints/sim # train split; best epoch on val is kept
|
|
188
|
+
vjepa-byor eval --data data/sim-run1 --finetuned checkpoints/sim --split test
|
|
189
|
+
```
|
|
190
|
+
|
|
191
|
+
- **encode** writes `<dataset>/latents/latents.pt`: every frame's layer-normalised latents (256 tokens × 1408,
|
|
192
|
+
bf16). A rerun after `collect --resume` encodes only the new frames.
|
|
193
|
+
- **finetune** trains on the latents, so the encoder is never in GPU memory. The objective is upstream's
|
|
194
|
+
teacher-forced step: before latents + mapped action + mapped state → after latents, L1 loss. bf16 autocast,
|
|
195
|
+
AdamW, warmup + cosine, activation checkpointing. `--lora-rank 16` trains LoRA adapters instead of every
|
|
196
|
+
weight; `--max-vram-gb 8` caps GPU memory to check a run fits an 8 GB card. The rig's `VJepaMapping` comes
|
|
197
|
+
from the rig the dataset was collected with (`--rig` / `--rig-config` to override); a mapping that is not
|
|
198
|
+
7-D gets fresh, custom-sized action and state input layers.
|
|
199
|
+
- **eval** reports the mean L1 error of the fine-tuned predictor, the original one and a "no change" baseline
|
|
200
|
+
(the before latents), over all samples and over those where the rig moved.
|
|
201
|
+
|
|
202
|
+
In code, `vjepa_byor.predictor.ActionPredictor(encoder, predictor, rig.mapping).predict_basket(photo, state,
|
|
203
|
+
rig.action_space)` predicts the next-state latents for every action in one batch.
|
|
@@ -0,0 +1,181 @@
|
|
|
1
|
+
# vjepa-byor
|
|
2
|
+
|
|
3
|
+
**Bring Your Own Robot:** fine-tune V-JEPA 2-AC to predict what your robot's actions will do.
|
|
4
|
+
|
|
5
|
+
## What is a rig?
|
|
6
|
+
|
|
7
|
+
A **rig** is your robot setup: the actuators that move, the controller that drives them, and the
|
|
8
|
+
camera that watches the result. A desk-sized XY stage with stepper motors, an ESP32 and a webcam is a
|
|
9
|
+
rig, and so is a simulated one. "Robot" is the public-facing word; "rig" is the term used in the code
|
|
10
|
+
(`Rig`, `RigState`, `--rig`).
|
|
11
|
+
|
|
12
|
+
The library has no hardware code of its own. To use it with your robot, you describe your rig to it:
|
|
13
|
+
which actions it can take, what its position looks like, how to talk to it, and how to see it.
|
|
14
|
+
|
|
15
|
+
## Adding your rig
|
|
16
|
+
|
|
17
|
+
A rig plugin implements the protocols in `vjepa_byor.interfaces`, bundles them in a `vjepa_byor.Rig`,
|
|
18
|
+
and registers it under the `vjepa_byor.rigs` entry-point group:
|
|
19
|
+
|
|
20
|
+
```toml
|
|
21
|
+
[project.entry-points."vjepa_byor.rigs"]
|
|
22
|
+
my-rig = "my_rig_package:make_rig"
|
|
23
|
+
```
|
|
24
|
+
|
|
25
|
+
The rig's name must match the name it is registered under: `make_rig` above returns
|
|
26
|
+
`Rig(name="my-rig", ...)`, otherwise `load_rig("my-rig")` refuses it.
|
|
27
|
+
|
|
28
|
+
There are two ways to talk to the hardware:
|
|
29
|
+
|
|
30
|
+
- **Text commands** (serial, WiFi, Bluetooth): give the rig a `CommandCodec`, which turns an action into
|
|
31
|
+
a command line and the reply back into a state, and a `Transport` that carries the lines.
|
|
32
|
+
- **Action tensors** (ROS 2, a vendor SDK, ...): give the rig an `ActionTransport` and no codec. Its
|
|
33
|
+
`execute(action, timeout)` carries out the move and returns the state the rig reports.
|
|
34
|
+
|
|
35
|
+
A ROS 2 robot needs no plugin: the built-in `ros2` rig is described by a YAML file
|
|
36
|
+
(`--rig ros2 --rig-config robot.yaml`, see `vjepa_byor.ros2.config`). Install with
|
|
37
|
+
`pip install vjepa-byor[ros2]` in a venv created with `--system-site-packages` from a sourced ROS 2
|
|
38
|
+
environment, since `rclpy` is not on PyPI. Without `rclpy`, `available_rigs()` does not list `ros2`.
|
|
39
|
+
|
|
40
|
+
Validate it without hardware by subclassing `vjepa_byor.testing.conformance.RigConformanceTests`
|
|
41
|
+
(install with `pip install vjepa-byor[testing]`) and backing the rig with a
|
|
42
|
+
`vjepa_byor.testing.FakeTransport` (or `FakeActionTransport`) that answers like your device.
|
|
43
|
+
|
|
44
|
+
## Trying it without hardware
|
|
45
|
+
|
|
46
|
+
The library includes `sim`, a simulated XY stage with a rendered top-down camera, so collection,
|
|
47
|
+
fine-tuning and planning run end to end with nothing attached. The examples below use it; with your own
|
|
48
|
+
rig, use its registered name in place of `sim`.
|
|
49
|
+
|
|
50
|
+
```
|
|
51
|
+
pip install vjepa-byor
|
|
52
|
+
vjepa-byor collect --rig sim --out data/sim-try --num-samples 200 --settle-time 0
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
- **Actions:** 8 compass directions × 3 step sizes (1, 5 and 20 mm by default), as `[dx_mm, dy_mm]`
|
|
56
|
+
- **State:** `XYState(x_mm, y_mm)`, reported after every move; travel is 200 × 150 mm, starting in the middle
|
|
57
|
+
- **V-JEPA mapping:** millimetres to metres in the end effector's x and y slots: `[dx, dy, 0, 0, 0, 0, 0]`
|
|
58
|
+
- **Wire format:** the JSON-line codec over an in-memory transport; out-of-range moves are refused with
|
|
59
|
+
`{"error": "out of range"}` and the carriage stays put
|
|
60
|
+
- **Camera:** 256 × 256 RGB `uint8` by default: an orange puck on a gridded mat, with a red and a blue
|
|
61
|
+
landmark block. Anti-aliased, so a 1 mm move changes the photo; deterministic, so a position always
|
|
62
|
+
gives the same photo
|
|
63
|
+
|
|
64
|
+
Options (`--rig-config sim.json`): `step_sizes_mm`, `travel_mm`, `start_mm`, `resolution` (`[height, width]`),
|
|
65
|
+
`reply_timeout`.
|
|
66
|
+
|
|
67
|
+
```python
|
|
68
|
+
from vjepa_byor import load_rig
|
|
69
|
+
|
|
70
|
+
with load_rig("sim") as rig:
|
|
71
|
+
print(rig.execute(rig.action_space.by_label("NE 5mm"))) # XYState(x_mm=105.0, y_mm=80.0)
|
|
72
|
+
photo = rig.camera.capture() # uint8 (3, 256, 256)
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
It is built only from the library's public building blocks (`vjepa_byor.sim`), so it doubles as an
|
|
76
|
+
example of a rig.
|
|
77
|
+
|
|
78
|
+
### The sim as a ROS 2 robot
|
|
79
|
+
|
|
80
|
+
`python -m vjepa_byor.sim.ros2_bridge` runs the sim as an ordinary ROS 2 node (needs a sourced ROS 2
|
|
81
|
+
environment and `vjepa-byor[ros2]`):
|
|
82
|
+
|
|
83
|
+
```bash
|
|
84
|
+
python -m vjepa_byor.sim.ros2_bridge --ros-args -p start_mm:="[40.0, 100.0]"
|
|
85
|
+
ros2 topic pub --once /rig/cmd geometry_msgs/msg/Vector3Stamped "{vector: {x: 0.02, y: -0.02}}"
|
|
86
|
+
ros2 topic echo /rig/state
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
| Name | Type | |
|
|
90
|
+
|---|---|---|
|
|
91
|
+
| `/camera/image_raw` | `sensor_msgs/Image` | `rgb8`, 256 × 256 |
|
|
92
|
+
| `/rig/state` | `sensor_msgs/JointState` | joints `x`, `y`, positions in metres |
|
|
93
|
+
| `/rig/cmd` (subscribed) | `geometry_msgs/Vector3Stamped` | relative move, `vector.x` / `vector.y` in metres |
|
|
94
|
+
| `/rig/reset` (service) | `std_srvs/Trigger` | carriage back to `start_mm` |
|
|
95
|
+
|
|
96
|
+
Photo and state share one `header.stamp` and go out at `rate_hz` (10) and after every command, including
|
|
97
|
+
a refused one, which leaves the position unchanged. Parameters: `travel_mm`, `start_mm`, `resolution`,
|
|
98
|
+
`rate_hz`, `camera_frame`, `stage_frame`.
|
|
99
|
+
|
|
100
|
+
## Collecting a fine-tuning dataset
|
|
101
|
+
|
|
102
|
+
`vjepa-byor collect` drives any registered rig through short episodes of random actions and records every
|
|
103
|
+
transition as (before photo, rig state, action, after photo):
|
|
104
|
+
|
|
105
|
+
```
|
|
106
|
+
vjepa-byor collect --rig sim --out data/sim-run1 \
|
|
107
|
+
--num-samples 1000 --steps-per-episode 20 --seed 0 --settle-time 0.5
|
|
108
|
+
vjepa-byor collect --rig sim --out data/sim-run1 --num-samples 2000 --resume # adds 1000 more
|
|
109
|
+
vjepa-byor collect --verify data/sim-run1
|
|
110
|
+
```
|
|
111
|
+
|
|
112
|
+
`--num-samples` is the target size of the whole dataset, so a resumed run only collects what is missing.
|
|
113
|
+
`--rig-config rig.json` passes options to the rig's factory, and `--state-bounds 0:200,0:150` tells the
|
|
114
|
+
random walk the limits of the state tensor so it steers away from them (without bounds it learns from
|
|
115
|
+
rejections). Ctrl-C finishes the step in flight and stops cleanly; a second Ctrl-C aborts.
|
|
116
|
+
|
|
117
|
+
The rig needs a camera (`Rig.camera`). Since the interfaces have no "home" or "read position" command,
|
|
118
|
+
the collector learns where the rig is from its reply to a first, unrecorded move.
|
|
119
|
+
|
|
120
|
+
### Dataset layout
|
|
121
|
+
|
|
122
|
+
```
|
|
123
|
+
<dataset>/
|
|
124
|
+
├── meta.json # schema_version, rig, action vocabulary, resolution [height, width], ...
|
|
125
|
+
├── manifest.jsonl # one sample per line, append-only
|
|
126
|
+
├── events.jsonl # rejections, failures, syncs, resumptions
|
|
127
|
+
├── splits.json # episode -> train / val / test (created on first read)
|
|
128
|
+
└── frames/ep000000/000.png, 001.png, ... # lossless PNGs exactly as captured
|
|
129
|
+
```
|
|
130
|
+
|
|
131
|
+
Step *k* of an episode has frame *k* as its before photo and frame *k+1* as its after photo, so photos are
|
|
132
|
+
stored once. The full format contract is in the docstring of `vjepa_byor.data.storage`.
|
|
133
|
+
|
|
134
|
+
Each manifest line is a sample:
|
|
135
|
+
|
|
136
|
+
| Field | Meaning |
|
|
137
|
+
|---|---|
|
|
138
|
+
| `sample_id`, `episode_id`, `step_index` | e.g. `ep000012_s0003`, `ep000012`, `3` |
|
|
139
|
+
| `before_frame`, `after_frame` | PNG paths relative to the dataset root |
|
|
140
|
+
| `state_before`, `state_after` | the rig's reported state, as `RigState.to_dict()` |
|
|
141
|
+
| `state_before_tensor`, `state_after_tensor` | the same, as `RigState.as_tensor()` |
|
|
142
|
+
| `action_index`, `action_name`, `action` | index into the vocabulary, its label, the action tensor |
|
|
143
|
+
| `command` | the wire line sent to the rig |
|
|
144
|
+
| `status` | `ok`, or `rejected` if the rig reported an error (the state is then unchanged) |
|
|
145
|
+
| `changed` | `state_after != state_before`, from the reported states |
|
|
146
|
+
| `resolution`, `timestamp`, `error` | `[height, width]`, ISO 8601 UTC, the rig's error message if rejected |
|
|
147
|
+
|
|
148
|
+
Read it back with `vjepa_byor.data.CollectedDataset(root, split="train", drop_no_change=False)`; splits are by
|
|
149
|
+
episode and never change once written. `vjepa_byor.data.export(root, out, split="train")` copies a split into
|
|
150
|
+
a new self-contained dataset.
|
|
151
|
+
|
|
152
|
+
Plugin authors can check the collector drives their rig by adding
|
|
153
|
+
`vjepa_byor.testing.collector_conformance.CollectorConformanceTests` next to `RigConformanceTests`; it needs
|
|
154
|
+
`rig.camera`, and `vjepa_byor.testing.fake_camera.MarkerCamera` can draw the emulated position.
|
|
155
|
+
|
|
156
|
+
## Fine-tuning the predictor
|
|
157
|
+
|
|
158
|
+
Only V-JEPA 2-AC's action-conditioned predictor is fine-tuned; the ViT-g encoder stays frozen. Install the
|
|
159
|
+
model code's dependencies with `pip install vjepa-byor[vjepa]`. The weights are not bundled (they have their
|
|
160
|
+
own license): download `vjepa2-ac-vitg.pt` (about 11 GB) from `https://dl.fbaipublicfiles.com/vjepa2/` to
|
|
161
|
+
`~/.cache/torch/hub/checkpoints/`, or pass `--checkpoint PATH`.
|
|
162
|
+
|
|
163
|
+
```
|
|
164
|
+
vjepa-byor encode --data data/sim-run1 # frozen encoder, once per frame
|
|
165
|
+
vjepa-byor finetune --data data/sim-run1 --out checkpoints/sim # train split; best epoch on val is kept
|
|
166
|
+
vjepa-byor eval --data data/sim-run1 --finetuned checkpoints/sim --split test
|
|
167
|
+
```
|
|
168
|
+
|
|
169
|
+
- **encode** writes `<dataset>/latents/latents.pt`: every frame's layer-normalised latents (256 tokens × 1408,
|
|
170
|
+
bf16). A rerun after `collect --resume` encodes only the new frames.
|
|
171
|
+
- **finetune** trains on the latents, so the encoder is never in GPU memory. The objective is upstream's
|
|
172
|
+
teacher-forced step: before latents + mapped action + mapped state → after latents, L1 loss. bf16 autocast,
|
|
173
|
+
AdamW, warmup + cosine, activation checkpointing. `--lora-rank 16` trains LoRA adapters instead of every
|
|
174
|
+
weight; `--max-vram-gb 8` caps GPU memory to check a run fits an 8 GB card. The rig's `VJepaMapping` comes
|
|
175
|
+
from the rig the dataset was collected with (`--rig` / `--rig-config` to override); a mapping that is not
|
|
176
|
+
7-D gets fresh, custom-sized action and state input layers.
|
|
177
|
+
- **eval** reports the mean L1 error of the fine-tuned predictor, the original one and a "no change" baseline
|
|
178
|
+
(the before latents), over all samples and over those where the rig moved.
|
|
179
|
+
|
|
180
|
+
In code, `vjepa_byor.predictor.ActionPredictor(encoder, predictor, rig.mapping).predict_basket(photo, state,
|
|
181
|
+
rig.action_space)` predicts the next-state latents for every action in one batch.
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "vjepa-byor"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Bring Your Own Robot: fine-tune V-JEPA 2-AC to predict what your robot's actions will do."
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
license = "MIT"
|
|
7
|
+
license-files = ["LICENSE"]
|
|
8
|
+
requires-python = ">=3.11"
|
|
9
|
+
dependencies = [
|
|
10
|
+
"pillow>=10",
|
|
11
|
+
"torch>=2.4",
|
|
12
|
+
]
|
|
13
|
+
|
|
14
|
+
[project.scripts]
|
|
15
|
+
vjepa-byor = "vjepa_byor.cli:main"
|
|
16
|
+
|
|
17
|
+
[project.optional-dependencies]
|
|
18
|
+
vjepa = [
|
|
19
|
+
"einops>=0.7",
|
|
20
|
+
"timm>=1.0",
|
|
21
|
+
]
|
|
22
|
+
testing = ["pytest>=8"]
|
|
23
|
+
ros2 = [
|
|
24
|
+
"numpy",
|
|
25
|
+
"pyyaml",
|
|
26
|
+
]
|
|
27
|
+
|
|
28
|
+
[project.urls]
|
|
29
|
+
Homepage = "https://github.com/ampyard/vjepa-byor"
|
|
30
|
+
Repository = "https://github.com/ampyard/vjepa-byor"
|
|
31
|
+
Issues = "https://github.com/ampyard/vjepa-byor/issues"
|
|
32
|
+
|
|
33
|
+
[build-system]
|
|
34
|
+
requires = ["uv_build>=0.8,<0.14"]
|
|
35
|
+
build-backend = "uv_build"
|
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "vjepa-byor"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Bring Your Own Robot: fine-tune V-JEPA 2-AC to predict what your robot's actions will do."
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
license = "MIT"
|
|
7
|
+
license-files = ["LICENSE"]
|
|
8
|
+
requires-python = ">=3.11"
|
|
9
|
+
dependencies = [
|
|
10
|
+
"pillow>=10",
|
|
11
|
+
"torch>=2.4",
|
|
12
|
+
]
|
|
13
|
+
|
|
14
|
+
[project.scripts]
|
|
15
|
+
vjepa-byor = "vjepa_byor.cli:main"
|
|
16
|
+
|
|
17
|
+
[project.optional-dependencies]
|
|
18
|
+
# Meta's V-JEPA 2-AC model code (loaded through torch.hub) needs these
|
|
19
|
+
vjepa = ["einops>=0.7", "timm>=1.0"]
|
|
20
|
+
testing = ["pytest>=8"]
|
|
21
|
+
# The built-in ros2 rig and ROS 2 tools: their pip-installable dependencies. rclpy and the message
|
|
22
|
+
# packages come from a sourced ROS 2 environment (not PyPI); unpinned so ROS's own numpy satisfies it
|
|
23
|
+
ros2 = ["numpy", "pyyaml"]
|
|
24
|
+
|
|
25
|
+
[project.urls]
|
|
26
|
+
Homepage = "https://github.com/ampyard/vjepa-byor"
|
|
27
|
+
Repository = "https://github.com/ampyard/vjepa-byor"
|
|
28
|
+
Issues = "https://github.com/ampyard/vjepa-byor/issues"
|
|
29
|
+
|
|
30
|
+
[build-system]
|
|
31
|
+
requires = ["uv_build>=0.8,<0.14"]
|
|
32
|
+
build-backend = "uv_build"
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
"""vjepa-byor — Bring Your Own Robot: fine-tune V-JEPA 2-AC to predict what your robot's actions will do."""
|
|
2
|
+
|
|
3
|
+
from vjepa_byor.interfaces import (
|
|
4
|
+
ActionSpace,
|
|
5
|
+
ActionTransport,
|
|
6
|
+
Camera,
|
|
7
|
+
CodecError,
|
|
8
|
+
CommandCodec,
|
|
9
|
+
RigError,
|
|
10
|
+
RigState,
|
|
11
|
+
Transport,
|
|
12
|
+
VJepaMapping,
|
|
13
|
+
)
|
|
14
|
+
from vjepa_byor.registry import (
|
|
15
|
+
RigNotFoundError,
|
|
16
|
+
available_rigs,
|
|
17
|
+
load_rig,
|
|
18
|
+
register_rig,
|
|
19
|
+
unregister_rig,
|
|
20
|
+
)
|
|
21
|
+
from vjepa_byor.rig import Rig
|
|
22
|
+
|
|
23
|
+
__all__ = [
|
|
24
|
+
"ActionSpace",
|
|
25
|
+
"ActionTransport",
|
|
26
|
+
"Camera",
|
|
27
|
+
"CodecError",
|
|
28
|
+
"CommandCodec",
|
|
29
|
+
"Rig",
|
|
30
|
+
"RigError",
|
|
31
|
+
"RigNotFoundError",
|
|
32
|
+
"RigState",
|
|
33
|
+
"Transport",
|
|
34
|
+
"VJepaMapping",
|
|
35
|
+
"available_rigs",
|
|
36
|
+
"load_rig",
|
|
37
|
+
"register_rig",
|
|
38
|
+
"unregister_rig",
|
|
39
|
+
]
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
"""Generic building blocks for a rig's actions, state, V-JEPA mapping and wire format.
|
|
2
|
+
|
|
3
|
+
Rig plugins combine these with their own settings; none of them is specific to
|
|
4
|
+
one piece of hardware.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from vjepa_byor.actions.codecs import JsonLineCodec
|
|
8
|
+
from vjepa_byor.actions.state import VectorState
|
|
9
|
+
from vjepa_byor.actions.vjepa_mapping import VJEPA_AC_DIM, SlotMapping
|
|
10
|
+
from vjepa_byor.actions.vocabulary import COMPASS, DiscreteActionSpace
|
|
11
|
+
|
|
12
|
+
__all__ = [
|
|
13
|
+
"COMPASS",
|
|
14
|
+
"VJEPA_AC_DIM",
|
|
15
|
+
"DiscreteActionSpace",
|
|
16
|
+
"JsonLineCodec",
|
|
17
|
+
"SlotMapping",
|
|
18
|
+
"VectorState",
|
|
19
|
+
]
|
|
@@ -0,0 +1,69 @@
|
|
|
1
|
+
"""A JSON-per-line wire format, for rigs whose controller can parse JSON (or are simulated).
|
|
2
|
+
|
|
3
|
+
Commands and replies are single-line JSON objects::
|
|
4
|
+
|
|
5
|
+
{"action": [5, -5]} command: one action tensor
|
|
6
|
+
{"state": {"x": 30, "y": 15}} reply: the new state, as RigState.to_dict()
|
|
7
|
+
{"error": "out of range"} reply: the rig refused or failed
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
import json
|
|
13
|
+
import math
|
|
14
|
+
from typing import Any
|
|
15
|
+
|
|
16
|
+
import torch
|
|
17
|
+
|
|
18
|
+
from vjepa_byor.interfaces import CodecError, RigError, RigState
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def _number(value: float) -> int | float:
|
|
22
|
+
return int(value) if float(value).is_integer() else value
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class JsonLineCodec:
|
|
26
|
+
"""Implements :class:`~vjepa_byor.interfaces.CommandCodec` for any state type and action size."""
|
|
27
|
+
|
|
28
|
+
def __init__(self, state_type: type[RigState], action_dim: int) -> None:
|
|
29
|
+
self.state_type = state_type
|
|
30
|
+
self.action_dim = action_dim
|
|
31
|
+
|
|
32
|
+
@staticmethod
|
|
33
|
+
def _parse(line: str) -> dict[str, Any]:
|
|
34
|
+
try:
|
|
35
|
+
data = json.loads(line)
|
|
36
|
+
except json.JSONDecodeError as e:
|
|
37
|
+
raise CodecError(f"not JSON: {line!r}") from e
|
|
38
|
+
if not isinstance(data, dict):
|
|
39
|
+
raise CodecError(f"not a JSON object: {line!r}")
|
|
40
|
+
return data
|
|
41
|
+
|
|
42
|
+
def encode_action(self, action: torch.Tensor) -> str:
|
|
43
|
+
if action.shape != (self.action_dim,):
|
|
44
|
+
raise ValueError(f"expected an action of shape ({self.action_dim},), got {tuple(action.shape)}")
|
|
45
|
+
return json.dumps({"action": [_number(v) for v in action.tolist()]})
|
|
46
|
+
|
|
47
|
+
def decode_action(self, line: str) -> torch.Tensor:
|
|
48
|
+
values = self._parse(line).get("action")
|
|
49
|
+
if not (
|
|
50
|
+
isinstance(values, list)
|
|
51
|
+
and len(values) == self.action_dim
|
|
52
|
+
and all(isinstance(v, int | float) and not isinstance(v, bool) and math.isfinite(v) for v in values)
|
|
53
|
+
):
|
|
54
|
+
raise CodecError(f"bad command: {line!r}")
|
|
55
|
+
return torch.tensor([float(v) for v in values])
|
|
56
|
+
|
|
57
|
+
def encode_state(self, state: RigState) -> str:
|
|
58
|
+
return json.dumps({"state": state.to_dict()})
|
|
59
|
+
|
|
60
|
+
def decode_state(self, line: str) -> RigState:
|
|
61
|
+
data = self._parse(line)
|
|
62
|
+
if "error" in data:
|
|
63
|
+
raise RigError(str(data["error"]))
|
|
64
|
+
if not isinstance(data.get("state"), dict):
|
|
65
|
+
raise CodecError(f"bad reply: {line!r}")
|
|
66
|
+
try:
|
|
67
|
+
return self.state_type.from_dict(data["state"])
|
|
68
|
+
except (KeyError, TypeError, ValueError) as e:
|
|
69
|
+
raise CodecError(f"bad state in reply {line!r}: {e}") from e
|
|
@@ -0,0 +1,47 @@
|
|
|
1
|
+
"""Rig state schema: a base class for states that are a few named numbers."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from dataclasses import asdict, dataclass, fields
|
|
6
|
+
from typing import Any, Self
|
|
7
|
+
|
|
8
|
+
import torch
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
@dataclass(frozen=True)
|
|
12
|
+
class VectorState:
|
|
13
|
+
"""Base for rig states made of numeric fields; implements :class:`~vjepa_byor.interfaces.RigState`.
|
|
14
|
+
|
|
15
|
+
Subclass it as a frozen dataclass; the fields, in declaration order, become
|
|
16
|
+
the state tensor::
|
|
17
|
+
|
|
18
|
+
@dataclass(frozen=True)
|
|
19
|
+
class XYState(VectorState):
|
|
20
|
+
x_mm: float
|
|
21
|
+
y_mm: float
|
|
22
|
+
"""
|
|
23
|
+
|
|
24
|
+
@classmethod
|
|
25
|
+
def field_names(cls) -> list[str]:
|
|
26
|
+
return [f.name for f in fields(cls)]
|
|
27
|
+
|
|
28
|
+
def as_tensor(self) -> torch.Tensor:
|
|
29
|
+
return torch.tensor([float(getattr(self, name)) for name in self.field_names()])
|
|
30
|
+
|
|
31
|
+
def to_dict(self) -> dict[str, Any]:
|
|
32
|
+
return asdict(self)
|
|
33
|
+
|
|
34
|
+
@classmethod
|
|
35
|
+
def from_dict(cls, data: dict[str, Any]) -> Self:
|
|
36
|
+
missing = [name for name in cls.field_names() if name not in data]
|
|
37
|
+
if missing:
|
|
38
|
+
raise ValueError(f"{cls.__name__} is missing {', '.join(missing)}")
|
|
39
|
+
return cls(**{name: data[name] for name in cls.field_names()})
|
|
40
|
+
|
|
41
|
+
@classmethod
|
|
42
|
+
def from_tensor(cls, tensor: torch.Tensor) -> Self:
|
|
43
|
+
"""Inverse of :meth:`as_tensor`; values come back as Python floats."""
|
|
44
|
+
names = cls.field_names()
|
|
45
|
+
if tensor.shape != (len(names),):
|
|
46
|
+
raise ValueError(f"{cls.__name__} needs a tensor of shape ({len(names)},), got {tuple(tensor.shape)}")
|
|
47
|
+
return cls(**dict(zip(names, tensor.tolist())))
|