DreamerVLA 🤖

Language-conditioned model-based RL agent for robot manipulation in simulation.

Combines Dreamer-style RSSM world models with CLIP-encoded language goals and an actor-critic policy trained entirely in latent imagination — no real environment steps needed during policy learning.

---

Demo

Real Camera	World Model Imagination

Agent trained on FetchReach-v3: "reach the red block"

---

Architecture

Image obs (RGB)  →  CNN Encoder  ─┐
                                   ├→  Feature Fusion  →  RSSM World Model  →  Imagined rollouts
Language goal    →  CLIP Encoder  ─┘         (h_t, z_t)          ↓
                                                           Actor-Critic
                                                                ↓
                                                        Actions → Environment

Key components

• CNN Encoder — encodes (3, 64, 64) image observations into a 256-dim embedding
• CLIP Encoder (frozen ViT-B/32) — encodes text goals into 256-dim embeddings
• Feature Fusion — concatenates visual + language embeddings
• RSSM — Recurrent State Space Model with deterministic h_t and stochastic z_t; learns to predict future latent states
• Decoder + Reward Head — reconstructs observations and predicts rewards from latent state (world model training signal)
• Actor-Critic — trained on K=15 step imagined rollouts inside the world model; never touches the real environment during policy updates

---

Results

Metric	Value
Avg episode return	~TBD
World model recon loss	~TBD
Training time	~4 hrs on RTX 3090
Environment	FetchReach-v3 (MuJoCo)

Training curves on Weights & Biases

---

Quickstart

Install

git clone https://github.com/yourusername/dreamervla
cd dreamervla
pip install torch torchvision gymnasium[mujoco] wandb imageio
pip install git+https://github.com/openai/CLIP.git

Train

python train.py

Watch a trained agent

python enjoy.py --checkpoint checkpoints/best.pt --goal "reach the red block"

---

Project structure

dreamervla/
├── model.py          # CNN encoder, CLIP encoder, RSSM, actor-critic
├── train.py          # main training loop
├── replay_buffer.py  # circular replay buffer
├── env_wrapper.py    # gym env → 64x64 RGB tensors
├── utils.py          # logging, GIF recorder, checkpointing
├── config.py         # all hyperparameters
└── enjoy.py          # visualize a trained agent

---

Hyperparameters

Parameter	Value
Latent dim	32
Hidden dim (RSSM)	512
Embed dim	256
Imagination horizon	15
Batch size	32
Sequence length	50
Learning rate	3e-4
Discount γ	0.99
KL weight	1.0

---

How it works

World model training

The RSSM is trained on real environment transitions stored in a replay buffer. Given a sequence of observations and actions, it learns to:

1. Encode observations into a compact latent state (h_t, z_t)
2. Predict the next latent state given an action (prior)
3. Reconstruct the observation from the latent state (decoder)
4. Predict the reward from the latent state (reward head)

Policy training (in imagination)

Once the world model is trained, the actor-critic never touches the real environment again during a policy update. Instead it:

1. Samples a starting latent state from the replay buffer
2. Rolls out K steps using the actor + RSSM prior (pure imagination)
3. Computes λ-returns from imagined rewards + critic values
4. Updates actor to maximize returns, critic to predict them accurately

Language conditioning

Goals are encoded once via frozen CLIP and fused with visual features at every timestep. The same agent can follow different instructions at inference time by swapping the goal string.

---

Inspiration

• Mastering Diverse Domains through World Models (DreamerV3) — Hafner et al., 2023
• Learning Latent Dynamics for Planning from Pixels (DreamerV1) — Hafner et al., 2019
• RT-2: Vision-Language-Action Models — Brohan et al., 2023
• CLIP: Learning Transferable Visual Models — Radford et al., 2021

---

License

MIT