World models
Let the agent dream: learn the world as a model and train inside the dream.
A world model has three parts: a VAE that compresses screen frames into latents, an MDN-RNN that predicts the next latent (memory), and a small controller that acts. Once the model is good enough, the agent can train its policy inside the model's "dream" and transfer it to the real world.
The example learns CartPole's dynamics from real experience, then "dreams" 10 steps and compares the pole angle with reality.
| Part | Role |
|---|---|
| V — VAE | Vision: compress each frame to z |
| M — MDN-RNN | Memory: predicts p(zₜ₊₁ | zₜ, aₜ, hₜ) |
| C — Controller | Acts from z and h, trained with CMA-ES |
Key takeaways
- A generative model can serve as a simulated world for RL
- Links to the Deep RL series: this is deep model-based RL
import numpy as np, torch, torch.nn as nn
import gymnasium as gym
torch.manual_seed(0)
env = gym.make("CartPole-v1")
# 1) Collect real experience with a random policy
S, A, S2 = [], [], []
for ep in range(200):
s, _ = env.reset(seed=ep); done = False
while not done:
a = env.action_space.sample()
s2, r, term, trunc, _ = env.step(a)
S.append(s); A.append(a); S2.append(s2); s, done = s2, term or trunc
S, A, S2 = map(lambda v: torch.tensor(np.array(v), dtype=torch.float32), (S, A, S2))
# 2) Learn the dynamics: (state, action) -> next state (the "memory" model)
model = nn.Sequential(nn.Linear(5, 128), nn.ReLU(), nn.Linear(128, 4))
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
inp = torch.cat([S, A[:, None]], 1)
for step in range(2000):
loss = ((S + model(inp) - S2) ** 2).mean() # predict the change in state
opt.zero_grad(); loss.backward(); opt.step()
print(f"dynamics model loss {loss.item():.6f} ({len(S)} real transitions)")
# 3) "Dream": roll out 10 steps inside the model and compare with reality
s, _ = env.reset(seed=999); dream = torch.tensor(s, dtype=torch.float32)
for t in range(10):
a = t % 2
s, *_ = env.step(a)
with torch.no_grad():
dream = dream + model(torch.cat([dream, torch.tensor([float(a)])]))
print("real pole angle:", round(float(s[2]), 4))
print("dream pole angle:", round(float(dream[2]), 4))Output
dynamics model loss 0.000003 (4204 real transitions) real pole angle: 0.0641 dream pole angle: 0.0662
- #GenerativeAI
- #DeepLearning
- #PythonສຳລັບຄົນLaos
- #AI
- #ບອກຕໍ່ຄວາມຮູ້