diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7429256..62c6548 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -16,7 +16,7 @@ jobs: python: name: Python ${{ matrix.python }} runs-on: ubuntu-latest - timeout-minutes: 20 + timeout-minutes: 10 strategy: fail-fast: false matrix: @@ -40,10 +40,6 @@ jobs: uv venv uv pip install --torch-backend cpu -c "$RUNNER_TEMP/constraints.txt" -e . --group dev - # The rotation test writes rotation metadata with the ffmpeg CLI. - - name: Install ffmpeg - run: sudo apt-get update -q && sudo apt-get install -y -q ffmpeg - - name: Lint run: | uv run --no-sync ruff check src tests diff --git a/.gitignore b/.gitignore index f8f029a..9d4be85 100644 --- a/.gitignore +++ b/.gitignore @@ -20,3 +20,7 @@ src/jepa_scope/static/ # Claude Code: per-machine settings and agent worktrees .claude/settings.local.json .claude/worktrees/ + +# eb_jepa checkpoints and planning videos stay in the lab directory, never in the repository. +*.pth.tar +*.gif diff --git a/CLAUDE.md b/CLAUDE.md index 4dc9b89..59947b6 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -81,7 +81,10 @@ and report formats live there. - Merging into the default branch. - Publishing a release: tags, GitHub releases, PyPI uploads. - Changing repository settings or visibility, or posting about the project anywhere public. -- Paid compute, such as renting a GPU for embedding runs or `eb_jepa` training. +- Paid compute outside the spending cap of a plan the owner approved. Within that cap, cloud + GPUs (RunPod or similar) need no further approval, provided each pod's GPU type and hourly + price from a live read go in the plan's log before it is created, a watchdog stops it if the + session is lost, and it is deleted when the job ends. ### Model routing diff --git a/docs/decisions.md b/docs/decisions.md index 996a4e2..3c1e7b6 100644 --- a/docs/decisions.md +++ b/docs/decisions.md @@ -27,17 +27,21 @@ anything the code already shows. ## Method · `docs/decisions/method.md` +- 2026-11-06 — Planning protocol: a Markov planning state as one fixed-shape tensor, one-step stateless predict, model-owned energy, summed objective by default - 2026-11-05 — Surprise = mean L1 between predictor output and layer_norm(target), with the context frames encoded on their own; never the HF full-clip forward - 2026-11-05 — Windows of 16 frames at 8 fps, overlapping by half: context is the first half, every later tubelet is predicted exactly once - 2026-11-05 — Any output-changing edit bumps OUTPUT_VERSION, which keys the result cache ## Product & data · `docs/decisions/product.md` +- 2026-11-06 — The planning protocol is drafted here (`src/jepa_scope/planning/`, `docs/planning-protocol.md`) and moves to latent-plan when that project starts - 2026-11-05 — Apache-2.0, public at ryangswe/jepa-scope; V-JEPA 2 weights are downloaded at run time, never redistributed - 2026-11-05 — `.latent` = safetensors plus a validated JSON header; the viewer refuses a video whose SHA-256 differs from the header - 2026-11-03 — jepa-scope is the first tool; the order of the other two is a proposal; tools target a model-agnostic protocol rather than V-JEPA 2 ## Process & tooling · `docs/decisions/process.md` +- 2026-11-05 — Cloud GPUs within the spending cap of an owner-approved plan need no further approval: log the price, run a watchdog, delete the pod +- 2026-11-05 — Milestones follow the roadmap's weekly order; Week 2's scope landed in Week 1, so Weeks 3 and 4 come next - 2026-11-05 — Adopted the workflow kit at the full level, with its process files, the roadmap and this record published in the public repo - 2026-11-05 — CI installs uv.lock versions as constraints with CPU torch; the real-model test stays opt-in and runs locally diff --git a/docs/decisions/method.md b/docs/decisions/method.md index d890ac0..fcdea4c 100644 --- a/docs/decisions/method.md +++ b/docs/decisions/method.md @@ -5,6 +5,48 @@ index is `docs/decisions.md`. --- +## 2026-11-06 — Planning protocol: a Markov planning state as one fixed-shape tensor, one-step stateless predict, model-owned energy, summed objective by default + +**Status:** live, indexed + +**Context.** The 2026-11-03 product decision said the tools target a model-agnostic protocol +rather than V-JEPA 2. Week 1 drafted that protocol and checked it against `eb_jepa`'s Two Rooms +model. latent-plan will build on it, with V-JEPA 2-AC as its second model, so its shape is costly +to change later. + +**Decision.** A world model has `state_shape`, `action_dim`, `encode(obs)`, `predict(z, a)` and +`energy(z, z_goal)`. `z` is a planning state: one fixed-shape tensor holding everything the next +prediction needs, which may be more than the latent `energy` compares (for V-JEPA 2-AC, its token +context and the arm's pose). `predict` takes one step and keeps no state. `encode` accepts the +model's own observation type. `plan_cost` sums the energy over all `T + 1` rollout states by +default, with `"last"` as a planner setting; `plan_cost` adds the energies step by step and keeps +only the current state. Adapters treat inputs as read-only, keep batch rows independent, refuse +to overflow a packed state's declared capacity, and promise inference only. The owner approved +the signatures, then a Fable architect's changes (planning state, `encode(obs: Any)`, the +objective setting), then an outside review's (renaming `latent_shape` to `state_shape`, the +adapter rules, streaming cost, input validation). The interface stays experimental until a +second adapter reproduces its own model's native rollouts. + +**Evidence.** Through an adapter, rollouts and costs match `eb_jepa`'s own, and the reference CEM +returns `eb_jepa`'s plans given the same noise. In a closed-loop engineering check it solved 16 +of 20 episodes against 18 of 20 natively, which does not establish equivalence (details in +`docs/planning-protocol.md`). vjepa2's planner (`notebooks/utils/ +mpc_utils.py`) starts from one image and grows its context from its own predictions, so its +state is Markov and fits a stateless `predict`. + +**Rejected.** +- *An explicit history argument or an `initial_state(obs_history)` hook.* Neither target model's + planner uses an observation history, and a fixed history length does not hold the pose. +- *An opaque state object with a separate latent.* `cem_plan` expands `z0` with tensor + operations, so it would also need an `expand_state` operation; two new operations before any + adapter needs them. +- *A summed objective hard-coded into `plan_cost`.* V-JEPA 2-AC's planner scores only the final + frame, and a parity check against it needs the same objective. +- *Deferred, not rejected:* an optional `latent(z) -> [B, *latent]` for drawing latents along a plan, to be + decided when latent-plan is planned. + +--- + ## 2026-11-05 — Surprise = mean L1 between predictor output and layer_norm(target), with the context frames encoded on their own; never the HF full-clip forward **Status:** live, indexed diff --git a/docs/decisions/process.md b/docs/decisions/process.md index a1053a2..7b89a13 100644 --- a/docs/decisions/process.md +++ b/docs/decisions/process.md @@ -5,6 +5,43 @@ Full rationale for this theme's decisions, newest first. The lean, auto-loaded i --- +## 2026-11-05 — Milestones follow the roadmap's weekly order; Week 2's scope landed in Week 1, so Weeks 3 and 4 come next + +**Status:** live, indexed + +**Context.** After Week 1's scaffold, the next milestone was the owner's choice among the +roadmap's candidates. + +**Decision.** Milestones follow the roadmap's planned weekly order, finishing a week's remaining +steps before moving on. Week 2's scope (inference, the grid, the overlay, the file format and the +CLI) landed in Week 1, so Weeks 3 and 4 come next. + +**Rejected.** +- *Re-choosing the order at each milestone.* The weekly order already reflects the roadmap's + reasoning; changing it is a roadmap edit, not a per-milestone choice. + +--- + +## 2026-11-05 — Cloud GPUs within the spending cap of an owner-approved plan need no further approval: log the price, run a watchdog, delete the pod + +**Status:** live, indexed + +**Context.** Training `eb_jepa`'s Two Rooms model on the laptop CPU projected to about 17 hours. +The owner has a RunPod account with credit and said cloud GPUs are fine if the cost is not +excessive. + +**Decision.** Paid compute outside an approved plan's spending cap stays owner-only. Within the +cap, a session may create cloud GPUs without asking again, provided the GPU type and hourly +price from a live read go in the plan's log before each pod is created, a watchdog stops the pod +if the session is lost, and the pod is deleted when the job ends. `CLAUDE.md` carries the rule. +The Week 1 plan's cap was $25; it spent about $1.65. + +**Rejected.** +- *Asking before every pod.* The cap and the logging already bound the cost, and asking again + adds owner time for no control. + +--- + ## 2026-11-05 — Adopted the workflow kit at the full level, with its process files, the roadmap and this record published in the public repo **Status:** live, indexed diff --git a/docs/decisions/product.md b/docs/decisions/product.md index 047b6da..51376da 100644 --- a/docs/decisions/product.md +++ b/docs/decisions/product.md @@ -5,6 +5,24 @@ lean, auto-loaded index is `docs/decisions.md`. --- +## 2026-11-06 — The planning protocol is drafted here (`src/jepa_scope/planning/`, `docs/planning-protocol.md`) and moves to latent-plan when that project starts + +**Status:** live, indexed + +**Context.** The roadmap's Week 1 included running `eb_jepa`'s Two Rooms example and drafting +the planning protocol, which belongs to latent-plan, the third tool, and does not exist yet. + +**Decision.** The draft lives in `src/jepa_scope/planning/` (protocol, reference CEM, `eb_jepa` +adapter) with its spec in `docs/planning-protocol.md`. It depends only on torch; `eb_jepa` is +imported only when the adapter is built and is never a dependency. The code moves to latent-plan +when that project starts. + +**Rejected.** +- *Starting the latent-plan repository now.* Its plan, UI and audience are not decided; a draft + in an existing tested package is cheaper to change. + +--- + ## 2026-11-05 — Apache-2.0, public at ryangswe/jepa-scope; V-JEPA 2 weights are downloaded at run time, never redistributed **Status:** live, indexed diff --git a/docs/planning-protocol.md b/docs/planning-protocol.md new file mode 100644 index 0000000..29f93e1 --- /dev/null +++ b/docs/planning-protocol.md @@ -0,0 +1,350 @@ +# The planning protocol (draft) + +This document describes a small, model-agnostic interface for planning with a JEPA world model. +A planner written against it can plan with any world model that implements three operations, +whatever its architecture. The code is in `src/jepa_scope/planning/`. The protocol is layer E of +the toolchain described in `docs/roadmap.md` and belongs to latent-plan, the third of its tools; +it is drafted here and moves there when that project starts. + +**Status: experimental.** The protocol has been checked against one model, eb_jepa's Two Rooms +world model. Treat the interface as provisional until a second adapter (V-JEPA 2-AC is the +intended one) reproduces its own model's native rollouts. The signatures may change until then. + +## What a world model implements + +A world model is any object with two attributes and three methods. They work on a **planning +state** `z`: one tensor of fixed shape that holds everything the model needs to predict the next +step. For many models the planning state is simply the latent. For others it holds more, such as +earlier predicted frames, past actions or a robot arm's pose, and the model's `energy` decides +which part to compare with the goal. Shapes put the batch first: `B` is the batch size, `A` is +the action size, and `*state` is the shape of one planning state, which can have any number of +dimensions. + +| Member | Shape | Meaning | +|---|---|---| +| `state_shape` | `tuple[int, ...]` | The shape of one planning state, without the batch dimension. | +| `action_dim` | `int` | The size of one action vector, `A`. | +| `encode(obs)` | `B` observations to `[B, *state]` | Turns observations into planning states. | +| `predict(z, a)` | `[B, *state]` and `[B, A]` to `[B, *state]` | Predicts the planning state one action later. | +| `energy(z, z_goal)` | `[B, *state]` twice to `[B]` | Scores how far each state is from its goal. | + +The three operations have these rules. + +- `encode` accepts whatever observation type the model needs: an image tensor, or an image + together with a robot's pose, for example. Planners pass observations through without looking + inside them. `encode` is used for the start observation and for the goal observation, and + planners never call it on predicted states. +- `predict` takes exactly one step and keeps no hidden state between calls. Everything the model + needs in order to predict must be inside `z`. A recurrent model whose hidden state is its + latent (like the eb_jepa Two Rooms model) satisfies this directly. A model that predicts from + several frames keeps them inside `z`. +- `energy` is non-negative and lower when `z` is closer to the goal. It is zero when the parts of + `z` that the model compares equal those of `z_goal`, so a goal state can leave the parts that + are not compared (such as a pose) as placeholders. Its scale belongs to the model: energies can + be compared within one model, never across two models. + +`WorldModel` in `protocol.py` is a `typing.Protocol`, so a model does not inherit from anything. +`isinstance(model, WorldModel)` checks that the members exist, not their shapes. + +### Rules every adapter follows + +Planners rely on these. The packed-context toy model in `tests/test_planning_protocol.py` follows +them, and its tests show the helpers working with such a model; there is not yet a check you can +run against your own adapter. + +- **Inputs are read-only.** `predict` and `energy` never write into the tensors they are given, + and a state they returned earlier never changes later. Planners branch many candidates from one + start state, often as an expanded view that shares memory, so writing into it would corrupt + every branch. +- **Inputs may be expanded or non-contiguous.** Copy before writing, as the packed-context toy + model does. +- **Batch rows are independent.** Each row's prediction and energy depend only on that row. No + normalisation or reduction may run across the batch. Rows can be at different points of a + rollout, so a packed state that records how much of itself is valid must be read row by row. +- **A packed state has a declared capacity and refuses to overflow.** If the state holds a + context sized for a maximum number of steps, the adapter documents that number, and `predict` + raises an error when a state is already full. It never truncates old steps, drops new ones or + resizes the tensor. Callers keep the planning horizon within the capacity. +- **Inference only.** The protocol promises derivative-free inference: planners call it without + gradients, and nothing guarantees that gradients flow through an adapter. The adapter puts its + own modules in evaluation mode and moves inputs to its own device and dtype. `predict` is + deterministic; stochastic predictors are an open point. + +### Actions + +`action_dim` gives only the length of an action vector. Every adapter documents what an action +means: the order of its coordinates, their units or normalisation, how much time one action +covers (including any repeat), its bounds, and how a planned action is converted into the action +an environment executes. The planner's settings must keep actions within those bounds; `energy` +does not enforce them. + +## Rolling out and scoring a plan + +Two functions are built on the three operations, and planners use only these two. + +`rollout(model, z0, actions)` applies `predict` once per action and keeps every state. `z0` is +`[B, *state]` and `actions` is `[B, T, A]`. It returns `[B, T + 1, *state]`, with `z0` at index 0 +and the state after action `t` at index `t + 1`. It is for callers that need the trajectory; +planners do not. + +`plan_cost(model, z0, actions, z_goal, objective="sum")` is the cost of an action sequence, and it +returns `[B]`. It keeps only the current state, so its memory does not grow with the number of +actions. The objective is a setting of the planner, not of the model, and has two values. + +- `"sum"`, the default, adds up the `energy` to the goal of all `T + 1` states, including the + start state, one step at a time in the dtype `energy` returns. It gives every step of the plan + a cost, which favours plans that reach the goal early and stay there. It is the objective + eb_jepa uses for Two Rooms (`sum_all_diffs: true`), and eb_jepa's README reports a Two Rooms + success rate of 97% for its MPPI planner with it against 89% with only the last state's cost. + That supports it for this task, not as a general rule: a plan that must detour away from the + goal is penalised for the detour, and a model's errors add up across the sum as well. The start + state adds the same amount to every candidate, so including it does not change which plan wins. +- `"last"` scores only the final state. Other native planners use it, V-JEPA 2-AC's among them, + and comparing a model with its native planner needs the same objective. + +## The reference planner: the Cross-Entropy Method + +`cem_plan(model, z0, z_goal, config, steps_left=None, generator=None)` in `cem.py` searches for an +action sequence with the Cross-Entropy Method. It plans from one start state (`z0` and `z_goal` are +`[1, *state]`, on the same device) and returns a `CEMResult` holding the planned actions `[T, A]` +and the lowest sampled cost in each iteration. + +Each iteration works as follows. + +1. It samples `num_samples` action sequences from a Gaussian with a per-step, per-dimension mean + and standard deviation. The mean starts at zero and the standard deviation at `var_scale`. +2. If `max_norm` is set, it rescales every action to a norm of at most `max_norm`. +3. It scores every sequence with `plan_cost`, using the configured `objective`. +4. It takes the `num_elites` lowest-cost sequences and sets the new mean and standard deviation + to their mean and their unbiased standard deviation. + +After the last iteration, the planned sequence is the elite mean, not the best sample seen. The +horizon is `plan_length`, or the steps left in the episode when `steps_left` is given and smaller. +Every call starts afresh from a zero mean. A caller that plans again after each executed action +(model-predictive control) gets eb_jepa's behaviour by executing the first action of each plan. + +**What it supports.** Continuous action vectors with an optional bound on each action's norm, +and nothing else. Actions are sampled and returned in float32 on `z0`'s device, whatever the +planning state's dtype. Because the elite mean of actions inside a norm bound is also inside it, +the returned plan respects the bound. Per-coordinate bounds, discrete parts such as a gripper, or +any constraint whose average is not itself valid would need the planner to project its samples +and its final plan; this planner does not. Settings must be finite, counts must be integers, and +a model that scores a candidate as NaN or infinite makes `cem_plan` raise an error rather than +return a plan. + +`CEMConfig`'s defaults are eb_jepa's settings for Two Rooms (`planning_cem.yaml`): a horizon of 90 +actions, 20 iterations, 200 samples, 20 elites, an initial standard deviation of 1.5, an action +norm limit of 2.45 and the summed objective. The planner mirrors `CEMPlanner.plan` in eb_jepa's +`eb_jepa/planning.py` at commit `966e61e` step by step, including its norm-rescaling formula, so +that the same model can be planned natively and through the protocol and the two can be compared. +Pass a `torch.Generator` for a reproducible plan. + +## Adding a model + +To plan with a new model: + +1. Write an adapter class with the five members in the table above, following the adapter rules. + It usually wraps the model's own encoder and predictor and picks an energy, often a mean + squared or absolute distance between latents. If the model needs more than one latent, or + past actions or poses, to predict, the planning state packs all of them into one fixed shape. + List every input the model's own predictor takes and show where each one lives in the state. +2. Document the model's actions as described under Actions, and the state's capacity if it has + one. +3. Check the adapter before trusting a plan. First, `rollout` through the adapter must give the + same states as the model's own multi-step prediction. Second, if the model has its own + planner, the reference planner fed the same random noise should return the same plan. Third, + planning through the adapter should succeed about as often as the model's own planner on the + same episodes. +4. Plan with `cem_plan`, then convert each planned action into an environment action as + documented. + +## Adapters + +### eb_jepa Two Rooms + +`EbJepaTwoRooms` in `eb_jepa_two_rooms.py` wraps the action-conditioned JEPA in +`examples/ac_video_jepa` of `facebookresearch/eb_jepa` (checked at commit `966e61e`). In that +task, a dot must cross a wall through a door to reach a goal position, and the model sees +two-channel 65 × 65 images. Its planning state is its 512-dimensional latent. The three operations +map onto the model as follows. + +- `encode` applies eb_jepa's Two Rooms `Normalizer` (a min-max scaling of each channel of each + image, followed by fixed per-channel statistics) to a raw `[B, 2, 65, 65]` observation, as + `DotWall` renders it. It then runs the Impala encoder on the result as a one-frame clip and + flattens its `[B, 512, 1, 1, 1]` output to `[B, 512]`. +- `predict` runs one step of the GRU predictor, with the action as the input and the latent as + the hidden state, and then applies the encoder's final LayerNorm, which the predictor shares. + The normalised output is the next latent. eb_jepa feeds it back as the next hidden state, so + the latent carries everything the next step needs and `predict` is stateless, as the protocol + requires. This is exactly one step of eb_jepa's `JEPA.unroll` in autoregressive mode. The state + has no capacity limit. +- `energy` is the mean squared difference over the 512 dimensions. It equals eb_jepa's + `ReprTargetDistMPCObjective` for one state. eb_jepa's README writes this cost with an L2 norm, + but its code takes the mean of the squares, and the adapter follows the code. + +**Actions.** An action is the dot's displacement for one environment step, in the +environment's position coordinates (the 65 × 65 pixel grid), in the same coordinate order as +`DotWall`'s positions, with no repeat. eb_jepa's planner limits its norm to 2.45. `DotWall` +declares a box of ±2.45 per coordinate but does not clip, so the bound comes from the planner. A +planned action is executed by passing it, as a float32 NumPy array, to `DotWall.step` unchanged. + +**Inference.** `encode` and `predict` run without gradients, with the modules in evaluation mode +on the device given when the adapter is built. + +`EbJepaTwoRooms.from_run(folder)` reads a run folder written by eb_jepa's training script. It +rebuilds the encoder and predictor from the run's config, the way the training script builds +them. It loads only their weights from `latest.pth.tar` (or another checkpoint you name), with +`torch.load(weights_only=True)` and strict key matching, so a wrong key mapping fails loudly +instead of leaving weights random. `EbJepaTwoRooms.untrained(folder)` builds the same +architecture with fresh weights, as a control. + +eb_jepa is not a dependency of jepa-scope and is not on PyPI. To use the adapter, install both +into one environment. The check below used Python 3.12 on macOS (Apple silicon) and these steps. + +```bash +git clone https://github.com/facebookresearch/eb_jepa && cd eb_jepa && git checkout 966e61e +uv venv --python 3.12 ../.venv +uv pip install --python ../.venv/bin/python torch==2.6.0 torchvision==0.21.0 einops fire \ + omegaconf pyyaml tqdm gymnasium scikit-learn pandas matplotlib seaborn opencv-python \ + imageio wandb pytest av +uv pip install --python ../.venv/bin/python -e . --no-deps +uv pip install --python ../.venv/bin/python -e --no-deps +``` + +`--no-deps` skips `decord`, which eb_jepa declares but the Two Rooms example never imports, and +which has no wheel for macOS on Apple silicon. On a Linux server, use `opencv-python-headless` in +place of `opencv-python`. + +## Checked against + +**eb_jepa Two Rooms, native planner.** The reference model was trained with eb_jepa's own +`examples/ac_video_jepa` script at commit `966e61e` (random wall positions, seed 1, 12 epochs, on +one NVIDIA L40S in about 1.6 hours). eb_jepa's own CEM evaluation (`planning_cem.yaml`, 20 +episodes of up to 200 steps) then succeeded in 18 of 20 episodes (90%). The evaluation ran on the +same GPU, at about 45 s per episode (921 s in all). With 20 episodes, a success rate near 90% has +a standard error of about 7 points. eb_jepa's README reports 96 ± 2% for CEM, where each figure +averages three training seeds and the last three checkpoints and ± 2 is their spread. One run +and one checkpoint is noisier than that, so 90% is consistent with it. + +eb_jepa counts an episode as a success by the agent's distance to the goal after the last allowed +step, not when it first reaches the goal, and the numbers here use the same rule. + +The run was trained and evaluated with these commands, from the eb_jepa clone, in the environment +described under Adapters (on a CUDA machine for training). Boolean overrides need a capital +`True` or `False`: eb_jepa merges command-line values into an untyped config, so a lowercase +`false` arrives as the string `"false"`, which counts as true. + +```bash +export EBJEPA_CKPTS= WANDB_DISABLED=true +python -m examples.ac_video_jepa.main --logging.log_wandb False --meta.enable_plan_eval False --meta.seed 1 +# is $EBJEPA_CKPTS/ac_video_jepa/dev_/_seed1, the folder holding latest.pth.tar +cp /config.yaml /config.train.yaml # the evaluation rewrites config.yaml +python -m examples.ac_video_jepa.main --meta.model_folder --meta.eval_only_mode True \ + --meta.enable_plan_eval True --eval.plan_cfg_path examples/ac_video_jepa/cfgs/planning_cem.yaml \ + --logging.log_wandb False +``` + +**Through the protocol.** `tests/test_eb_jepa_two_rooms.py` checks the adapter against that same +run on a laptop CPU (Apple M3 Max, float32, torch 2.6.0). + +- **Same computation.** On 8 observations and 8 random actions, `rollout` through the adapter + matches eb_jepa's own `JEPA.unroll` exactly (maximum difference 0.0), and `plan_cost` matches + eb_jepa's `ReprTargetDistMPCObjective` to a relative difference of 9e-8 (from the order of + addition). The reference is a separate `JEPA` loaded by eb_jepa's own checkpoint loader, so the + check does not compare the adapter with itself. +- **Same planner.** Given the same Gaussian noise, the reference CEM returns the same plan as + eb_jepa's `CEMPlanner` (maximum difference 6e-7 with a 25-step horizon, and 0.0 when the + horizon is cut to the steps left). Changing the reference CEM's elite update to a biased + standard deviation makes this check fail by 0.45, so it would catch the planner drifting from + eb_jepa's. +- **Same starting frames.** The test rebuilds eb_jepa's evaluation environment and replays its + 20 episodes with eb_jepa's reset procedure. Each replayed start frame (wall, door and dot) + matches the start frame eb_jepa saved for that episode, and the test also checks that pairing + each replay with the next episode instead makes every frame differ (by 356 to 1,344 pixels). Goals are not compared directly: eb_jepa saves them only in a PDF. + They come from the same generator calls as the start, but the test does not check them, and + nothing here shows that the two runs' trajectories or wall-collision noise were the same. +- **Closed-loop planning, an engineering check.** Planning through the protocol solved 16 of the + 20 episodes, against 18 of 20 for eb_jepa's own planner, and passed a threshold fixed before + the run: a protocol rate of at least 70%, and at most two more episodes solved only natively + than only through the protocol. That second quantity equals the difference in total successes, + so the rule is a net-count threshold, not a statistical test. The two runs also differed in + device (CPU against GPU), random draws and possibly precision. The planners disagreed on two + episodes, both solved only natively; under equal planners a split that favours the native + planner this much happens about one time in four. So this small experiment does not show that + the protocol loses no planning quality, nor does it measure a gap: 20 episodes cannot rule out + a gap of 10 points or more. Replanned with two other seeds, both disputed episodes succeed, + which suggests sampling luck, but it does not change the full comparison. +- **Baselines.** A policy that steps straight towards the goal solved 9 of 20 episodes (45%), and + the same planner on an untrained model of the same architecture solved none of 10. These show + the learned model and planner are doing useful work; they say nothing about equivalence with + the native planner. + +In short: on the tested CPU cases, the adapter reproduces native rollout and cost values, and the +reference CEM reproduces native plans given the same noise. In a separate closed-loop check on a +different device, it solved 16 of 20 tasks against 18 of 20 natively and passed the pre-set +threshold. That does not establish non-inferiority. + +**Where the numbers come from.** The closed-loop run (82 minutes) used the code at commit `00e6919`. `plan_cost` has +since been changed to add energies one step at a time, which changes its result by about 1e-7, +and the same-noise check was added afterwards; both short checks were rerun on the later code. A +protocol episode took 166 s on average. eb_jepa's own CEM took 805 s per episode on the same +laptop in an earlier timing with an untrained 30-step checkpoint over two episodes; the +computation does not depend on the weights, but that timing is not a controlled benchmark. +eb_jepa's planner encodes the start observation once for each of its 200 samples, and the +protocol's once per re-plan, which likely accounts for much of the difference. + +To reproduce the checks, train and evaluate a run with the commands above, then run this from the +eb_jepa clone, in the same environment: + +```bash +JEPA_SCOPE_TWO_ROOMS_CHECKPOINT= python -m pytest /tests/test_eb_jepa_two_rooms.py -s +``` + +The run folder is the one eb_jepa's training script writes, and it must also hold that run's +native CEM evaluation (`plan_eval/*_eval_only`), which the test compares against. The test uses +the folder's `latest.pth.tar`. Add `-k "match_eb_jepa or same_noise or replay"` for the short +checks, which take seconds; the full test takes over an hour on a laptop CPU. + +**Not yet checked.** The closed-loop comparison on one device with shared candidate noise; the +replayed goals; a packed state's capacity on a real model; and a bounded-action planner whose +executed actions are checked against the ones it predicted with. + +## Open points + +- **Fitting V-JEPA 2-AC.** V-JEPA 2-AC is the next model this protocol should carry. Its own + planner (`notebooks/utils/mpc_utils.py` in `facebookresearch/vjepa2`) starts from one image, + grows the context with its own predicted frames during a rollout, and updates the robot arm's + pose from each action. Its predictor is called with the token context, the actions so far and + the poses so far. That is a Markov state, so it can fit `predict(z, a)` without a history + argument, but only if the planning state holds all three. A sketch, not yet built: the state + holds a token context, the actions and poses so far, and the number of valid frames, sized for + a declared capacity; `energy` is the mean absolute difference over the latest valid frame's + tokens, as in vjepa2's `l1`. The adapter must list the predictor's inputs against the pinned + vjepa2 code and match its native rollout before this counts as compatible. If packing turns + out costly, a structured state object is still an option; it would need an operation to + repeat a state across candidates. Whether vjepa2 layer-normalises predicted tokens the same way + as encoded ones (its `normalize_reps` setting) must also be checked then, because the energy + depends on it. +- **A representation for drawing.** A tool that draws latents along a plan (latent-plan's + playground, for example) needs the comparable part of a planning state, not the packed state, + which may mix tokens with poses, counts and padding. An optional operation such as + `latent(z) -> [B, *latent]`, with its own documented shape (`latent_shape`), would give it: the + identity for Two Rooms, a slice for V-JEPA 2-AC. Planners do not need it, so it is not part of + this draft. Decide it when latent-plan is planned, before its interface is built. +- **Memory.** `plan_cost` keeps only the current state, so planning memory is roughly the + candidates times one planning state, plus the model's own working memory. A state that holds a + growing context is still large: its size grows with the capacity, and in float32 it takes + twice the bytes of bf16. Measure the peak on the target device when such an adapter is built, + and score candidates in chunks if needed. +- **Action spaces.** The reference planner handles continuous actions with a norm bound only. + V-JEPA 2-AC's planner instead clips each coordinate to a box, blends each new mean and + standard deviation with the previous ones (momentum), keeps rotation at zero, and bounds and + rounds the gripper. A planner matching it needs those settings and a projection of its final + plan, not only the model's energy. +- **Stochastic predictors.** `predict` returns one next state. Models that predict a distribution + would need samples or an expected energy. +- **Energy across models.** Energies are only comparable within one model. Comparing planners + across models would need a calibrated scale. jepa-scope's surprise view meets the same problem + when it compares prediction errors between clips. +- **Goals per sample.** `plan_cost` takes one goal per batch row. Planning towards several goals + at once works by repeating the start state, which costs memory. diff --git a/docs/plans/2026-11-05-week1-two-rooms-protocol.md b/docs/plans/2026-11-05-week1-two-rooms-protocol.md new file mode 100644 index 0000000..7fe115c --- /dev/null +++ b/docs/plans/2026-11-05-week1-two-rooms-protocol.md @@ -0,0 +1,263 @@ +# Week 1 closeout: run eb_jepa Two Rooms and draft the planning protocol + +- **Date:** 2026-11-05 +- **Status:** Done +- **Planned by:** planner (Fable 5.1), briefed by the lead session +- **Spec / source:** `docs/roadmap.md`, Week 1 table, last open step; layer E of "The missing toolchain"; the 2026-11-03 product decision that tools target a protocol rather than one model +- **Branch:** `week1-two-rooms-protocol` + +## Goal + +Finish the one Week 1 step that has not started. Train eb_jepa's action-conditioned Two Rooms world model on a rented RunPod GPU (with the laptop's CPU as the fallback), reproduce its planning success rate, and draft the model-agnostic planning protocol (`encode(obs) → z`, `predict(z, a) → z'`, `energy(z, z_goal)`) as a spec plus a small typed Python module. "Drafted" here means checked against reality: an adapter wraps the trained eb_jepa model in the protocol, a planner written only against the protocol reaches a success rate on the same task that matches eb_jepa's own planner, and a model with untrained weights does not. The numbers go in the spec. + +This plan does not build latent-plan. There is no browser UI, no V-JEPA 2-AC adapter, no MPPI beyond what the check needs, no new repository, and no change to jepa-scope's inspector behaviour. Weeks 3 and 4 of jepa-scope follow this milestone and get their own plan. + +## Approach + +All eb_jepa work happens in a lab directory outside the repository, `~/Workspace/eb_jepa-lab/`, holding a clone of `facebookresearch/eb_jepa` pinned to commit `966e61e`, a Python 3.12 virtual environment with the pinned `torch==2.6.0`, and the checkpoints (`EBJEPA_CKPTS=~/Workspace/eb_jepa-lab/checkpoints`). The scratchpad is not durable (macOS purges it after a few days), so the lab is a normal directory the owner can see. Nothing in it is committed anywhere: no weights, no GIFs, no eval CSVs, no eb_jepa code. + +Training runs on a RunPod pod by default, because that is the setting eb_jepa's numbers come from (CUDA, `torch.compile`, bf16) and it costs a few dollars, while CPU training on the laptop is untested, slower and not comparable with the README. The owner has a RunPod account with credit, its SSH keys already include this laptop's, and `runpodctl` is installed. The pod runs the same clone, venv and commands as the lab, under `/workspace/eb_jepa-lab/`; the finished run folder is copied back to `~/Workspace/eb_jepa-lab/checkpoints/` with `scp` (or `runpodctl send`), and the pod is deleted and the deletion confirmed by a fresh read. torch 2.6.0's CUDA builds have no kernels for Blackwell GPUs (RTX 5090, RTX PRO 6000, B200), so those are excluded; the create sets a CUDA floor of 12.4, which current hosts on CUDA 13.x drivers satisfy. The pod needs at least 16 vCPUs, because the data loader generates every trajectory on the CPU with 16 workers. + +eb_jepa's `pyproject.toml` requires Python `==3.12.*` (Homebrew has 3.12.15) and declares `decord`, which has no macOS arm64 wheel but is imported by no file in the repository. The install therefore uses `uv pip install -e --no-deps` followed by an explicit list of what the Two Rooms example actually imports. Device selection is hard-coded to CUDA-or-CPU in three places (`examples/ac_video_jepa/main.py` and the CEM and MPPI planners in `eb_jepa/planning.py`), so the first runs are on CPU; an MPS variant is tried only if CPU is too slow, as a three-line edit in the lab clone that the reproduction notes describe and that is never committed. Training runs with wandb logging off (no account needed; `wandb` must still be importable because `main.py` imports it at module level) and with the in-training planning evaluation switched off, since each evaluation is 20 episodes of planning. The planning evaluation runs once at the end with `--meta.eval_only_mode True`, using eb_jepa's CEM configuration so the later protocol check compares like with like. + +Three rules apply to every eb_jepa command in this plan. First, boolean overrides are written `True` or `False` with a capital letter: `main.py` passes arguments through `fire.Fire` and merges them with `OmegaConf.merge` into an untyped config, so a lowercase `false` arrives as the string `"false"`, which Python treats as true, and wandb, the in-training planning evaluation and AMP would all stay on. `WANDB_DISABLED=true` is also exported as a backstop. Second, every command runs from the clone root, `~/Workspace/eb_jepa-lab/eb_jepa`, because `examples` is not an installed package and `--fname` defaults to a relative path; `EBJEPA_CKPTS=~/Workspace/eb_jepa-lab/checkpoints` is exported first. Third, any config file that lives in the lab rather than in the clone is passed by absolute path. On a RunPod pod the lab is `/workspace/eb_jepa-lab/` instead of `~/Workspace/eb_jepa-lab/`, and the exports must be repeated in every SSH shell. + +In jepa-scope, the draft lives in a new subpackage `src/jepa_scope/planning/` (`protocol.py` with the `WorldModel` protocol and the `rollout` and `plan_cost` helpers, `cem.py` with a reference CEM planner, `eb_jepa_two_rooms.py` with the adapter behind a lazy import) and a spec at `docs/planning-protocol.md`. The subpackage depends only on torch. The spec is the document other people and the later latent-plan project write adapters against; it records the shapes, the cost definition, the reproduction steps, the measured numbers and the open points. Tests split the way `tests/test_real_model.py` does: `tests/test_planning_protocol.py` runs in CI on a toy point-mass world model and needs nothing but torch; `tests/test_eb_jepa_two_rooms.py` skips unless `JEPA_SCOPE_TWO_ROOMS_CHECKPOINT` points at a trained run, and is executed from the lab venv, into which jepa-scope is installed with `--no-deps` (the planning subpackage imports only torch, and running it on torch 2.6.0 is the lower-risk direction; the reverse, eb_jepa on torch 2.15, is not attempted). + +The adapter is where the protocol meets the model. `encode` normalises a `[B, 2, 65, 65]` observation with eb_jepa's `Normalizer.normalize_state`, runs `ImpalaEncoder` on it as a one-frame clip and flattens `[B, 512, 1, 1, 1]` to `[B, 512]`. `predict` calls `RNNPredictor(z.view(B, 512, 1, 1, 1), a.unsqueeze(-1))` and flattens again, which is exactly one step of `JEPA.unroll` in autoregressive mode (the GRU hidden state is the latent). `energy` is the mean squared difference over the 512 dimensions, which is `ReprTargetDistMPCObjective` with H = W = 1. (eb_jepa's README writes this cost with an L2 norm, but the code takes an MSE mean, `planning.py:495-496`; the protocol follows the code and the spec says so.) Strictly, the latent is the GRU output after the shared LayerNorm, which is then fed back as the hidden state; the spec says this rather than "the hidden state is the latent". The model is rebuilt from the run's saved `config.yaml` the way `main.py:174-218` builds it, keeping `final_ln=encoder.final_ln` for the predictor (it shares the encoder's LayerNorm, and `None` crashes at `architectures.py:452`). The checkpoint is read with `torch.load(..., weights_only=True, map_location=device)` (it was saved on a CUDA machine), because the adapter ships in a public wheel and eb_jepa's own loader uses `weights_only=False`; only `model_state_dict` is used, the `regularizer.*` keys and any `_orig_mod.` prefix that `torch.compile` adds on CUDA are dropped, and the rest is loaded with `strict=True` into a module holding the encoder and predictor, so a wrong key mapping fails loudly instead of leaving weights random. If `weights_only=True` refuses the file (the checkpoint also holds optimizer, scheduler and probe states), the named globals are allowlisted with `torch.serialization.add_safe_globals` when they are plain library types; otherwise T4 escalates. The adapter never falls back to `weights_only=False`. The adapter imports only `eb_jepa.architectures`, `eb_jepa.jepa` and `eb_jepa.datasets.two_rooms`, never `eb_jepa.planning`, whose plotting imports (cv2, imageio, seaborn, matplotlib, pandas) it does not need. + +The shared contract in this plan is between `protocol.py`, the spec, the toy test and the adapter: a change to the protocol's signatures changes all four together. + +## Alternatives considered + +- **Create the latent-plan repository now and draft the protocol there.** Rejected for now. The 2026-11-03 decision keeps the order of the other two tools a proposal until one is planned, a repository holding one draft file is a weaker home than a package with tests and CI, and creating it is owner-only. The spec says the subpackage moves to latent-plan when that project is planned. This is an owner question (below), not an assumption. +- **Keep the draft out of the package, under `experiments/`.** Rejected because it would need its own lint, test and import wiring, and the subpackage is tiny and torch-only. If the owner prefers the wheel not to carry a draft API, the tasks below move the three files to `experiments/two_rooms/` and add that path to the lint commands; nothing else changes. +- **Write the spec from the paper and the code without training.** Rejected. "A signal that runs is not a signal that works" applies to a protocol as much as to an energy; without a trained model there is no way to know the adapter is right, and the roadmap wants the protocol drafted against a real world model. +- **Pick the fixed-wall task because it is easier.** Rejected. Training cost is the same (the data size does not change), the README's reference numbers are for the random wall, and the random wall is what the README calls the final setup. The fixed wall remains a fallback if the random-wall run fails its threshold. +- **Train on the laptop's CPU overnight.** Rejected as the default now that a RunPod account is available. It is free but untested for this example, needs AMP and loader settings eb_jepa never measured, and gives a result less comparable with the README. It remains the fallback. +- **Patch eb_jepa for MPS by default.** Rejected as the first move: bf16 autocast and `GradScaler` on MPS are not verified for torch 2.6.0, and the planner is latency-bound where MPS launch overhead can lose to CPU. Measure CPU first; try MPS only for training if CPU training is over the threshold. + +## Risks and open questions + +- **Running external code on the laptop needs the owner's approval (question).** Installing and running eb_jepa's code was not pre-approved, and the planning session's safety check stopped it. Approving this plan is the approval for running eb_jepa in the lab directory on the laptop (T1, T4 and the fallback) and on RunPod pods (T2), with RunPod spend up to $25 in total. A pod is also the lower-risk place to run outside code, since it is a disposable machine. Say so explicitly when approving. +- **Where the draft lives (question).** Recommendation: `src/jepa_scope/planning/` as above, with a decision line recording that it moves to latent-plan later. Alternative: `experiments/two_rooms/`. The owner picks at approval. This is more than a file location: the subpackage puts a draft public API, and a first piece of latent-plan, into jepa-scope's published wheel. The spec and the module docstring mark it as a draft that may change without notice. +- **Compute cost.** The configs are tuned for an H100 (batch 384, 12 epochs, 100k trajectories per epoch, so 260 steps per epoch and 3,120 steps in all), and the README says training takes up to a few hours on one GPU. RunPod's catalog on 2026-11-05 listed community-cloud prices per hour of $0.79 for an L40S (48 GB, 16 vCPUs minimum), $1.39 for an A100 SXM 80 GB, $1.99 for an H100 PCIe and $2.69 for an H100 SXM, all with low stock. The default is an L40S, falling back to an A100 or an H100 PCIe if it is out of stock; a loader-bound run needs a host with more vCPUs, not a bigger GPU. Expected spend is about $3 to $10 for training and evaluation. **Budget:** the owner allowed RunPod or another cloud GPU "as needed if the required expenditure isn't too crazy"; this plan's cap is $25 in total, which covers a retry and the MPPI evaluation (T4, including its contingency, runs on the laptop and costs nothing). Within the cap, each pod create needs no separate approval, but its GPU type and hourly price from a live catalog read go in the log before the create. Anything that would pass the cap stops for the owner. The laptop path (CPU overnight if T1 projects at most 10 hours) is the fallback only when no suitable GPU is in stock. +- **A single run may land below the README's numbers.** The README's 97 ± 2 % (MPPI) and 96 ± 2 % (CEM) average three seeds and three checkpoints. One seed and one checkpoint can be lower. T2 escalates below 80 %; between 80 and 90 the run is accepted and the spread is noted in the spec. +- **Twenty episodes are a small sample.** A rate measured on 20 episodes carries about ±7 points of sampling error, and two independent 20-episode runs of equally good planners differ by more than 10 points about a third of the time at an 80 % rate. T4 therefore does not compare two independent rates. eb_jepa's evaluation episodes are deterministic (`DotWall` seeds its generator with 0, `env.py:52`, and `main_eval` calls `reset()` once before its loop, `planning.py:177-178`), so T4 replays the same 20 episodes and compares the two planners episode by episode, with the acceptance rule fixed before anything runs. Stepping the environment does not draw from that generator (wall-bump noise uses the global torch generator), so the replay stays aligned whatever the planner does. The two planners' trajectories still differ through sampling noise, so the rule tests the direction of disagreement (does the protocol lose episodes the native planner wins more often than the reverse?), not its total, which mostly measures noise and would fail a correct adapter most of the time at an 80 % rate. *(Revised after the outside review: native-only minus protocol-only successes equals the difference in total successes, so this rule is a net-count engineering threshold, not a test of direction; see the completion summary.)* +- **MPS unknowns if it is needed.** `torch.compile` is already skipped off CUDA. bf16 autocast, `GradScaler("mps")` and GRU kernels on MPS are unverified for torch 2.6.0; the MPS attempt in T1 runs with `--training.use_amp False` and stops at the first error rather than debugging it. +- **The protocol's single-step `predict` fits an RNN but not a context-window predictor.** V-JEPA 2-AC's predictor needs a history of latents. The draft keeps `predict(z, a)` and records the extension (a context of K latents, or an opaque planner state) as an open point to settle when that adapter is written. This is deliberate scope, not an oversight. +- **Torch version skew in the check.** jepa-scope's planning module runs on torch 2.6.0 in the lab venv. It uses only basic tensor ops; if any call fails there, the module is simplified rather than the pin changed. +- **The milestone order itself is a decision.** The owner chose "milestones follow the roadmap's weekly order, so the remaining Week 1 step comes before Weeks 3 and 4". A later session could re-ask it, so it meets the bar: one line under Process in `docs/decisions.md` and an entry in `docs/decisions/process.md`, written in T5. The roadmap's "Next" section is updated at the same time. + +## Dependency summary + +```mermaid +flowchart LR + T1 --> T2 + T1 --> T3 + T2 --> T4 + T3 --> T4 + T4 --> T5 +``` + +- **Parallel group A:** T2 (training, wall-clock bound, external process) and T3 (protocol draft, in the repository). +- **Integration order:** T1, then T2 beside T3, then T4, then T5. + +## Task graph + +### T1 — Set up the lab, measure the laptop, and pick the RunPod GPU + +- **Status:** Done +- **Dependencies:** None +- **Parallel group:** Sequential +- **Complexity:** ROUTINE +- **Execution:** Lead session +- **Owner time:** None +- **Expected files/workspaces:** `~/Workspace/eb_jepa-lab/` (clone, `.venv`, `checkpoints/`, `cfgs/`); `docs/plans/2026-11-05-week1-two-rooms-protocol-log.md` +- **Objective:** Clone eb_jepa at `966e61e` into `~/Workspace/eb_jepa-lab/eb_jepa`, create a Python 3.12 venv at `~/Workspace/eb_jepa-lab/.venv`, install `torch==2.6.0` and the imports the Two Rooms example needs (`torchvision einops fire omegaconf pyyaml tqdm gymnasium scikit-learn pandas matplotlib seaborn opencv-python imageio wandb pytest av`, then `uv pip install -e . --no-deps` on the clone; add whatever `python -c "import examples.ac_video_jepa.main"` still reports missing). PyAV (`av`) is there because jepa-scope's `tests/conftest.py` imports it, and T4 runs pytest from this venv. Write `~/Workspace/eb_jepa-lab/cfgs/eval_2.yaml` (a copy of `examples/ac_video_jepa/cfgs/eval.yaml` with `num_eval_episodes: 2`). Then, from the clone root with `EBJEPA_CKPTS` and `WANDB_DISABLED=true` exported, time a 30-step training run on CPU with no code edits: `python -m examples.ac_video_jepa.main --logging.log_wandb False --meta.enable_plan_eval False --data.num_workers 8 --data.size 11520 --optim.epochs 1` (11,520 = 30 × 384 and the loader drops the last partial batch, so one epoch is 30 steps and a checkpoint is saved). Run it a second time with `--training.use_amp False`, because bf16 autocast on an Apple CPU may be slower than float32. Time planning with those untrained weights, since planning cost does not depend on the weights: `python -m examples.ac_video_jepa.main --meta.model_folder --meta.eval_only_mode True --meta.enable_plan_eval True --eval.eval_cfg_path ~/Workspace/eb_jepa-lab/cfgs/eval_2.yaml --eval.plan_cfg_path examples/ac_video_jepa/cfgs/planning_cem.yaml --logging.log_wandb False`. Try `num_workers` 4 and 8 if the first run looks loader-bound (the data is generated per trajectory in `__getitem__`). The laptop timings serve T4, which plans on this CPU, and the fallback; if T4's projected episodes are long, they run unattended overnight on the laptop. Only if the laptop becomes the training path and the CPU projection exceeds 10 hours, apply the three-line MPS edit in the lab clone (`main.py` line 95 and `planning.py` lines 560 and 683, `"mps"` when `torch.backends.mps.is_available()`), rerun the training timing with `--training.use_amp False`, and revert the edit if it fails. With wandb off, `main.py` logs no per-step times, so step time comes from tqdm's rate over the last 20 steps (or the wall clock of the whole run divided by 30), and episode time from `avg_episode_time` in the eval CSV. Record the chip and memory (`sysctl -n machdep.cpu.brand_string hw.memsize`), seconds per step with AMP on and off, seconds per planning episode, the projections for 3,120 training steps, for the T2 evaluations (20 CEM episodes, plus 20 MPPI if run) for T4's episodes (20 protocol, 20 greedy at negligible cost, 10 untrained), and for T4's contingency (50 native CEM plus 50 protocol episodes on the laptop), and the recommendation. Note that eb_jepa's planner re-encodes the start observation for every sample, so its episode time is not comparable with the protocol planner's and is not used in any comparison. +- **Success criteria:** + - The example imports and runs 30 steps on CPU and writes `latest.pth.tar`. The log has no "Planning eval:" line, and the run completing at all shows `log_wandb` was read as false (a truthy value calls `wandb.log` without an init at step 0, `main.py:387`; the "W&B logging disabled" line comes from `WANDB_DISABLED` and proves nothing about the override). The AMP-off run logs "AMP disabled" (`main.py:125`). + - The log has measured seconds per training step and per planning episode, the projections for T4 and for the laptop fallback, and the GPU type picked for T2 from a live RunPod catalog read with its hourly price and projected total. +- **Verification:** + - `ls ~/Workspace/eb_jepa-lab/checkpoints/ac_video_jepa/*/*/latest.pth.tar` + - The log entry with the numbers. +- **Escalate if:** + - torch 2.6.0 or any required package has no macOS arm64 wheel for Python 3.12, or the example fails to import after the explicit install. + - The projected RunPod total would pass the $25 cap: stop and ask the owner. + +### T2 — Train the random-wall model on a RunPod pod and reproduce eb_jepa's planning success rate + +- **Status:** Done +- **Dependencies:** T1 +- **Parallel group:** A +- **Complexity:** ROUTINE +- **Execution:** Lead session, on one RunPod pod driven with `runpodctl` and SSH from the laptop. The laptop fallback (only when no suitable GPU is in stock) runs the training and evaluation commands below in the lab instead, with T1's loader and AMP settings. +- **Owner time:** None, unless the cap would be passed +- **Expected files/workspaces:** on the pod, `/workspace/eb_jepa-lab/` (clone, venv, checkpoints, logs); on the laptop, `~/Workspace/eb_jepa-lab/checkpoints/ac_video_jepa//` (weights, `config.yaml`, `config.train.yaml`, `plan_eval/`, `unroll_eval/`, `train.log`, `eval_cem.log`); the plan log +- **Objective:** Follow these steps in order. + 1. **Pick the image and GPU from live reads.** List official images with `runpodctl template list --type official` and pick a `runpod/pytorch` image built on CUDA 12.4 (its tag names `cuda12.4` or `cu124`; the image's own torch is not used). Official images run sshd and include gcc and `runpodctl`. Read the catalog (`runpodctl gpu list`, or the RunPod `list-gpu-types` tool with availability) and pick the cheapest in-stock non-Blackwell GPU from L40S, A100 80 GB and H100 PCIe whose hosts offer at least 16 vCPUs. Write the image tag, GPU, tier, hourly price, projected hours and the running total against the $25 cap into the log before creating anything. + 2. **Create the pod.** `runpodctl pod create --name jepa-scope-two-rooms --image --gpu-id "" --cloud-type COMMUNITY --public-ip --min-cuda-version 12.4 --container-disk-in-gb 20 --volume-in-gb 40 --volume-mount-path /workspace --ports 22/tcp --wait --wait-timeout 30m` (use the image's CUDA version as the floor if it is higher than 12.4; use SECURE instead of COMMUNITY only if no community host is in stock, and log its price). The volume keeps the lab across a pod restart; the container disk does not. Record the pod id. If the create returns any error that carries a pod id (for example a `--wait` timeout while a large image pulls), the pod exists and is billing: record the id, start both watchdogs and read the pod back before doing anything else, and never create a second pod in its place. Read the pod back and confirm GPU, image, ports, the vCPU count and the billed `costPerHr`, which is the price the log uses. `runpodctl pod create` has no vCPU filter, so a pod with fewer than 16 vCPUs is deleted and re-picked; to require more vCPUs, create with the RunPod `create-pod` tool and its `gpu.minVcpuCountPerGpu` field instead. + 3. **Start two spending watchdogs at once.** The primary one runs on the pod, so it works even if the laptop sleeps or the lead session ends. RunPod gives every pod a pod-scoped `RUNPOD_API_KEY` and `RUNPOD_POD_ID` in PID 1's environment (not in SSH shells), so over SSH run `export $(tr '\0' '\n' ; runpodctl pod stop \$RUNPOD_POD_ID" > /workspace/eb_jepa-lab/watchdog.log 2>&1 &`. The second runs on the laptop as a backstop: `nohup caffeinate -i sh -c 'sleep <2 × projected seconds + 1800>; runpodctl pod stop ' > ~/Workspace/eb_jepa-lab/watchdog-.log 2>&1 &`. Log both PIDs. A stop frees the GPU and keeps the 40 GB volume (whose storage costs about a cent an hour), so the run survives. Restart the on-pod watchdog with a new limit whenever the projection changes, and kill the laptop one after the pod is deleted. + 4. **Set up over direct SSH.** RunPod's proxied SSH does not carry `scp`, so use the direct endpoint from `runpodctl ssh info ` (`ssh root@ -p `, and `scp -P ` later). On the pod, install tmux (`apt-get update && apt-get install -y tmux`) and uv, then create the venv with uv's own Python, which ships the headers Triton needs to compile its launcher. uv keeps its Pythons and cache under the home directory, which is on the container disk and is wiped by a restart, so put them on the volume first: `export UV_PYTHON_INSTALL_DIR=/workspace/eb_jepa-lab/python UV_CACHE_DIR=/workspace/eb_jepa-lab/.uv-cache`, then `uv venv --python 3.12 --python-preference only-managed /workspace/eb_jepa-lab/.venv`. Clone eb_jepa at `966e61e` into `/workspace/eb_jepa-lab/eb_jepa` and install as in T1, with `torch==2.6.0` and `torchvision==0.21.0` in the same command so the resolver cannot pull a newer torch. Environment variables set at create time do not reach SSH shells, so every command below starts with `export EBJEPA_CKPTS=/workspace/eb_jepa-lab/checkpoints WANDB_DISABLED=true UV_PYTHON_INSTALL_DIR=/workspace/eb_jepa-lab/python UV_CACHE_DIR=/workspace/eb_jepa-lab/.uv-cache` and runs from the clone root. After a restart, reinstall tmux and uv (both on the container disk); the venv and its Python survive on the volume. + 5. **Smoke-test the CUDA path.** Run T1's 30-step command on the pod, but with the README's loader and precision (`--data.num_workers 16`, AMP and `torch.compile` left on as `train.yaml` sets them): `python -m examples.ac_video_jepa.main --logging.log_wandb False --meta.enable_plan_eval False --data.size 11520 --optim.epochs 1`. The log must show "Using device: cuda" and "Compiling model with torch.compile". This turns the GPU projection into a measurement: update the log's projected hours, the running total and, if the projection grew, the watchdog. If compile fails, retry once; if it fails again, run with `--model.compile False` and log it as a deviation. If `nvidia-smi` shows the GPU mostly idle, the loader is CPU-bound; the fix is a host with more vCPUs (24 or 32), not a bigger GPU. If the GPU runs out of memory, move to an 80 GB card rather than change the batch size. + 6. **Train.** In a `tmux` session (which survives an SSH disconnect), run the full training with the README's best configuration, unchanged except for logging and the in-training evaluation: `python -m examples.ac_video_jepa.main --logging.log_wandb False --meta.enable_plan_eval False --meta.seed 1 2>&1 | tee -a /workspace/eb_jepa-lab/train.log` (`-a`, so a resumed run appends rather than overwriting the log). Check in from the laptop about once an hour. Training has finished when `e-11.pth.tar` exists in the run folder, the Python process has exited and the log's last line is the end of epoch 11. If the pod restarts or the process dies, resume with the same command plus `--meta.model_folder `, which reloads `latest.pth.tar` (`main.py:249-258`); a restart counts as one strike, and its cost counts against the cap. + 7. **Evaluate on the pod.** Copy `/config.yaml` to `/config.train.yaml`, because an eval-only run rewrites `config.yaml` with the eval-time config (`main.py:159-161`). Then run the CEM evaluation on 20 episodes and keep its output: `python -m examples.ac_video_jepa.main --meta.model_folder --meta.eval_only_mode True --meta.enable_plan_eval True --eval.plan_cfg_path examples/ac_video_jepa/cfgs/planning_cem.yaml --logging.log_wandb False 2>&1 | tee /eval_cem.log`. Check that the log contains "Loaded model state from": a mistyped folder is created silently and evaluated with random weights (`main.py:86`, `training_utils.py:195-197`). If the cap allows, run the MPPI evaluation too (`planning_mppi.yaml`, logging to `/eval_mppi.log`), since that is the README's headline number; first rename `plan_eval/step-*_eval_only/` to `..._cem`, because both evaluations write to the same folder and `eval.csv` is appended to. + 8. **Copy back and tear down.** Move `/workspace/eb_jepa-lab/train.log` into the run folder, then copy the run folder to the laptop with `scp -r -P ` (an estimated 0.3 to 0.6 GB). Check that the laptop copy has `latest.pth.tar`, all twelve `e-N.pth.tar` files and both logs. Delete the pod with `runpodctl pod delete `, confirm with a fresh read that its id is absent from the RunPod `list-pods` tool, kill the watchdog, and log the pod's uptime, billed cost and the running total (the RunPod `list-pod-billing` tool; billing can lag by up to an hour, so until it catches up, log uptime × the billed `costPerHr`). + 9. **Record.** Log `success_rate`, `mean_state_dist` and `avg_episode_time` from `plan_eval/step-*_eval_only*/eval.csv`, the unroll MSE from `unroll_eval/step-*_eval_only/eval.csv`, the per-episode outcomes (from the "GIF saved to …/ep_k/agent_steps_succ|fail.gif" lines of `eval_cem.log`, which are also the GIF file names), the run folder, the GPU, the wall-clock time and the seed. +- **Success criteria:** + - The CEM evaluation on 20 random-wall episodes reports a success rate of at least 80 %. + - The laptop copy of the run folder has `config.train.yaml`, `latest.pth.tar`, the twelve `e-N.pth.tar` files, `train.log` and `eval_cem.log`, and nothing from it is in the repository. + - None of the pod ids this plan created remain in `list-pods`, and the total RunPod spend is within the $25 cap. Pods this plan did not create are never touched. +- **Verification:** + - `cat /plan_eval/step-*_eval_only*/eval.csv /unroll_eval/step-*_eval_only/eval.csv` + - The pod ids this plan created are absent from the RunPod `list-pods` tool; the log has each pod's cost and the running total. + - `git -C ~/Workspace/jepa-scope status --short` shows no lab files. +- **Escalate if:** + - Success is below 80 %: report the loss curves from the log and the unroll MSE, and propose one of: a second seed or the fixed-wall task, within the cap. Do not tune hyperparameters. + - Training stops with an error that one retry does not clear (two strikes), or the pod is lost twice. + - Any step would take the running total past $25. + +### T3 — Draft the protocol: spec, typed module, reference CEM and a CI test on a toy world model + +- **Status:** Done +- **Dependencies:** T1 (so the draft is written knowing the run exists and the adapter target is fixed) +- **Parallel group:** A +- **Complexity:** STRONG +- **Execution:** Executor (Opus) on the branch, or the lead if it is idle while training runs +- **Owner time:** None until review +- **Expected files/workspaces:** `src/jepa_scope/planning/__init__.py`, `src/jepa_scope/planning/protocol.py`, `src/jepa_scope/planning/cem.py`, `docs/planning-protocol.md`, `tests/test_planning_protocol.py`, `.gitignore` (add `*.pth.tar` and `*.gif` as a backstop) +- **Objective:** Write the protocol as `typing.Protocol` classes and free functions, torch-only, with full docstrings. `WorldModel` has `latent_shape: tuple[int, ...]`, `action_dim: int`, `encode(obs: Tensor[B, *obs]) -> Tensor[B, *latent]`, `predict(z: Tensor[B, *latent], a: Tensor[B, action_dim]) -> Tensor[B, *latent]` (one step, stateless) and `energy(z: Tensor[B, *latent], z_goal: Tensor[B, *latent]) -> Tensor[B]` (non-negative, lower is better, unitless and comparable only within one model). `rollout(model, z0, actions: Tensor[B, T, A]) -> Tensor[B, T + 1, *latent]` includes `z0` at index 0. `plan_cost(model, z0, actions, z_goal) -> Tensor[B]` sums `energy` over all T + 1 latents including index 0; that is eb_jepa's `sum_all_diffs` objective and the spec says why the sum beats the last-step cost (the README's "MPPI last" ablation). `cem.py` implements the Cross-Entropy Method against the protocol only: Gaussian proposal over `[T, A]`, `num_samples`, `num_elites`, `n_iters`, an initial std (`var_scale`), an optional per-step action norm clip (`max_norm`), a `torch.Generator` for determinism, and a result holding the mean action sequence and the per-iteration best cost. Mirror eb_jepa's `CEMPlanner` semantics exactly so T4 can compare: the `planning_cem.yaml` numbers (plan length 90, 20 iterations, 200 samples, 20 elites, std 1.5, norm clip 2.45, re-plan every step and execute one action); the horizon is `min(plan_length, steps_left)`, so it shrinks over the last 90 steps (`planning.py:572-575`); the mean restarts at zero on every re-plan, with no warm start (`planning.py:578`); samples are norm-rescaled before they are scored (`planning.py:607-618`); the new mean and std are the mean and unbiased std of the elites (`planning.py:635-636`); the returned plan is the elite mean, not the best sample (`planning.py:651`). `max_norm_dims` is ignored by eb_jepa and is not reproduced. The spec `docs/planning-protocol.md` has: purpose and the roadmap pointer; the three operations with shapes and invariants; `rollout` and `plan_cost`; the reference CEM; an "Adapters" section listing the eb_jepa Two Rooms adapter (filled in by T4); a "Checked against" section with the reproduction steps, the measured numbers and their roughly ±7-point sampling error, and the note that success is measured after the last allowed step rather than on first arrival (filled in by T2 and T4); and "Open points": context-window predictors (V-JEPA 2-AC), stochastic predictors, energy calibration across models (jepa-scope's surprise floor is the precedent), goal batching, and action-space specs. The CI test defines a point-mass model (obs = z ∈ R², `predict` adds the action clipped to a max norm, `energy` is the squared distance) and checks that CEM from the origin reaches a goal at (5, 5) within 0.5 units in 12 steps, that `plan_cost` of the planned sequence is below that of random actions, and that two runs with the same generator seed give the same actions. It must finish in under 5 seconds on CPU. +- **Success criteria:** + - `uv run pytest -q tests/test_planning_protocol.py` passes; the full suite and lint pass. + - `docs/planning-protocol.md` is readable by someone who has never opened eb_jepa and tells them what to implement to add a model. + - The shapes in the spec, the module and the test agree. +- **Verification:** + - `uv run ruff check src tests && uv run ruff format --check src tests && uv run pytest -q` + - `python3 .claude/workflow/check.py` +- **Escalate if:** + - The design needs a second operation beyond the three (for example a separate history argument) to make CEM work on the toy or on the eb_jepa shapes: stop and bring the question to the lead rather than widening the protocol silently. + +### T4 — Wrap the trained model in the protocol and check parity and planning success through it + +- **Status:** Done +- **Dependencies:** T2, T3 +- **Parallel group:** Sequential +- **Complexity:** STRONG +- **Execution:** Lead session or Executor (Opus); the opt-in test runs from the lab venv on the owner's machine +- **Owner time:** None +- **Expected files/workspaces:** `src/jepa_scope/planning/eb_jepa_two_rooms.py`, `tests/test_eb_jepa_two_rooms.py`, `docs/planning-protocol.md` (the Adapters and Checked-against sections) +- **Objective:** Implement `EbJepaTwoRooms` as described in the approach, with `from_run(folder: Path, device="cpu")` reading `config.yaml` and `latest.pth.tar` (or a named `e-N.pth.tar`), importing eb_jepa lazily and raising a clear `ImportError` naming the lab setup if it is absent. Install jepa-scope into the lab venv with `uv pip install -e ~/Workspace/jepa-scope --no-deps`. Write `tests/test_eb_jepa_two_rooms.py`, skipped unless `JEPA_SCOPE_TWO_ROOMS_CHECKPOINT` is set, with four tests. + + - **Parity.** The reference is independent of the adapter: a full `JEPA(encoder, nn.Identity(), predictor, regularizer, SquareLossSeq())` constructed exactly as `main.py:174-218` does and loaded with eb_jepa's own `load_checkpoint(..., strict=True)`. (The adapter's encoder-and-predictor module cannot serve, because `unroll` is defined on `JEPA`, not `JEPAbase`, and comparing the adapter with its own modules proves nothing.) On a batch of 8 observations from `DotWall` and random actions `[8, 8, 2]`, `rollout(adapter, adapter.encode(obs), actions)[:, 1:]` must match `ref.unroll(obs5d, actions.permute(0, 2, 1), nsteps=8, unroll_mode="autoregressive", compute_loss=False)[0][:, :, 1:].flatten(3).squeeze(-1).permute(0, 2, 1)` to a max absolute difference under 1e-5, in float32 on CPU in eval mode. The permute matters: `unroll` returns `[B, 512, 9, 1, 1]` and `rollout` returns `[B, 9, 512]`, and comparing them without it compares values in a different order. `plan_cost` must match eb_jepa's real `ReprTargetDistMPCObjective(target_enc, sum_all_diffs=True)`, imported from `eb_jepa.planning` in the test (the lab venv has its plotting dependencies; the adapter itself still never imports it). + - **Paired success.** Build the environment exactly as `main.py:139-153` does, not from the dataclass defaults (which have `fix_wall=True`, `img_size=64` and `device="cuda"`; the fixed-wall branch also skips a generator draw, so every later episode would shift without any error): `env_config = init_data("two_rooms", cfg_data=["data"])[2]` and `env = DotWall(config=env_config, **["env"])` with the default generator. Call `reset()` once before the loop, as `main_eval` does (`planning.py:177-178`), so the 20 episodes are the same (wall, start, goal) tuples eb_jepa's own evaluation used in T2. Log each replayed episode's wall, start and goal, and check the replay against T2 by comparing the replayed initial and goal frames of episodes 0 and 1 with T2's `ep_0/state.pdf` and `ep_1/state.pdf` (or the first GIF frame). Run each episode the way `main_eval` does: the zero-action first step (`planning.py:211-214`), re-planning every step with the reference CEM and the `planning_cem.yaml` numbers, all 200 steps with no early stop on `done`, and success judged by `env.eval_state` (distance under 4.5) after the last step (`planning.py:249-296`). Print a per-episode table beside T2's native outcomes (read from the "GIF saved to …/ep_k/agent_steps_succ|fail.gif" lines of T2's log) and count native-only and protocol-only successes. + - **Greedy baseline.** The same 20 episodes with a policy that steps straight toward the goal's true position at norm 2.45. Start and goal are always in opposite rooms at this level, so this succeeds only when the door is roughly on the line; it measures how much the task needs planning. (Doing nothing scores 0 % by construction and would carry no information.) + - **Untrained control.** The same loop on an adapter built from the run's config with freshly initialised weights (seed 0), capped at 10 episodes for time. + + Then fill the spec's Adapters section (how the three operations map onto eb_jepa, including the post-LayerNorm latent fed back as the GRU state, the shared LayerNorm and the normaliser) and Checked-against section (the T2 native numbers, the parity result, the paired table, the baseline and the control, the machine and wall-clock times, and the exact commands). +- **Success criteria (fixed before anything runs; never tuned afterwards):** + - Parity holds under 1e-5. + - On the 20 paired episodes, the protocol planner succeeds on at least 70 %, and native-only successes minus protocol-only successes is at most 2. On a marginal 3 or 4, both planners are rerun on 50 episodes from the same generator before any conclusion, and the margin must be at most 4. Under equally good planners those two counts are symmetric however noisy the outcomes are. + - The untrained adapter scores at most 20 %. The greedy baseline's rate is reported; it is not a pass condition, but a greedy rate within 10 points of the protocol's is an escalation. + - `uv run pytest -q` in the jepa-scope venv still passes with the new test skipped, and CI does not run it. +- **Verification:** + - `JEPA_SCOPE_TWO_ROOMS_CHECKPOINT= ~/Workspace/eb_jepa-lab/.venv/bin/python -m pytest tests/test_eb_jepa_two_rooms.py -s` + - `uv run pytest -q -rs` shows the test as skipped. +- **Escalate if:** + - Parity holds but the success rate misses the band: the planner or the episode loop differs from eb_jepa's, not the protocol; compare against `main_eval` step by step before changing the protocol. + - Parity fails: the adapter's `predict` or normalisation is wrong; do not adjust tolerances. + - The replayed frames of episodes 0 and 1 differ from T2's `state.pdf` images (most likely an environment config that is not the one `main.py` builds): stop and fix the replay before reading any rate, since an unpaired comparison is not evidence. + - The untrained control scores above 20 %, or the greedy baseline comes within 10 points of the protocol: the task is easier than the README implies and the success criterion is not evidence; report before writing numbers into the spec. + - `weights_only=True` refuses the checkpoint for a reason an allowlist of plain library types does not fix. + - The 50-episode contingency projects beyond about 12 hours on the laptop (eb_jepa's native CEM re-encodes the start observation for every sample, so it is slow on a CPU): ask before running it. + +### T5 — Close out: roadmap, decisions, plan status, review, owner summary + +- **Status:** Done +- **Dependencies:** T4 +- **Parallel group:** Sequential +- **Complexity:** ROUTINE +- **Execution:** Lead session, with `reviewer` (Opus, fresh context) on the diff and this plan +- **Owner time:** 0.5 h to read the spec's numbers and the summary, then merge (owner-only) +- **Expected files/workspaces:** `docs/roadmap.md`, `docs/decisions.md`, `docs/decisions/product.md`, `docs/decisions/method.md`, `docs/decisions/process.md`, `docs/plans/README.md`, this plan and its log +- **Objective:** Mark the Week 1 row "Done" with a one-line pointer to the spec and the measured rate; add a lesson under "What Week 1 taught" only if one emerged (the protocol-versus-native parity check is a candidate: a wrapper that runs is not a wrapper that is right until its outputs are compared to the original's). Replace the roadmap's "Next" list with the owner's order (Weeks 3 and 4 next) and the note that Week 2's scope landed in Week 1. Add decision lines that meet the bar: under Product, where the protocol draft lives and that it moves to latent-plan; under Method, the protocol's shape and cost definition as checked; under Process, that milestones follow the roadmap's weekly order, and that cloud GPU spend within a cap the owner approved in a plan needs no further approval. Update the plan's `Status:` line, remove it from the active list, compress the log into the completion summary, delete the log file, run `/checkpoint`, run the reviewer, fix material findings, and write the owner summary with the numbers, the RunPod spend and the gates taken (plan approval, T3's signatures, the cap if it was reached, and merge). +- **Success criteria:** + - `python3 .claude/workflow/check.py` passes; CI is green on the PR. + - The reviewer's verdict is Pass or its material findings are fixed and checks rerun. +- **Verification:** + - `python3 .claude/workflow/check.py`; the CI run on the pull request. +- **Escalate if:** + - The reviewer disputes the protocol's shape: resolve with `architect` before the owner reviews; do not merge a contested contract. + +## Execution strategy + +- **Lead/orchestrator owns:** the lab directory and every run in it, the owner gates, integration, the roadmap and decisions edits. +- **Delegated tasks:** T3 to an executor (Opus) because it is self-contained and runs while the pod trains; T4 to an executor only if the lead's context is large by then, since it needs the lab venv and the owner's machine either way. +- **Model routing:** T1 ROUTINE, T2 ROUTINE, T3 STRONG, T4 STRONG, T5 ROUTINE with the reviewer at its default (Opus, high). +- **Owner time, total:** about 45 minutes of attention (approval 0.25 h, review and merge 0.5 h). RunPod spend is capped at $25. +- **Approval gates:** + - Owner approves this plan, which also approves installing and running eb_jepa in the lab directory on the laptop and on RunPod pods, with RunPod spend up to $25 in total. + - Anything that would pass the $25 cap is a separate owner approval with a quoted price. + - Owner approves T3's first small change (the protocol module's signatures) before the executor writes the rest, per the milestone rule. + - Merging is owner-only. +- **Review strategy:** the reviewer runs once at T5 on the whole diff and the plan, with the spec's numbers and the test output in front of it. If T3 lands long before T4, a quick reviewer pass on the protocol module alone is allowed so T4 does not build on a shape that will change. + +### Concurrency and isolation + +| Task | Actually concurrent with | Worker/session | Writes files? | Branch/base | Isolation | Shared mutable resources | Integration | +|---|---|---|---:|---|---|---|---| +| T1 | None | Lead | Lab only | `main` at `4ab37dd` (read-only in the repo) | External environment | CPU | Report only | +| T2 | T3 | Lead (RunPod pod; laptop fallback) | Lab only | n/a | External environment (RunPod pod) | The pod; laptop CPU on the fallback | Artifact handoff (checkpoint copied to the lab) | +| T3 | T2 | Executor | Yes | `week1-two-rooms-protocol` from `main` | Main checkout, single writer (a worktree if the lead also edits) | None | Direct | +| T4 | None | Lead or Executor | Yes | `week1-two-rooms-protocol` | Main checkout | Lab venv, checkpoint | Direct | +| T5 | None | Lead | Yes | `week1-two-rooms-protocol` | Main checkout | None | PR to `main` | + +- **T2 and T3 timing:** Concurrent. Training is wall-clock bound and touches nothing in the repository; T3 is the only writer in the checkout. The lead does not edit repository files while the executor holds the checkout; if that becomes necessary, the executor gets a worktree (`uv sync` there first) and integrates by cherry-pick. T3's inputs are tracked (nothing from the lab is needed). T2's output is a path handed to T4. +- **Single-instance resources:** the RunPod pod belongs to T2 alone, and this plan has at most one pod at a time. On the laptop fallback, the CPU is shared by training and by the executor's test runs, which only slows both, and the Apple GPU is claimed only on the MPS path, during which no jepa-scope embedding runs. No port is used. + +## Execution log + +Kept in `docs/plans/2026-11-05-week1-two-rooms-protocol-log.md` while the work ran, then compressed into the completion summary below and deleted (git keeps it). It was kept in the template's entry format. Material items expected there: the T1 measurements and the GPU picked; for each pod, the image, GPU, tier, hourly price before creation, uptime, billed cost, deletion confirmation and the running total against the $25 cap; the T2 run folder and native numbers, any MPS attempt and its outcome, the T4 parity and success numbers, and any deviation from the planner settings. + +## Integration and verification + +### Integrated checkpoint + +- [x] `/checkpoint` completed (lint → tests → build → workflow check) +- [x] Any checkpoint failures resolved and the checkpoint rerun (the build's wheel check failed at random; fixed in 115f87c) + +### End-to-end acceptance + +- [x] eb_jepa's own CEM evaluation on the trained model reports at least 80 % on 20 random-wall episodes, recorded in the spec with the machine and time (18/20 on an L40S, 45 s per episode). +- [x] The opt-in test passes from the lab venv: parity under 1e-5 (0.0), the paired rule met (16/20, margin 2), the untrained control at most 20 % (0/10), and the greedy baseline more than 10 points below the protocol (45 % against 80 %). +- [ ] CI passes with the opt-in test skipped and the toy test running (checked on the pull request). +- [x] `docs/planning-protocol.md` can be followed by a reader without eb_jepa to implement an adapter, and by a reader with the lab to reproduce the numbers. + +### Independent review + +- **Reviewer:** `reviewer` (Opus, high), fresh context, on the diff plus this plan; two passes. +- **Verdict:** Pass. +- **Findings:** four Minor. (1) No test pinned the CEM to eb_jepa's: added a same-noise opt-in test (plans agree to 6e-7; a planted biased std fails it by 0.45). (2) The spec called the protocol the second tool; it is layer E and belongs to latent-plan, the third. (3) The spec lacked the native evaluation's machine and time, misdescribed the run folder and the normaliser, and the log overstated how often a native-favouring 2–0 split happens under equal planners. All fixed. +- [x] Material findings resolved +- [x] Relevant verification rerun after fixes + +## Completion summary + +- **Delivered:** + - eb_jepa's Two Rooms model trained on one RunPod L40S (12 epochs, about 1.6 h) and evaluated with eb_jepa's own CEM: 18 of 20 episodes (90%), against the README's 96 ± 2% averaged over seeds and checkpoints. + - A draft planning protocol (`src/jepa_scope/planning/`): the `WorldModel` protocol, `rollout`, `plan_cost` with a `"sum"` or `"last"` objective, a reference CEM that mirrors eb_jepa's, and an `eb_jepa` Two Rooms adapter; specified in `docs/planning-protocol.md`, with 35 toy tests in CI. + - Checked through the adapter on the laptop CPU. Rollouts match eb_jepa's exactly, costs to 9e-8 relative, and given the same noise the reference CEM returns eb_jepa's plans (6e-7); a planted biased std fails that check by 0.45. All 20 replayed start frames match eb_jepa's (goals not compared directly). In a closed-loop engineering check (82 minutes, code at `00e6919`), 16 of 20 were solved through the protocol against 18 of 20 natively (native only: episodes 7 and 12), passing the pre-set threshold; that rule is a net count, and the result does not establish non-inferiority. Greedy 45%; untrained 0 of 10. Exploratory: episodes 7 and 12 succeed with two other planner seeds. The short checks were rerun at `3ac29ff` (log `~/Workspace/eb_jepa-lab/t4_fast_checks_3ac29ff.log`, with the checkpoint and config hashes; the planted-deviation run is in `t4_fast_checks_4c6087e.log`). + - An outside review (ChatGPT Astra, round 1: no Critical findings, five Major, five Minor, one suggestion) led to: `latent_shape` renamed to `state_shape` (owner-approved); adapter rules in the spec (read-only inputs, independent rows, declared capacity that refuses to overflow, inference only) checked by a packed-context toy model; `plan_cost` adding energies step by step; `cem_plan` rejecting non-finite or non-integer settings, mismatched start and goal, and NaN or infinite costs, with float32 actions; the action contract and the planner's narrower scope written down; and the closed-loop result reworded as an engineering check. + - RunPod spend: one pod, about 2.08 h at $0.79/h, about $1.65 of the $25 cap; deleted and confirmed. +- **Deviations from approved plan:** + - The lead wrote T3 instead of an executor while training ran. + - `opencv-python-headless` on the pod; watchdogs as script files after a `pkill -f` matched its own SSH command; MPPI not run (optional). + - The replay check compares GIF first frames with the label masked rather than `state.pdf`. + - After a Fable architect review, the owner approved `encode(obs: Any)`, the planning-state wording, the objective option and deferring `latent(z)`. + - Added: the same-noise CEM test, and a fix for the build's random wheel-check failure. + - After the outside review: the changes listed under Delivered, and the adapter's `predict` now moves its inputs to its device and dtype. +- **Known risks / follow-up:** + - 20 episodes cannot rule out a gap of 10 points or more between the planners; under equal planners a 2–0 split favouring the native planner happens a quarter of the time. Not yet checked: a closed-loop comparison on one device with shared candidate noise, the replayed goals, a packed state's capacity on a real model, and a bounded-action planner. + - The protocol has been checked against one model and stays experimental. V-JEPA 2-AC (a planning state with the token context, actions and poses so far; final-frame L1; box-clipped actions) is the next test of its shape and must match its native rollout; an optional `latent(z)` is to be decided when latent-plan is planned. + - The lab (`~/Workspace/eb_jepa-lab`, about 0.5 GB of checkpoints) stays outside the repository. +- **Decision records updated:** `docs/decisions.md`, Method (the protocol's shape), Product (where the draft lives), Process (milestone order; cloud GPUs within an approved cap). +- **Lesson:** cue: a remote shell command that both kills and starts a process by pattern / lesson: send remote helper scripts as files (`ssh host 'cat > x.sh'`) and kill by PID or session, never `pkill -f` on a pattern that also appears in the SSH command itself; this belongs in a RunPod or remote-compute skill if one is written / evidence: T2's first watchdog reset killed its own SSH session and the training launch with it / confidence: medium diff --git a/docs/plans/README.md b/docs/plans/README.md index 5113f5c..92f5400 100644 --- a/docs/plans/README.md +++ b/docs/plans/README.md @@ -13,6 +13,8 @@ Add a plan as soon as its file is committed (a Draft counts); remove it when its Done, Closed, Superseded or Shelved. Every other file here is a record. `.claude/workflow/check.py` fails when this list and a plan's `Status:` line disagree. +None at the moment. + ## Lifecycle - **Draft** while planning; reviewed by the owner before any code is touched. diff --git a/docs/roadmap.md b/docs/roadmap.md index 992dfe2..790a567 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -1,7 +1,8 @@ # Roadmap - **Written:** 2026-11-03, from the planning conversation that started this project. -- **Last updated:** 2026-11-05, after the Week 1 scaffold and CI landed. +- **Last updated:** 2026-11-06, after the `eb_jepa` Two Rooms run and the planning protocol draft + closed Week 1. - **Status:** living document. Executable milestone plans go in `docs/plans/`. Settled choices go in `docs/decisions.md`. @@ -117,7 +118,10 @@ shaped like familiar model APIs, such as `POST /v1/video/embeddings` and This project is a planning runtime written against the protocol in layer E, with a browser UI. A user clicks a goal and watches CEM or MPPI sample action sequences, roll them out in latent space and converge on the energy landscape. The first adapter is an `eb_jepa` Two Rooms world model -trained locally. The second is V-JEPA 2-AC, run offline on a rented GPU. +trained locally. The second is V-JEPA 2-AC, run offline on a rented GPU. The protocol draft, a +reference CEM planner and the `eb_jepa` Two Rooms adapter already exist in this repository +(`src/jepa_scope/planning/`, specified in `docs/planning-protocol.md`) and move here when this +project starts. - **Frontend leverage.** The energy over a two-dimensional slice of action space as an animated WebGL heatmap, sampled trajectories with the best ones highlighted, and the real rollout beside @@ -149,7 +153,7 @@ of GPU time for `eb_jepa` runs. ## Progress -### Week 1 (2026-11-03 to 2026-11-05) +### Week 1 (2026-11-03 to 2026-11-06) | Step | Status | | --- | --- | @@ -159,7 +163,7 @@ of GPU time for `eb_jepa` runs. | Surprise from masked prediction | Done. A cut in a spliced clip is detected; see below | | `.latent` format and CLI | Done | | Web viewer with overlay and timeline | Done | -| Run the `eb_jepa` Two Rooms example and draft the planning protocol | **Not started** | +| Run the `eb_jepa` Two Rooms example and draft the planning protocol | Done. Rollouts and same-noise plans match `eb_jepa`'s; in a closed-loop check `eb_jepa`'s planner solved 18 of 20 episodes and the protocol's 16; see `docs/planning-protocol.md` | | Publish under Apache-2.0 | Done: `ryangswe/jepa-scope`, with CI | The first local milestone was "open a clip on a MacBook without CUDA and see a PCA overlay and a @@ -177,15 +181,20 @@ within each scene and rose to between 0.65 and 0.67 for the second after the cut showed nothing useful. Natural video with motion is the minimum for judging a signal. - **Expect a floor.** Surprise never reaches zero, and what carries information is the change relative to the clip's own floor. +- **A wrapper that runs is not a wrapper that is right.** The `eb_jepa` adapter was trusted only + after its rollouts matched `eb_jepa`'s own to the last digit and, given the same random noise, + its planner returned `eb_jepa`'s plans. Closed-loop success on 20 episodes (16 against 18) was a + useful smoke test but could not show the two planners equal; an outside review pointed out that + the "paired" pass rule reduced to a net count. Controlled comparisons answer "is it the same + computation?"; small success-rate runs do not. ### Next -Choosing the next milestone is the owner's call. These are the candidates: +Milestones follow the weekly order. Week 2's scope landed in Week 1, so the next milestones are: -1. The remaining Week 1 step: run the `eb_jepa` Two Rooms example and draft the planning protocol. -2. jepa-scope Weeks 3 and 4: folder indexing, nearest-neighbour search and the latent +1. jepa-scope Weeks 3 and 4: folder indexing, nearest-neighbour search and the latent trajectory view. -3. jepa-scope Weeks 5 and 6: embedding health metrics, polish and a launch demo. +2. jepa-scope Weeks 5 and 6: embedding health metrics, polish and a launch demo. Each milestone gets a plan in `docs/plans/` before code changes, written by the `planner` agent and approved by the owner. diff --git a/scripts/build.sh b/scripts/build.sh index 0182aae..52d88bf 100755 --- a/scripts/build.sh +++ b/scripts/build.sh @@ -9,7 +9,10 @@ cd "$(dirname "$0")/.." uv build wheel=$(ls -t dist/*.whl | head -1) -if ! unzip -l "$wheel" | grep -q "jepa_scope/static/index.html"; then +# Read the listing first: piping unzip into `grep -q` fails at random under pipefail, because grep +# exits at the first match and unzip then dies writing to the closed pipe. +listing=$(unzip -l "$wheel") +if ! grep -q "jepa_scope/static/index.html" <<<"$listing"; then echo "error: $wheel does not contain the built viewer" >&2 exit 1 fi diff --git a/src/jepa_scope/planning/__init__.py b/src/jepa_scope/planning/__init__.py new file mode 100644 index 0000000..825ea23 --- /dev/null +++ b/src/jepa_scope/planning/__init__.py @@ -0,0 +1,16 @@ +"""Planning with JEPA world models: a draft protocol and a reference planner.""" + +from jepa_scope.planning.cem import ACTION_DTYPE, CEMConfig, CEMResult, cem_plan, clip_action_norm +from jepa_scope.planning.protocol import Objective, WorldModel, plan_cost, rollout + +__all__ = [ + "ACTION_DTYPE", + "CEMConfig", + "CEMResult", + "Objective", + "WorldModel", + "cem_plan", + "clip_action_norm", + "plan_cost", + "rollout", +] diff --git a/src/jepa_scope/planning/cem.py b/src/jepa_scope/planning/cem.py new file mode 100644 index 0000000..4fb2910 --- /dev/null +++ b/src/jepa_scope/planning/cem.py @@ -0,0 +1,166 @@ +"""A reference Cross-Entropy Method planner written only against the planning protocol (draft). + +It mirrors ``CEMPlanner.plan`` in ``eb_jepa/planning.py`` of ``facebookresearch/eb_jepa`` (commit +``966e61e``), so a model planned through the protocol can be compared with the same model planned +natively: + +- The horizon is ``min(plan_length, steps_left)``, so it shrinks near the end of an episode. +- The proposal mean starts at zero and the standard deviation at ``var_scale`` on every call; there + is no warm start from the previous plan. +- Each sampled action is rescaled to a norm of at most ``max_norm`` before it is scored, and the + rescaled actions are the ones the elites are drawn from. +- The new mean and standard deviation are the mean and the unbiased standard deviation of the + elites. +- The returned plan is the final elite mean, not the best sample seen. + +eb_jepa's ``max_norm_dims`` setting is not used by its planner, and it is not reproduced here. +Executing the plan (for example, one action and then planning again) is the caller's job. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import torch +from torch import Tensor + +from jepa_scope.planning.protocol import Objective, WorldModel, plan_cost + +_NORM_EPS = 1e-6 +ACTION_DTYPE = torch.float32 +"""The dtype the planner samples and returns actions in, whatever the planning state's dtype.""" + + +@dataclass(frozen=True) +class CEMConfig: + """Settings for :func:`cem_plan`. The defaults are eb_jepa's ``planning_cem.yaml``.""" + + plan_length: int = 90 + """Longest horizon, in actions.""" + + n_iters: int = 20 + """Number of sample-and-refit iterations.""" + + num_samples: int = 200 + """Action sequences sampled per iteration.""" + + num_elites: int = 20 + """Lowest-cost samples the proposal is refitted to. At least 2, so their spread is defined.""" + + var_scale: float = 1.5 + """Initial standard deviation of every action component. Despite the name, it is not a variance.""" + + max_norm: float | None = 2.45 + """Largest allowed norm of one action, or ``None`` for no limit.""" + + objective: Objective = "sum" + """How :func:`~jepa_scope.planning.protocol.plan_cost` scores a sample: ``"sum"`` over every + state of its rollout (eb_jepa's Two Rooms setting) or the ``"last"`` state only.""" + + def __post_init__(self) -> None: + for name in ("plan_length", "n_iters", "num_samples", "num_elites"): + value = getattr(self, name) + if not isinstance(value, int) or isinstance(value, bool): + raise ValueError(f"{name} must be an integer, not {value!r}") + for name in ("var_scale", "max_norm"): + value = getattr(self, name) + if value is not None and not math.isfinite(value): + raise ValueError(f"{name} must be finite, not {value!r}") + if self.plan_length < 1 or self.n_iters < 1: + raise ValueError("plan_length and n_iters must be at least 1") + if not 2 <= self.num_elites <= self.num_samples: + raise ValueError("num_elites must be at least 2 and at most num_samples") + if self.var_scale <= 0: + raise ValueError("var_scale must be positive") + if self.max_norm is not None and self.max_norm <= 0: + raise ValueError("max_norm must be positive or None") + if self.objective not in ("sum", "last"): + raise ValueError("objective must be 'sum' or 'last'") + + +@dataclass(frozen=True) +class CEMResult: + """What :func:`cem_plan` returns.""" + + actions: Tensor + """The planned action sequence ``[T, A]``: the elite mean after the last iteration.""" + + best_costs: Tensor + """The lowest sampled ``plan_cost`` in each iteration, ``[n_iters]``, on the CPU.""" + + +def clip_action_norm(actions: Tensor, max_norm: float) -> Tensor: + """Rescale each action (the last dimension) to a norm of at most ``max_norm``. + + This is eb_jepa's formula, ``min(norm, max_norm) / (norm + 1e-6)``, so actions under the limit + also shrink by a factor of ``norm / (norm + 1e-6)``. + """ + norms = actions.norm(dim=-1, keepdim=True) + return actions * (norms.clamp(max=max_norm) / (norms + _NORM_EPS)) + + +@torch.no_grad() +def cem_plan( + model: WorldModel, + z0: Tensor, + z_goal: Tensor, + config: CEMConfig | None = None, + *, + steps_left: int | None = None, + generator: torch.Generator | None = None, +) -> CEMResult: + """Plan an action sequence from ``z0`` towards ``z_goal`` with the Cross-Entropy Method. + + ``z0`` and ``z_goal`` are single planning states, ``[1, *state]``, on the same device. Each + sample is scored with :func:`~jepa_scope.planning.protocol.plan_cost` using + ``config.objective``. ``steps_left``, when given, caps the horizon at the steps remaining in + the episode. Pass a ``generator`` on ``z0``'s device for a reproducible plan. + + Actions are continuous vectors, sampled and returned in ``ACTION_DTYPE`` on ``z0``'s device, + and limited only by ``config.max_norm``. The returned plan is the elite mean, which stays + within the norm bound because the bound is convex; converting it into environment actions is + the caller's job. A rollout that scores a candidate as NaN or infinite raises ``ValueError``. + """ + config = config or CEMConfig() + if z0.shape[0] != 1 or z_goal.shape[0] != 1: + raise ValueError("cem_plan plans for one start state: z0 and z_goal must have batch size 1") + if z0.shape != z_goal.shape or tuple(z0.shape[1:]) != tuple(model.state_shape): + raise ValueError( + f"z0 {tuple(z0.shape)} and z_goal {tuple(z_goal.shape)} must both be [1, *state_shape], " + f"with state_shape {tuple(model.state_shape)}" + ) + if z0.device != z_goal.device: + raise ValueError(f"z0 is on {z0.device} but z_goal is on {z_goal.device}") + if steps_left is not None and steps_left < 1: + raise ValueError("steps_left must be at least 1") + horizon = config.plan_length if steps_left is None else min(config.plan_length, steps_left) + + device = z0.device + mean = torch.zeros(horizon, model.action_dim, dtype=ACTION_DTYPE, device=device) + std = torch.full((horizon, model.action_dim), config.var_scale, dtype=ACTION_DTYPE, device=device) + starts = z0.expand(config.num_samples, *z0.shape[1:]) + goals = z_goal.expand(config.num_samples, *z_goal.shape[1:]) + + best_costs = [] + for _ in range(config.n_iters): + noise = torch.randn( + config.num_samples, horizon, model.action_dim, generator=generator, dtype=ACTION_DTYPE, device=device + ) + actions = mean + std * noise # [num_samples, T, A] + if config.max_norm is not None: + actions = clip_action_norm(actions, config.max_norm) + + costs = plan_cost(model, starts, actions, goals, config.objective) + if costs.shape != (config.num_samples,): + raise ValueError(f"energy must return one cost per candidate, got shape {tuple(costs.shape)}") + if not torch.isfinite(costs).all(): + bad = int((~torch.isfinite(costs)).sum()) + raise ValueError(f"{bad} of {config.num_samples} candidate costs are NaN or infinite") + best_costs.append(costs.min().item()) + + elites = actions[torch.topk(costs, config.num_elites, largest=False).indices] + mean = elites.mean(dim=0) + std = elites.std(dim=0) + + return CEMResult(actions=mean, best_costs=torch.tensor(best_costs)) diff --git a/src/jepa_scope/planning/eb_jepa_two_rooms.py b/src/jepa_scope/planning/eb_jepa_two_rooms.py new file mode 100644 index 0000000..d68f871 --- /dev/null +++ b/src/jepa_scope/planning/eb_jepa_two_rooms.py @@ -0,0 +1,142 @@ +"""An adapter that plans with eb_jepa's Two Rooms world model through the planning protocol (draft). + +The model is the action-conditioned JEPA in ``examples/ac_video_jepa`` of +``facebookresearch/eb_jepa`` (checked at commit ``966e61e``): an Impala encoder whose output is a +512-dimensional latent, and a one-layer GRU predictor. The three operations map onto it as +follows. + +- ``encode`` applies eb_jepa's Two Rooms ``Normalizer`` to a raw ``[B, 2, 65, 65]`` observation, + runs the encoder on it as a one-frame clip and flattens ``[B, 512, 1, 1, 1]`` to ``[B, 512]``. +- ``predict`` runs one GRU step with the action as input and the latent as the hidden state, then + the encoder's final LayerNorm, which the predictor shares. The normalised output is the next + latent, and eb_jepa feeds it back as the next hidden state, so the latent carries everything + the next step needs. This is one step of ``JEPA.unroll`` in autoregressive mode. +- ``energy`` is the mean squared difference over the 512 dimensions, which is eb_jepa's + ``ReprTargetDistMPCObjective`` for one state. (eb_jepa's README writes this cost with an L2 + norm, but its code takes the mean of squares, and this adapter follows the code.) + +eb_jepa is not a dependency of jepa-scope. It is imported when an adapter is built, and it must be +installed in the same environment (see ``docs/planning-protocol.md``). +""" + +from __future__ import annotations + +from pathlib import Path + +import torch +from torch import Tensor, nn + +IMG_SIZE = 65 +LATENT_DIM = 512 +ACTION_DIM = 2 + +_INSTALL_HINT = ( + "The eb_jepa Two Rooms adapter needs eb_jepa (github.com/facebookresearch/eb_jepa, commit " + "966e61e) installed in this environment. See 'Adapters' in docs/planning-protocol.md." +) + + +def _import_eb_jepa(): + try: + from eb_jepa.architectures import ImpalaEncoder, RNNPredictor + from eb_jepa.datasets.two_rooms.normalizer import Normalizer + except ImportError as error: + raise ImportError(_INSTALL_HINT) from error + return ImpalaEncoder, RNNPredictor, Normalizer + + +def _read_model_config(folder: Path) -> dict: + import yaml # eb_jepa depends on PyYAML + + # eb_jepa's evaluation rewrites config.yaml; config.train.yaml, when present, is the + # training-time copy. Their model sections are the same. + for name in ("config.train.yaml", "config.yaml"): + path = folder / name + if path.exists(): + with path.open() as f: + return yaml.safe_load(f)["model"] + raise FileNotFoundError(f"No config.yaml in {folder}") + + +class _Modules(nn.Module): + """The encoder and predictor, under the names they have inside eb_jepa's ``JEPA``.""" + + def __init__(self, model_cfg: dict): + super().__init__() + impala_encoder, rnn_predictor, _ = _import_eb_jepa() + dobs = model_cfg["dobs"] + # As examples/ac_video_jepa/main.py builds them. + self.encoder = impala_encoder( + width=1, + stack_sizes=(16, model_cfg["henc"], model_cfg["dstc"]), + num_blocks=2, + dropout_rate=None, + layer_norm=False, + input_channels=dobs, + final_ln=True, + mlp_output_dim=LATENT_DIM, + input_shape=(dobs, IMG_SIZE, IMG_SIZE), + ) + self.predictor = rnn_predictor(hidden_size=LATENT_DIM, final_ln=self.encoder.final_ln) + + +class EbJepaTwoRooms: + """eb_jepa's Two Rooms world model behind the planning protocol. + + Build it with :meth:`from_run` for trained weights or :meth:`untrained` for a control. + """ + + state_shape = (LATENT_DIM,) + action_dim = ACTION_DIM + + def __init__(self, modules: nn.Module, device: str | torch.device = "cpu"): + _, _, normalizer = _import_eb_jepa() + self.device = torch.device(device) + self.modules = modules.to(self.device).eval() + self.normalizer = normalizer() + + @classmethod + def from_run( + cls, folder: str | Path, checkpoint: str = "latest.pth.tar", device: str | torch.device = "cpu" + ) -> EbJepaTwoRooms: + """Load a run folder written by eb_jepa's ``examples/ac_video_jepa/main.py``. + + Only the encoder and predictor weights are read. The checkpoint is loaded with + ``weights_only=True`` and strict key matching, so a file that is not a plain state dict, + or a key layout that does not match, raises instead of leaving weights random. + """ + folder = Path(folder) + modules = _Modules(_read_model_config(folder)) + saved = torch.load(folder / checkpoint, weights_only=True, map_location="cpu") + state = {} + for key, value in saved["model_state_dict"].items(): + key = key.removeprefix("_orig_mod.") # added by torch.compile + if key.startswith(("encoder.", "predictor.")): + state[key] = value + modules.load_state_dict(state, strict=True) + return cls(modules, device) + + @classmethod + def untrained(cls, folder: str | Path, seed: int = 0, device: str | torch.device = "cpu") -> EbJepaTwoRooms: + """The same architecture as the run in ``folder``, with freshly initialised weights.""" + with torch.random.fork_rng(devices=[]): + torch.manual_seed(seed) + modules = _Modules(_read_model_config(Path(folder))) + return cls(modules, device) + + @torch.no_grad() + def encode(self, obs: Tensor) -> Tensor: + """Encode raw observations ``[B, 2, 65, 65]`` (as ``DotWall`` renders them) into ``[B, 512]``.""" + obs = self.normalizer.normalize_state(obs.to(self.device, torch.float32)) + return self.modules.encoder(obs.unsqueeze(2)).flatten(1) + + @torch.no_grad() + def predict(self, z: Tensor, a: Tensor) -> Tensor: + """One GRU step: ``z`` ``[B, 512]`` and action ``a`` ``[B, 2]`` to the next latent ``[B, 512]``.""" + z = z.to(self.device, torch.float32) + a = a.to(self.device, torch.float32) + return self.modules.predictor(z.reshape(z.shape[0], LATENT_DIM, 1, 1, 1), a.unsqueeze(-1)).flatten(1) + + def energy(self, z: Tensor, z_goal: Tensor) -> Tensor: + """Mean squared difference over the latent dimensions, ``[B]``.""" + return ((z - z_goal) ** 2).mean(dim=1) diff --git a/src/jepa_scope/planning/protocol.py b/src/jepa_scope/planning/protocol.py new file mode 100644 index 0000000..b28da91 --- /dev/null +++ b/src/jepa_scope/planning/protocol.py @@ -0,0 +1,109 @@ +"""A model-agnostic protocol for planning with a JEPA world model (draft). + +A world model that implements three operations can be planned with, whatever its architecture: + +- ``encode(obs) -> z`` turns an observation into a planning state. +- ``predict(z, a) -> z'`` predicts the planning state after one action. +- ``energy(z, z_goal)`` scores how far a planning state is from a goal; lower is better. + +The planning state ``z`` is a single tensor of fixed shape (``model.state_shape``). For many +models it is simply the latent. For others it also carries what the next prediction needs, such +as earlier predicted frames or a robot's pose, and ``energy`` decides which part it compares. + +Planners in this package are written only against these operations, so a new model needs an +adapter, not a new planner. This is a draft: the signatures may change without notice until the +protocol has been checked against more than one model. See ``docs/planning-protocol.md``, which +also states the rules an adapter must follow (read-only inputs, independent batch rows, and so +on). + +Shapes use ``B`` for the batch, ``T`` for the number of actions, ``A`` for the action size, and +``*state`` for the shape of one planning state. +""" + +from __future__ import annotations + +from typing import Any, Literal, Protocol, runtime_checkable + +import torch +from torch import Tensor + +Objective = Literal["sum", "last"] +"""How :func:`plan_cost` combines the energies along a rollout.""" + + +@runtime_checkable +class WorldModel(Protocol): + """What a planner needs from a world model. + + The methods must treat their input tensors as read-only, keep each batch row independent of + the others, and accept expanded or non-contiguous inputs. A state they returned earlier must + not change later. They are called for inference only, without gradients. + """ + + state_shape: tuple[int, ...] + """Shape of one planning state, without the batch dimension.""" + + action_dim: int + """Size of one action vector.""" + + def encode(self, obs: Any) -> Tensor: + """Encode a batch of ``B`` observations into planning states ``[B, *state]``. + + The observation type is the model's own (an image tensor, or an image with a robot pose, + for example). Planners pass observations through without looking inside them. + """ + ... + + def predict(self, z: Tensor, a: Tensor) -> Tensor: + """Predict the planning states ``[B, *state]`` one action ``[B, A]`` after ``z`` ``[B, *state]``. + + One step and stateless: everything the model needs to predict must be in ``z``. + """ + ... + + def energy(self, z: Tensor, z_goal: Tensor) -> Tensor: + """Score planning states ``[B, *state]`` against goals ``[B, *state]``, giving ``[B]``. + + Non-negative, lower is better, and zero when the parts of ``z`` that the model compares + equal those of ``z_goal``. Each row's score depends only on that row. The scale is the + model's own: energies are comparable within one model, never across models. + """ + ... + + +def rollout(model: WorldModel, z0: Tensor, actions: Tensor) -> Tensor: + """Apply ``predict`` once per action and keep every state. + + ``z0`` is ``[B, *state]`` and ``actions`` is ``[B, T, A]``. Returns ``[B, T + 1, *state]``, + with ``z0`` at index 0 and the state after action ``t`` at index ``t + 1``. Planners do not + need the whole trajectory; :func:`plan_cost` does not keep it. + """ + states = [z0] + z = z0 + for t in range(actions.shape[1]): + z = model.predict(z, actions[:, t]) + states.append(z) + return torch.stack(states, dim=1) + + +def plan_cost(model: WorldModel, z0: Tensor, actions: Tensor, z_goal: Tensor, objective: Objective = "sum") -> Tensor: + """Cost of an action sequence, from the ``energy`` of its rollout's states to the goal. + + ``z0`` is ``[B, *state]``, ``actions`` is ``[B, T, A]`` and ``z_goal`` is ``[B, *state]``. + Returns ``[B]``. + + With ``objective="sum"`` (the default) the cost is the sum of the energies of all ``T + 1`` + states, including ``z0``, added one step at a time in the dtype ``energy`` returns. It gives + every step of the plan a cost, and it is the objective of eb_jepa's Two Rooms planner. With + ``objective="last"`` it is the energy of the final state only, which other native planners + use. Only the current state is kept, so memory does not grow with ``T``. + """ + if objective not in ("sum", "last"): + raise ValueError(f"objective must be 'sum' or 'last', not {objective!r}") + z = z0 + total = model.energy(z, z_goal) if objective == "sum" else None + for t in range(actions.shape[1]): + z = model.predict(z, actions[:, t]) + if total is not None: + total = total + model.energy(z, z_goal) + return total if total is not None else model.energy(z, z_goal) diff --git a/tests/conftest.py b/tests/conftest.py index 693c0af..cf01eb0 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -38,12 +38,18 @@ def tiny_encoder(): return Encoder(model, processor, "test/tiny-vjepa2", torch.device("cpu"), torch.float32) -def write_video(path, frames: np.ndarray, fps: int = 24) -> None: - """Write uint8 ``[T, H, W, 3]`` frames to an mp4 with PyAV.""" +def write_video(path, frames: np.ndarray, fps: int = 24, rotation: int | None = None) -> None: + """Write uint8 ``[T, H, W, 3]`` frames to an mp4 with PyAV. + + ``rotation`` (degrees counter-clockwise) adds a display matrix, as a phone does. It is the + same matrix ``ffmpeg -display_rotation`` writes. + """ with av.open(str(path), "w") as out: stream = out.add_stream("libx264", rate=fps) stream.width, stream.height = frames.shape[2], frames.shape[1] stream.pix_fmt = "yuv420p" + if rotation is not None: + stream.set_display_rotation(rotation) for f in frames: for packet in stream.encode(av.VideoFrame.from_ndarray(f, format="rgb24")): out.mux(packet) diff --git a/tests/test_eb_jepa_two_rooms.py b/tests/test_eb_jepa_two_rooms.py new file mode 100644 index 0000000..e84b27a --- /dev/null +++ b/tests/test_eb_jepa_two_rooms.py @@ -0,0 +1,411 @@ +"""eb_jepa's Two Rooms model planned through the protocol, checked against eb_jepa's own planner. + +Opt-in: set ``JEPA_SCOPE_TWO_ROOMS_CHECKPOINT`` to a run folder written by eb_jepa's +``examples/ac_video_jepa/main.py`` that also holds that run's native CEM evaluation +(``plan_eval/*_eval_only``), and run from an environment with eb_jepa installed: + + JEPA_SCOPE_TWO_ROOMS_CHECKPOINT= python -m pytest tests/test_eb_jepa_two_rooms.py -s + +The pass bars were fixed before anything ran. Planning 20 episodes on a laptop CPU takes hours. +""" + +from __future__ import annotations + +import os +import time +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import pytest +import torch + +CHECKPOINT = os.environ.get("JEPA_SCOPE_TWO_ROOMS_CHECKPOINT") +pytestmark = pytest.mark.skipif( + not CHECKPOINT, reason="set JEPA_SCOPE_TWO_ROOMS_CHECKPOINT to an eb_jepa Two Rooms run folder" +) + +EPISODES = 20 +CONTROL_EPISODES = 10 +STEP_NORM = 2.45 +LABEL_ROWS = 20 # of the 130-pixel GIF frame; the label ends at row 17 + + +@dataclass +class Episode: + success: bool + distance: float + wall: tuple[float, float] + start: tuple[float, float] + goal: tuple[float, float] + first_obs: torch.Tensor + seconds: float + + +def _cfg_dir() -> Path: + import eb_jepa + + package = Path(next(iter(eb_jepa.__path__))).resolve() # a namespace package: no __file__ + return package.parent / "examples" / "ac_video_jepa" / "cfgs" + + +def _yaml(path: Path) -> dict: + import yaml + + with path.open() as f: + return yaml.safe_load(f) + + +def _make_env(): + """The evaluation environment exactly as examples/ac_video_jepa/main.py builds it.""" + from eb_jepa.datasets.two_rooms.env import DotWall + from eb_jepa.datasets.utils import init_data + + eval_cfg = _yaml(_cfg_dir() / "eval.yaml") + _, _, env_config = init_data(env_name="two_rooms", cfg_data=dict(eval_cfg.get("data", {}))) + return DotWall(config=env_config, **eval_cfg["env"]) + + +def _cem_config(): + from jepa_scope.planning import CEMConfig + + planner = _yaml(_cfg_dir() / "planning_cem.yaml")["planner"] + assert planner["num_act_stepped"] == 1 + assert planner["planning_objective"]["sum_all_diffs"] is True + (max_norm,) = planner["max_norms"] + return CEMConfig( + plan_length=planner["plan_length"], + n_iters=planner["n_iters"], + num_samples=planner["num_samples"], + num_elites=planner["num_elites"], + var_scale=planner["var_scale"], + max_norm=max_norm, + ) + + +def _run_episodes(policy, count: int) -> list[Episode]: + """Replay eb_jepa's evaluation episodes the way ``main_eval`` runs them (planning.py:176-296). + + One ``reset`` before the loop, then per episode a ``reset``, a zero-action step, all + ``n_allowed_steps`` steps with no early stop, and success judged after the last step. + """ + torch.manual_seed(0) # wall-collision noise uses torch's global generator + env = _make_env() + env.reset() + episodes = [] + for index in range(count): + started = time.time() + env.reset() + obs, _, _, _, info = env.step(np.zeros(env.action_space.shape[0])) + first_obs = obs.clone() + start = tuple(float(v) for v in info["dot_position"]) + goal = tuple(float(v) for v in info["target_position"]) + policy.reset(info["target_obs"]) + for steps_left in range(env.n_allowed_steps, 0, -1): + action = policy(obs, info, steps_left) + obs, _, _, _, info = env.step(action) + result = env.eval_state(info["target_position"], info["dot_position"]) + episode = Episode( + success=bool(result["success"]), + distance=float(result["state_dist"]), + wall=(float(env.wall_x), float(env.hole_y)), + start=start, + goal=goal, + first_obs=first_obs, + seconds=time.time() - started, + ) + print( + f" episode {index:2d}: wall x={episode.wall[0]:.0f} door y={episode.wall[1]:.0f} " + f"start=({start[0]:.1f}, {start[1]:.1f}) goal=({goal[0]:.1f}, {goal[1]:.1f}) " + f"{'success' if episode.success else 'fail '} dist={episode.distance:.2f} " + f"{episode.seconds:.0f}s", + flush=True, + ) + episodes.append(episode) + return episodes + + +class ProtocolPolicy: + """Plans with the reference CEM through the protocol and executes the first action.""" + + def __init__(self, model, seed: int = 0): + self.model = model + self.config = _cem_config() + self.generator = torch.Generator().manual_seed(seed) + + def reset(self, goal_obs): + self.z_goal = self.model.encode(goal_obs[None].float()) + + def __call__(self, obs, info, steps_left): + from jepa_scope.planning import cem_plan + + z = self.model.encode(obs[None].float()) + plan = cem_plan(self.model, z, self.z_goal, self.config, steps_left=steps_left, generator=self.generator) + return plan.actions[0].cpu().numpy().astype(np.float32) + + +class GreedyPolicy: + """Steps straight towards the goal's true position, at most ``STEP_NORM`` per step.""" + + def reset(self, goal_obs): + pass + + def __call__(self, obs, info, steps_left): + offset = (info["target_position"] - info["dot_position"]).double() + norm = float(offset.norm()) + scale = min(1.0, STEP_NORM / norm) if norm > 0 else 0.0 + return (offset * scale).cpu().numpy().astype(np.float32) + + +@pytest.fixture(scope="module") +def run() -> Path: + return Path(CHECKPOINT) + + +@pytest.fixture(scope="module") +def model(run): + from jepa_scope.planning.eb_jepa_two_rooms import EbJepaTwoRooms + + return EbJepaTwoRooms.from_run(run) + + +@pytest.fixture(scope="module") +def native_episodes(run) -> list[Path]: + """The ``ep_k`` folders of the run's own CEM evaluation, in episode order.""" + (eval_dir,) = sorted(run.glob("plan_eval/*_eval_only")) + folders = [eval_dir / f"ep_{k}" for k in range(EPISODES)] + assert all(folder.is_dir() for folder in folders) + return folders + + +def _native_success(folder: Path) -> bool: + (gif,) = list(folder.glob("agent_steps_*.gif")) + return gif.stem.endswith("succ") + + +@pytest.fixture(scope="module") +def protocol_episodes(model): + print(f"\nProtocol CEM, {EPISODES} episodes:") + return _run_episodes(ProtocolPolicy(model), EPISODES) + + +def _reference_jepa(run: Path): + """An eb_jepa ``JEPA`` built exactly as main.py:174-218 builds it, loaded by eb_jepa's own loader.""" + from eb_jepa.architectures import ImpalaEncoder, InverseDynamicsModel, RNNPredictor + from eb_jepa.jepa import JEPA + from eb_jepa.losses import SquareLossSeq, VC_IDM_Sim_Regularizer + from eb_jepa.training_utils import load_checkpoint + + cfg = _yaml(run / "config.train.yaml" if (run / "config.train.yaml").exists() else run / "config.yaml")["model"] + img_size = 65 + encoder = ImpalaEncoder( + width=1, + stack_sizes=(16, cfg["henc"], cfg["dstc"]), + num_blocks=2, + dropout_rate=None, + layer_norm=False, + input_channels=cfg["dobs"], + final_ln=True, + mlp_output_dim=512, + input_shape=(cfg["dobs"], img_size, img_size), + ) + _, f, _, h, w = encoder(torch.rand(1, cfg["dobs"], 1, img_size, img_size)).shape + predictor = RNNPredictor(hidden_size=encoder.mlp_output_dim, final_ln=encoder.final_ln) + assert not cfg["regularizer"]["use_proj"] + idm = InverseDynamicsModel(state_dim=h * w * f, hidden_dim=256, action_dim=2) + reg = cfg["regularizer"] + regularizer = VC_IDM_Sim_Regularizer( + cov_coeff=reg["cov_coeff"], + std_coeff=reg["std_coeff"], + sim_coeff_t=reg["sim_coeff_t"], + idm_coeff=reg.get("idm_coeff", 0.1), + idm=idm, + first_t_only=reg.get("first_t_only"), + projector=None, + spatial_as_samples=reg["spatial_as_samples"], + idm_after_proj=reg["idm_after_proj"], + sim_t_after_proj=reg["sim_t_after_proj"], + ) + jepa = JEPA(encoder, torch.nn.Identity(), predictor, regularizer, SquareLossSeq()) + load_checkpoint(run / "latest.pth.tar", jepa, strict=True) + return jepa.eval() + + +def test_rollout_and_cost_match_eb_jepa(run, model): + from eb_jepa.planning import ReprTargetDistMPCObjective + + from jepa_scope.planning import plan_cost, rollout + + reference = _reference_jepa(run) + env = _make_env() + obs = [] + for _ in range(9): + env.reset() + obs.append(env.step(np.zeros(2))[0]) + obs, goal_obs = torch.stack(obs[:8]).float(), obs[8][None].float() + actions = 1.5 * torch.randn(8, 8, 2, generator=torch.Generator().manual_seed(0)) + + with torch.no_grad(): + obs5d = env.normalizer.normalize_state(obs).unsqueeze(2) + native, _ = reference.unroll( + obs5d, actions.permute(0, 2, 1), nsteps=8, unroll_mode="autoregressive", compute_loss=False + ) # [8, 512, 9, 1, 1] + native_states = native.flatten(3).squeeze(-1).permute(0, 2, 1) # [8, 9, 512] + ours = rollout(model, model.encode(obs), actions) + parity = (ours - native_states).abs().max().item() + print(f"\nParity: max |protocol rollout - eb_jepa unroll| = {parity:.2e}") + assert parity < 1e-5 + + goal_5d = reference.encode(env.normalizer.normalize_state(goal_obs).unsqueeze(2)) + native_cost = ReprTargetDistMPCObjective(goal_5d, sum_all_diffs=True)(native) + z_goal = model.encode(goal_obs).expand(8, -1) + ours_cost = plan_cost(model, model.encode(obs), actions, z_goal) + print(f"Cost: max relative difference = {((ours_cost - native_cost).abs() / native_cost).max().item():.2e}") + torch.testing.assert_close(ours_cost, native_cost, rtol=1e-5, atol=1e-6) + + +class _SameNoise: + """Stands in for the ``torch`` module inside a planner, so ``torch.randn`` returns given noise.""" + + def __init__(self, noises, layout): + self._noises = iter(noises) + self._layout = layout + + def __getattr__(self, name): + return getattr(torch, name) + + def randn(self, *shape, **kwargs): + return self._layout(next(self._noises)).clone() + + +@pytest.mark.parametrize("steps_left", [None, 7]) +def test_cem_matches_eb_jepa_planner_on_the_same_noise(run, model, monkeypatch, steps_left): + """Fed the same Gaussian noise, the reference CEM returns eb_jepa's ``CEMPlanner`` plan.""" + import eb_jepa.planning as eb_planning + from eb_jepa.planning import CEMPlanner, ReprTargetDistMPCObjective + + import jepa_scope.planning.cem as cem_module + from jepa_scope.planning import CEMConfig, cem_plan + + reference = _reference_jepa(run) + env = _make_env() + env.reset() + obs, _, _, _, info = env.step(np.zeros(2)) + goal = info["target_obs"].float() + samples, length, iters, elites = 200, 25, 6, 20 + horizon = length if steps_left is None else min(length, steps_left) + generator = torch.Generator().manual_seed(123) + noises = [torch.randn(samples, horizon, 2, generator=generator) for _ in range(iters)] + + def unroll(obs_init, actions): + states, _ = reference.unroll( + obs_init.repeat(actions.shape[0], 1, 1, 1, 1), + actions, + nsteps=actions.shape[2], + unroll_mode="autoregressive", + compute_loss=False, + ) + return states + + with torch.no_grad(): + obs5d = env.normalizer.normalize_state(obs.float()).unsqueeze(0).unsqueeze(2) + goal5d = reference.encode(env.normalizer.normalize_state(goal).unsqueeze(0).unsqueeze(2)) + native_planner = CEMPlanner( + unroll, + n_iters=iters, + num_samples=samples, + plan_length=length, + action_dim=2, + var_scale=1.5, + num_elites=elites, + max_norms=[STEP_NORM], + decode_each_iteration=False, + ) + native_planner.set_objective(ReprTargetDistMPCObjective(goal5d, sum_all_diffs=True)) + # eb_jepa samples in a [T, N, A] layout; the reference CEM in [N, T, A]. + monkeypatch.setattr(eb_planning, "torch", _SameNoise(noises, lambda n: n.permute(1, 0, 2))) + native = native_planner.plan(obs5d, steps_left=steps_left).actions + monkeypatch.setattr(cem_module, "torch", _SameNoise(noises, lambda n: n)) + config = CEMConfig( + plan_length=length, n_iters=iters, num_samples=samples, num_elites=elites, max_norm=STEP_NORM + ) + ours = cem_plan(model, model.encode(obs[None].float()), model.encode(goal[None]), config, steps_left=steps_left) + difference = (ours.actions - native).abs().max().item() + print(f"\nSame-noise CEM, steps_left={steps_left}: max |protocol plan - eb_jepa plan| = {difference:.2e}") + assert ours.actions.shape == native.shape == (horizon, 2) + assert difference < 1e-5 + + +def test_replay_matches_the_native_evaluation(tmp_path, native_episodes): + """The replayed episodes are the ones eb_jepa evaluated: compare every first frame, drawn the + same way eb_jepa drew its GIFs.""" + from eb_jepa.vis_utils import save_gif + from PIL import Image + + class StandStill: + def reset(self, goal_obs): + pass + + def __call__(self, obs, info, steps_left): + return np.zeros(2, dtype=np.float32) + + def frame(path): + # The top rows hold the frame label ("Frame 1/201" there, "Frame 1/1" here); skip them. + return np.asarray(Image.open(path).convert("RGB"), dtype=np.int16)[LABEL_ROWS:] + + def differing(a, b): + return int((np.abs(a - b).max(axis=-1) > 16).sum()) + + replayed = _run_episodes(StandStill(), EPISODES) + natives, ours = [], [] + for index, (episode, folder) in enumerate(zip(replayed, native_episodes, strict=True)): + (native_gif,) = list(folder.glob("agent_steps_*.gif")) + ours_gif = tmp_path / f"ep_{index}.gif" + save_gif(episode.first_obs[None], save_path=str(ours_gif), show_frame_numbers=True, fps=20) + natives.append(frame(native_gif)) + ours.append(frame(ours_gif)) + for index in range(EPISODES): + assert differing(natives[index], ours[index]) <= 4, f"episode {index}: first frames differ" + # Negative control: pairing each replay with the next native episode must fail clearly. + shifted = [differing(natives[index + 1], ours[index]) for index in range(EPISODES - 1)] + print(f"\nReplay: shifted-by-one pairs differ in {min(shifted)} to {max(shifted)} pixels") + assert min(shifted) > 100 + + +def test_protocol_planning_matches_native_success(protocol_episodes, native_episodes): + native = [_native_success(folder) for folder in native_episodes] + ours = [episode.success for episode in protocol_episodes] + native_only = sum(n and not o for n, o in zip(native, ours, strict=True)) + protocol_only = sum(o and not n for n, o in zip(native, ours, strict=True)) + print(f"\n{'episode':>8} {'native':>7} {'protocol':>9}") + for index, (n, o) in enumerate(zip(native, ours, strict=True)): + print(f"{index:>8} {'success' if n else 'fail':>7} {'success' if o else 'fail':>9}") + print( + f"Native {sum(native)}/{EPISODES}, protocol {sum(ours)}/{EPISODES}; " + f"native only {native_only}, protocol only {protocol_only}; " + f"mean episode time {np.mean([e.seconds for e in protocol_episodes]):.0f}s" + ) + assert sum(ours) >= 0.7 * EPISODES + margin = native_only - protocol_only + if margin in (3, 4): + pytest.fail(f"Marginal: native-only minus protocol-only is {margin}. Rerun both on 50 episodes.") + assert margin <= 2 + + +def test_greedy_baseline_is_well_below_planning(protocol_episodes): + print("\nGreedy baseline:") + greedy = _run_episodes(GreedyPolicy(), EPISODES) + greedy_rate = sum(e.success for e in greedy) / EPISODES + protocol_rate = sum(e.success for e in protocol_episodes) / EPISODES + print(f"Greedy {greedy_rate:.0%}, protocol {protocol_rate:.0%}") + assert greedy_rate < protocol_rate - 0.10, "the task needs less planning than assumed: escalate" + + +def test_untrained_model_fails(run): + from jepa_scope.planning.eb_jepa_two_rooms import EbJepaTwoRooms + + print(f"\nUntrained control, {CONTROL_EPISODES} episodes:") + episodes = _run_episodes(ProtocolPolicy(EbJepaTwoRooms.untrained(run, seed=0)), CONTROL_EPISODES) + rate = sum(e.success for e in episodes) / CONTROL_EPISODES + print(f"Untrained {rate:.0%}") + assert rate <= 0.20 diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index b1d5c5c..b32aab5 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -137,20 +137,10 @@ def test_server_endpoints(latent, moving_square_video): def test_rotated_video_is_decoded_upright(tmp_path): """Phone clips carry a display matrix; frames must come out the way a player shows them.""" - import shutil - import subprocess - - if not shutil.which("ffmpeg"): - pytest.skip("needs the ffmpeg CLI to write rotation metadata") frames = np.zeros((16, 64, 96, 3), np.uint8) frames[:, :16, :16] = 255 # a white block in the top-left corner of the stored frame - plain = tmp_path / "plain.mp4" - write_video(plain, frames) rotated = tmp_path / "rotated.mp4" - subprocess.run( - ["ffmpeg", "-v", "error", "-y", "-display_rotation", "90", "-i", str(plain), "-c", "copy", str(rotated)], - check=True, - ) + write_video(rotated, frames, rotation=90) v = decode(rotated, sample_fps=8) assert v.rotation == 90 assert (v.width, v.height) == (64, 96) diff --git a/tests/test_planning_protocol.py b/tests/test_planning_protocol.py new file mode 100644 index 0000000..a77d6f9 --- /dev/null +++ b/tests/test_planning_protocol.py @@ -0,0 +1,307 @@ +"""The planning protocol and the reference CEM, checked on a point mass where the answer is known.""" + +import dataclasses + +import pytest +import torch + +from jepa_scope.planning import ACTION_DTYPE, CEMConfig, WorldModel, cem_plan, clip_action_norm, plan_cost, rollout + +GOAL = (5.0, 5.0) +STEPS = 12 +SPEED = 1.0 # largest distance the point moves per step + + +class PointMass: + """The latent is the position in R², each action moves it by at most ``SPEED``, and the + energy is the squared distance to the goal.""" + + state_shape = (2,) + action_dim = 2 + + def encode(self, obs: torch.Tensor) -> torch.Tensor: + return obs + + def predict(self, z: torch.Tensor, a: torch.Tensor) -> torch.Tensor: + return z + clip_action_norm(a, SPEED) + + def energy(self, z: torch.Tensor, z_goal: torch.Tensor) -> torch.Tensor: + return ((z - z_goal) ** 2).sum(dim=-1) + + +SMALL = CEMConfig(plan_length=STEPS, n_iters=10, num_samples=100, num_elites=10, var_scale=1.5, max_norm=SPEED) + + +def _generator(seed: int = 0) -> torch.Generator: + return torch.Generator().manual_seed(seed) + + +def test_point_mass_implements_the_protocol(): + assert isinstance(PointMass(), WorldModel) + + +def test_rollout_keeps_the_start_state_and_one_state_per_action(): + z0 = torch.zeros(3, 2) + actions = torch.tensor([[1.0, 0.0], [0.0, 1.0]]).expand(3, 2, 2) + states = rollout(PointMass(), z0, actions) + assert states.shape == (3, 3, 2) + assert torch.equal(states[:, 0], z0) + assert torch.allclose(states[0, 2], torch.tensor([1.0, 1.0]), atol=1e-5) + + +def test_plan_cost_sums_energy_over_every_state_including_the_start(): + model = PointMass() + z0 = torch.zeros(1, 2) + z_goal = torch.tensor([[3.0, 4.0]]) # energy 25 from the origin + standing_still = torch.zeros(1, 4, 2) + assert plan_cost(model, z0, standing_still, z_goal).item() == pytest.approx(25.0 * 5) + + +def test_last_objective_scores_only_the_final_state(): + model = PointMass() + z0 = torch.zeros(1, 2) + z_goal = torch.tensor([[3.0, 4.0]]) + actions = torch.tensor([[[1.0, 0.0], [1.0, 0.0]]]) # ends at (2, 0): energy 1 + 16 + assert plan_cost(model, z0, actions, z_goal, "last").item() == pytest.approx(17.0, rel=1e-5) + with pytest.raises(ValueError, match="objective"): + plan_cost(model, z0, actions, z_goal, "mean") + + +def test_cem_reaches_the_goal_with_either_objective(): + model = PointMass() + z0, z_goal = torch.zeros(1, 2), torch.tensor([GOAL]) + for objective in ("sum", "last"): + config = dataclasses.replace(SMALL, objective=objective) + plan = cem_plan(model, z0, z_goal, config, generator=_generator()) + final = rollout(model, z0, plan.actions.unsqueeze(0))[0, -1] + assert torch.dist(final, z_goal[0]).item() < 0.5, objective + + +def test_cem_reaches_the_goal_when_replanning_every_step(): + model = PointMass() + z = model.encode(torch.zeros(1, 2)) + z_goal = model.encode(torch.tensor([GOAL])) + generator = _generator() + for step in range(STEPS): + plan = cem_plan(model, z, z_goal, SMALL, steps_left=STEPS - step, generator=generator) + assert plan.actions.shape == (STEPS - step, 2) + z = model.predict(z, plan.actions[:1]) + assert torch.dist(z, z_goal).item() < 0.5 + + +def test_planned_sequence_beats_random_sequences(): + model = PointMass() + z0, z_goal = torch.zeros(1, 2), torch.tensor([GOAL]) + plan = cem_plan(model, z0, z_goal, SMALL, generator=_generator()) + random_actions = 1.5 * torch.randn(64, STEPS, 2, generator=_generator(1)) + random_costs = plan_cost(model, z0.expand(64, 2), random_actions, z_goal.expand(64, 2)) + planned_cost = plan_cost(model, z0, plan.actions.unsqueeze(0), z_goal).item() + assert planned_cost < random_costs.min().item() + assert plan.best_costs.shape == (SMALL.n_iters,) + assert plan.best_costs[-1] < plan.best_costs[0] + + +def test_same_seed_gives_the_same_plan(): + model = PointMass() + z0, z_goal = torch.zeros(1, 2), torch.tensor([GOAL]) + first = cem_plan(model, z0, z_goal, SMALL, generator=_generator(7)) + second = cem_plan(model, z0, z_goal, SMALL, generator=_generator(7)) + assert torch.equal(first.actions, second.actions) + + +def test_horizon_shrinks_to_the_steps_left_and_actions_respect_the_norm_limit(): + plan = cem_plan(PointMass(), torch.zeros(1, 2), torch.tensor([GOAL]), SMALL, steps_left=3, generator=_generator()) + assert plan.actions.shape == (3, 2) + assert (plan.actions.norm(dim=-1) <= SPEED).all() + + +def test_clip_action_norm_matches_the_eb_jepa_formula(): + actions = torch.tensor([[3.0, 4.0], [0.3, 0.4], [0.0, 0.0]]) + clipped = clip_action_norm(actions, 2.45) + assert clipped[0].norm().item() == pytest.approx(2.45, rel=1e-5) + assert torch.allclose(clipped[1], actions[1] * (0.5 / (0.5 + 1e-6))) + assert torch.equal(clipped[2], actions[2]) + + +@pytest.mark.parametrize( + "settings", + [ + {"num_elites": 1}, + {"num_elites": 300}, + {"plan_length": 0}, + {"var_scale": 0.0}, + {"max_norm": -1.0}, + {"objective": "mean"}, + ], +) +def test_bad_settings_are_refused(settings): + with pytest.raises(ValueError): + CEMConfig(**settings) + + +def test_cem_refuses_a_batch_of_start_states(): + with pytest.raises(ValueError, match="batch size 1"): + cem_plan(PointMass(), torch.zeros(2, 2), torch.zeros(2, 2), SMALL) + + +class PackedContext: + """A planning state that packs a history: rows 0..CAPACITY-1 hold positions, and the last row + holds the number of valid positions. ``energy`` compares only the latest valid position, so + it must read a different row of each state when states have different lengths.""" + + CAPACITY = 6 + state_shape = (CAPACITY + 1, 2) + action_dim = 2 + + def encode(self, obs: torch.Tensor) -> torch.Tensor: + z = torch.zeros(obs.shape[0], *self.state_shape) + z[:, 0] = obs + z[:, -1, 0] = 1 + return z + + @staticmethod + def _counts(z: torch.Tensor) -> torch.Tensor: + return z[:, -1, 0].long() + + def _latest(self, z: torch.Tensor) -> torch.Tensor: + return z[torch.arange(z.shape[0]), self._counts(z) - 1] + + def predict(self, z: torch.Tensor, a: torch.Tensor) -> torch.Tensor: + counts = self._counts(z) + if (counts >= self.CAPACITY).any(): + raise ValueError("context is full: the plan is longer than this model's capacity") + out = z.clone() # never write into the input, which may be an expanded view + rows = torch.arange(z.shape[0]) + out[rows, counts] = self._latest(z) + clip_action_norm(a, SPEED) + out[:, -1, 0] = counts + 1 + return out + + def energy(self, z: torch.Tensor, z_goal: torch.Tensor) -> torch.Tensor: + return ((self._latest(z) - self._latest(z_goal)) ** 2).sum(dim=-1) + + +def _cost_from_rollout(model, z0, actions, z_goal, objective): + states = rollout(model, z0, actions) + energies = torch.stack([model.energy(states[:, t], z_goal) for t in range(states.shape[1])], dim=1) + return energies.sum(dim=1) if objective == "sum" else energies[:, -1] + + +@pytest.mark.parametrize("model", [PointMass(), PackedContext()], ids=["point_mass", "packed_context"]) +@pytest.mark.parametrize("objective", ["sum", "last"]) +def test_plan_cost_matches_scoring_the_whole_rollout(model, objective): + z0 = model.encode(torch.tensor([[0.0, 0.0], [1.0, -2.0], [3.0, 3.0]])) + z_goal = model.encode(torch.tensor([[GOAL], [GOAL], [GOAL]]).squeeze(1)) + actions = torch.randn(3, 5, 2, generator=_generator(3)) + expected = _cost_from_rollout(model, z0, actions, z_goal, objective) + torch.testing.assert_close(plan_cost(model, z0, actions, z_goal, objective), expected) + zero = torch.zeros(3, 0, 2) + torch.testing.assert_close(plan_cost(model, z0, zero, z_goal, objective), model.energy(z0, z_goal)) + + +def test_packed_context_keeps_inputs_and_earlier_states_unchanged(): + model = PackedContext() + assert isinstance(model, WorldModel) + z0 = model.encode(torch.tensor([[0.0, 0.0]])) + z0_copy = z0.clone() + actions = torch.tensor([[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]]) + first = rollout(model, z0, actions) + first_copy = first.clone() + second = rollout(model, z0, actions) # branch again from the same start + rollout(model, first[:, -1], actions[:, :1]) # continue from a returned state + assert torch.equal(z0, z0_copy) + assert torch.equal(first, first_copy) + assert torch.equal(first, second) + + +def test_packed_context_accepts_an_expanded_start_and_scores_rows_independently(): + model = PackedContext() + z0 = model.encode(torch.tensor([[0.0, 0.0]])) + z_goal = model.encode(torch.tensor([GOAL])) + actions = torch.randn(4, 3, 2, generator=_generator(4)) + expanded = plan_cost(model, z0.expand(4, *model.state_shape), actions, z_goal.expand(4, *model.state_shape)) + rows = torch.cat([plan_cost(model, z0, actions[i : i + 1], z_goal) for i in range(4)]) + torch.testing.assert_close(expanded, rows) + + # States with different valid lengths, scored together and one at a time. + mixed = torch.cat([rollout(model, z0, actions[:1])[:, t] for t in range(4)]) + goals = z_goal.expand(4, *model.state_shape) + together = model.energy(mixed, goals) + alone = torch.cat([model.energy(mixed[i : i + 1], z_goal) for i in range(4)]) + torch.testing.assert_close(together, alone) + + +def test_packed_context_refuses_to_overflow_and_plans_within_capacity(): + model = PackedContext() + z0 = model.encode(torch.tensor([[0.0, 0.0]])) + full = torch.zeros(1, model.CAPACITY - 1, 2) + rollout(model, z0, full) # exactly at capacity + with pytest.raises(ValueError, match="capacity"): + rollout(model, z0, torch.zeros(1, model.CAPACITY, 2)) + config = dataclasses.replace(SMALL, plan_length=model.CAPACITY - 1, max_norm=SPEED) + plan = cem_plan(model, z0, model.encode(torch.tensor([[3.0, 3.0]])), config, generator=_generator()) + final = rollout(model, z0, plan.actions.unsqueeze(0))[:, -1] + assert torch.dist(model._latest(final), torch.tensor([[3.0, 3.0]])).item() < 0.5 + + +@pytest.mark.parametrize( + "settings", + [ + {"var_scale": float("nan")}, + {"var_scale": float("inf")}, + {"max_norm": float("nan")}, + {"plan_length": 1.5}, + {"num_elites": 2.5}, + {"n_iters": True}, + ], +) +def test_non_finite_or_non_integer_settings_are_refused(settings): + with pytest.raises(ValueError): + CEMConfig(**settings) + + +class _BadEnergy(PointMass): + def __init__(self, energy): + self._energy = energy + + def energy(self, z, z_goal): + return self._energy(super().energy(z, z_goal)) + + +@pytest.mark.parametrize( + ("energy", "message"), + [ + (lambda e: torch.full_like(e, float("nan")), "100 of 100 candidate costs are NaN or infinite"), + (lambda e: e[:1], "one cost per candidate"), + ], + ids=["all_nan", "wrong_shape"], +) +def test_invalid_costs_fail_loudly(energy, message): + with pytest.raises(ValueError, match=message): + cem_plan(_BadEnergy(energy), torch.zeros(1, 2), torch.tensor([GOAL]), SMALL, generator=_generator()) + + +def test_a_few_invalid_costs_fail_loudly(): + """Only the candidates whose final state ends far from the goal are infinite.""" + last = dataclasses.replace(SMALL, objective="last") + some_infinite = _BadEnergy(lambda e: torch.where(e > e.median(), torch.full_like(e, float("inf")), e)) + with pytest.raises(ValueError, match="candidate costs are NaN or infinite") as caught: + cem_plan(some_infinite, torch.zeros(1, 2), torch.tensor([GOAL]), last, generator=_generator()) + invalid = int(str(caught.value).split(" of ")[0]) + assert 0 < invalid < SMALL.num_samples + + +def test_mismatched_start_and_goal_are_refused(): + with pytest.raises(ValueError, match="state_shape"): + cem_plan(PointMass(), torch.zeros(1, 2), torch.zeros(1, 3), SMALL) + + +def test_actions_are_float32_whatever_the_default_or_state_dtype(): + previous = torch.get_default_dtype() + torch.set_default_dtype(torch.float64) # without an explicit dtype, the planner would follow this + try: + plan = cem_plan( + PointMass(), torch.zeros(1, 2, dtype=torch.float64), torch.tensor([GOAL], dtype=torch.float64), SMALL + ) + finally: + torch.set_default_dtype(previous) + assert plan.actions.dtype == ACTION_DTYPE == torch.float32