World Models

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))