Oneiros

Status: complete research artifact — results verified. v0.2 adds image-based perception as a demonstrated result (the four pixel gates below), closing v0.1's documented next step; the V-JEPA 2 scale-up remains future work (needs GPU + pretrained weights).

A JEPA-style latent world model that an agent calls as a tool — over MCP — to plan.

Oneiros is a small, CPU-only research artifact at the intersection of agentic systems and world models. It trains a Joint Embedding Predictive Architecture (JEPA) world model on a 2D point-mass environment, then exposes that model as a set of Model Context Protocol tools. An agent plans by calling the learned predictive world model as a tool — encoding observations into latents, rolling latent dynamics forward, and running model-predictive control entirely in latent space — rather than the usual pattern of an LLM calling hand-written functions.

Why predict in latent space

A JEPA predicts the future in a learned latent space, not in pixels or tokens. Given an observation o_t and an action a_t, it learns an encoder f and a predictor g such that g(f(o_t), a_t) matches f(o_{t+1}) — there is no decoder and no pixel-reconstruction loss. This matters because reconstruction wastes capacity modeling perceptually salient but control-irrelevant detail (texture, lighting, background), while a latent predictor is free to discard everything that does not help it anticipate the future. The cost is a well-known failure mode: latent prediction can collapse to a constant (every observation maps to the same point, making prediction trivially perfect). Oneiros defeats collapse with an EMA target encoder, stop-gradients, and a VICReg-style variance + covariance penalty, and then uses the resulting latent dynamics for planning.

Related MCP server: MuJoCo MCP Server

Architecture

flowchart LR
    subgraph Env["Point-mass environment (numpy)"]
        O["obs o_t"]
        ON["obs o_t+1"]
    end

    subgraph WM["JEPA world model (torch, CPU)"]
        F["encoder f"]
        FT["EMA target f_target<br/>(stop-grad)"]
        G["predictor g"]
        O --> F --> Z["latent z_t"]
        Z --> G
        A["action a_t"] --> G
        G --> ZH["z_hat_t+1"]
        ON --> FT --> ZT["z_t+1 (target)"]
        ZH -. "MSE + VICReg<br/>variance/covariance" .-> ZT
    end

    subgraph MCP["MCP server (agentic interface)"]
        T1["encode_observation"]
        T2["predict_rollout"]
        T3["plan_to_goal"]
        T4["reset_env / step_env"]
    end

    subgraph Agent["Agent loop"]
        P["latent MPC planner<br/>(CEM over g)"]
    end

    WM --> MCP
    MCP <--> Agent
    P -->|"first action"| Env

The agent never sees the environment's dynamics. It calls plan_to_goal, which encodes the current and goal observations, searches action sequences by rolling the predictor g forward H steps in latent space (cross-entropy method), scores each candidate by predicted-latent distance to the goal, and returns the first action. The planner replans every step (receding-horizon MPC).

Verified results

Numbers below are from an actual run on this machine (CPU only, seed 0). Train with python -m oneiros.train and reproduce the diagnostics with python -m oneiros.demo_agent.

Training takes about 21 seconds for 4000 steps. The checkpoint (oneiros/checkpoint.pt, ~240 KB) is committed so the demo, MCP server, and planning tests run without retraining.

The pixel gates (v0.2): image-based perception, demonstrated

v0.1 documented why the image encoder could not beat the identity baseline. The diagnosis had two parts, and each got a principled fix rather than a knob-twiddle:

1. Consecutive frames were nearly identical (the blob moves ~a pixel per step), so "predict no change" was already an excellent predictor. Fix: the swift environment preset (PointMassConfig.swift()) — larger dt and acceleration so the agent moves several pixels per frame, plus speed-proportional drag so the dynamics are genuinely non-linear.

2. A single frame hides velocity — the dynamics are second-order, so no single-frame predictor can recover the next state, and the convolutional encoder exploited this by temporal smoothing (mapping consecutive frames to nearly identical latents; the dataset-wide variance penalty does not forbid it — more training made it worse, 0.89 → 0.94 MSE ratio). Fixes: two-frame stacking (velocity becomes observable from pixels, the same reason pixel world models from DQN to V-JEPA consume clips, not stills) and a delta-variance penalty (VICReg-style hinge on the std of z_{t+1} - z_t) that forbids the temporal collapse outright.

Results from the committed checkpoint_image.pt (~735 KB), heldout data, enforced as tests in tests/test_gates_image.py:

One honest negative, measured and deliberately not gated: a privileged linear-dynamics MPC reading the true 4D state still reaches the goal about twice as fast as the pixel planner (median ~5 vs ~9.5 steps). Planning from pixels has not caught planning from privileged state, and gate (g) shows open-loop imagination degrading by horizon 12 — which is exactly why the planner replans every step. The baselines live in oneiros/baselines.py; reproduce with the image-training command under Run.

The agent drives the point-mass to the goal (green star) using only the world-model planning interface.

What this is — and isn't

This is a genuine, end-to-end demonstration that (1) a non-trivial latent dynamics model can be learned without collapse and without reconstruction, and (2) planning in that latent space solves a control task far better than chance, all behind an agentic tool interface.

It is not at scale. The environments are toy 2D point-masses, the models are tiny (a few hundred KB). The default committed model uses a vector-state observation on the original environment; the committed image model (v0.2, checkpoint_image.pt) perceives stacked 32x32 frames on the swift environment and clears its own four gates above — v0.1's "image-based perception is the documented next step" is now a demonstrated result, with the diagnosis (frame similarity + hidden velocity + temporal smoothing) and fixes documented rather than hand-waved. What remains future work is the scale-up: a frozen pretrained perception encoder such as V-JEPA 2 with a learned latent dynamics head on top, exactly the recipe this toy mirrors — that step needs a GPU and pretrained weights, which this CPU-only artifact deliberately does not assume.

Install

Requires Python 3.12. A project-local virtual environment is recommended.

python -m venv .venv
# Windows: .venv\Scripts\activate    |    Unix: source .venv/bin/activate
pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install -e ".[dev]"

torch is installed from the CPU wheel index — no GPU is needed or used.

Run

# Train the JEPA world model (writes oneiros/checkpoint.pt). ~21s on CPU.
python -m oneiros.train --obs-mode vector --steps 4000

# Train the image world model on the swift environment (~3 min on CPU).
# This is the exact recipe of the committed checkpoint: two-frame stacking,
# the delta-variance penalty, and an explicit --checkpoint so the vector
# model's default path is not clobbered.
python -m oneiros.train --obs-mode image --env swift --steps 6000 \
    --frame-stack 2 --delta-var-weight 25 --checkpoint oneiros/checkpoint_image.pt

# Run the scripted agent: drive to the goal via latent MPC, write the GIF +
# diagnostic plots to assets/.
python -m oneiros.demo_agent --seed 0 --k 20

# Tests, including the three honesty gates (uses the committed checkpoint).
pytest -q

# Lint.
ruff check oneiros tests

MCP server

The world model is exposed over MCP as a stdio server:

# vector model (default)
python -m oneiros.mcp_server

# the committed image model: stacked-frame observations, swift environment
python -m oneiros.mcp_server --obs-mode image

# any other checkpoint
python -m oneiros.mcp_server --checkpoint /path/to/checkpoint.pt

--obs-mode selects which committed world model the server exposes. The environment tools serve observations in the loaded model's format — 6D vectors, or the stacked rendered frames the image model trains on — and reset_env returns a ready-made goal_observation for the planning tools. The served environment always uses the configuration the checkpoint was trained on (the swift preset for the image model).

Tools:

To register the server with Claude Desktop, add this to claude_desktop_config.json (use absolute paths for your checkout):

{
  "mcpServers": {
    "oneiros": {
      "command": "C:/path/to/Oneiros/.venv/Scripts/python.exe",
      "args": ["-m", "oneiros.mcp_server"],
      "cwd": "C:/path/to/Oneiros"
    }
  }
}

On Unix the command is .venv/bin/python. Append "--obs-mode", "image" to args to serve the committed image world model instead of the vector one.

Repository layout

oneiros/
  env.py          # deterministic 2D point-mass environment (numpy)
  model.py        # JEPA encoder + predictor + VICReg regularizers
  data.py         # random-policy rollout replay buffer
  train.py        # JEPA training loop, evaluation, checkpoint I/O
  planner.py      # latent-space MPC (CEM / random shooting)
  baselines.py    # pixel episode runners + privileged linear-MPC baseline
  mcp_server.py   # MCP tools exposing the world model
  demo_agent.py   # scripted agent + diagnostics (GIF, plots)
  checkpoint.pt        # committed vector model (~240 KB)
  checkpoint_image.pt  # committed image model, swift env (~735 KB)
tests/            # determinism, shapes, and the three honesty gates
assets/           # generated GIF and diagnostic figures

See SYNERGY.md for how the same latent-dynamics idea connects to regime-aware modeling in time series.

License

MIT — see LICENSE.