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.
Files changed (54) hide show
  1. vjepa_byor-0.1.0/LICENSE +21 -0
  2. vjepa_byor-0.1.0/PKG-INFO +203 -0
  3. vjepa_byor-0.1.0/README.md +181 -0
  4. vjepa_byor-0.1.0/pyproject.toml +35 -0
  5. vjepa_byor-0.1.0/pyproject.toml.orig +32 -0
  6. vjepa_byor-0.1.0/src/vjepa_byor/__init__.py +39 -0
  7. vjepa_byor-0.1.0/src/vjepa_byor/actions/__init__.py +19 -0
  8. vjepa_byor-0.1.0/src/vjepa_byor/actions/codecs/__init__.py +5 -0
  9. vjepa_byor-0.1.0/src/vjepa_byor/actions/codecs/json_line.py +69 -0
  10. vjepa_byor-0.1.0/src/vjepa_byor/actions/state.py +47 -0
  11. vjepa_byor-0.1.0/src/vjepa_byor/actions/vjepa_mapping.py +88 -0
  12. vjepa_byor-0.1.0/src/vjepa_byor/actions/vocabulary.py +84 -0
  13. vjepa_byor-0.1.0/src/vjepa_byor/cli.py +340 -0
  14. vjepa_byor-0.1.0/src/vjepa_byor/dashboard/__init__.py +11 -0
  15. vjepa_byor-0.1.0/src/vjepa_byor/dashboard/hub.py +211 -0
  16. vjepa_byor-0.1.0/src/vjepa_byor/dashboard/static/index.html +473 -0
  17. vjepa_byor-0.1.0/src/vjepa_byor/dashboard/web.py +196 -0
  18. vjepa_byor-0.1.0/src/vjepa_byor/data/__init__.py +40 -0
  19. vjepa_byor-0.1.0/src/vjepa_byor/data/collector.py +402 -0
  20. vjepa_byor-0.1.0/src/vjepa_byor/data/dataset.py +309 -0
  21. vjepa_byor-0.1.0/src/vjepa_byor/data/policy.py +140 -0
  22. vjepa_byor-0.1.0/src/vjepa_byor/data/sample.py +98 -0
  23. vjepa_byor-0.1.0/src/vjepa_byor/data/storage.py +291 -0
  24. vjepa_byor-0.1.0/src/vjepa_byor/interfaces.py +183 -0
  25. vjepa_byor-0.1.0/src/vjepa_byor/planning/__init__.py +46 -0
  26. vjepa_byor-0.1.0/src/vjepa_byor/planning/planner.py +268 -0
  27. vjepa_byor-0.1.0/src/vjepa_byor/planning/ranking.py +109 -0
  28. vjepa_byor-0.1.0/src/vjepa_byor/planning/scorer.py +145 -0
  29. vjepa_byor-0.1.0/src/vjepa_byor/predictor/__init__.py +58 -0
  30. vjepa_byor-0.1.0/src/vjepa_byor/predictor/encoding.py +63 -0
  31. vjepa_byor-0.1.0/src/vjepa_byor/predictor/evaluate.py +81 -0
  32. vjepa_byor-0.1.0/src/vjepa_byor/predictor/finetune.py +217 -0
  33. vjepa_byor-0.1.0/src/vjepa_byor/predictor/inference.py +122 -0
  34. vjepa_byor-0.1.0/src/vjepa_byor/predictor/input_layers.py +31 -0
  35. vjepa_byor-0.1.0/src/vjepa_byor/predictor/latent_cache.py +162 -0
  36. vjepa_byor-0.1.0/src/vjepa_byor/predictor/loading.py +115 -0
  37. vjepa_byor-0.1.0/src/vjepa_byor/predictor/lora.py +80 -0
  38. vjepa_byor-0.1.0/src/vjepa_byor/registry.py +120 -0
  39. vjepa_byor-0.1.0/src/vjepa_byor/rig.py +93 -0
  40. vjepa_byor-0.1.0/src/vjepa_byor/ros2/__init__.py +94 -0
  41. vjepa_byor-0.1.0/src/vjepa_byor/ros2/config.py +288 -0
  42. vjepa_byor-0.1.0/src/vjepa_byor/ros2/dashboard.py +156 -0
  43. vjepa_byor-0.1.0/src/vjepa_byor/ros2/link.py +327 -0
  44. vjepa_byor-0.1.0/src/vjepa_byor/ros2/reach_goal.py +253 -0
  45. vjepa_byor-0.1.0/src/vjepa_byor/sim/__init__.py +21 -0
  46. vjepa_byor-0.1.0/src/vjepa_byor/sim/camera.py +115 -0
  47. vjepa_byor-0.1.0/src/vjepa_byor/sim/rig.py +44 -0
  48. vjepa_byor-0.1.0/src/vjepa_byor/sim/ros2_bridge.py +169 -0
  49. vjepa_byor-0.1.0/src/vjepa_byor/sim/stage.py +68 -0
  50. vjepa_byor-0.1.0/src/vjepa_byor/testing/__init__.py +9 -0
  51. vjepa_byor-0.1.0/src/vjepa_byor/testing/collector_conformance.py +61 -0
  52. vjepa_byor-0.1.0/src/vjepa_byor/testing/conformance.py +208 -0
  53. vjepa_byor-0.1.0/src/vjepa_byor/testing/fake_camera.py +50 -0
  54. vjepa_byor-0.1.0/src/vjepa_byor/testing/fake_transport.py +91 -0
@@ -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,5 @@
1
+ """Reusable wire formats for :class:`~vjepa_byor.interfaces.CommandCodec`."""
2
+
3
+ from vjepa_byor.actions.codecs.json_line import JsonLineCodec
4
+
5
+ __all__ = ["JsonLineCodec"]
@@ -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())))