File size: 3,922 Bytes
9ede8c0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | """
The four learned pieces: an encoder (real state -> latent), a dynamics model
(latent + action -> predicted next latent + reward), a value head (latent ->
estimated cost-to-go), and a diagnostic decoder (latent -> reconstructed
state, used only to sanity-check the latent isn't collapsing -- never
consulted for a real decision at inference time).
Config-driven: nothing here is hardcoded to Connect-4 beyond `state_dim`/
`num_actions`, so this file would work for any domain with a fixed-length
tuple state and a fixed discrete action space.
"""
import torch
import torch.nn as nn
def mlp(dims, out_activation=None):
layers = []
for i in range(len(dims) - 1):
layers.append(nn.Linear(dims[i], dims[i + 1]))
is_last = i == len(dims) - 2
if not is_last:
layers.append(nn.ReLU())
elif out_activation is not None:
layers.append(out_activation)
return nn.Sequential(*layers)
class Encoder(nn.Module):
def __init__(self, state_dim, latent_dim, hidden_dim=128):
super().__init__()
self.net = mlp([state_dim, hidden_dim, hidden_dim, latent_dim])
def forward(self, state):
return self.net(state)
class DynamicsModel(nn.Module):
def __init__(self, latent_dim, num_actions, hidden_dim=128):
super().__init__()
self.num_actions = num_actions
self.trunk = mlp([latent_dim + num_actions, hidden_dim, hidden_dim])
self.next_latent_head = nn.Linear(hidden_dim, latent_dim)
self.reward_head = nn.Linear(hidden_dim, 1)
def forward(self, z, action_idx):
action_onehot = nn.functional.one_hot(action_idx, self.num_actions).float()
h = self.trunk(torch.cat([z, action_onehot], dim=-1))
next_z = self.next_latent_head(h)
reward = self.reward_head(h).squeeze(-1)
return next_z, reward
class ValueHead(nn.Module):
def __init__(self, latent_dim, hidden_dim=128):
super().__init__()
self.net = mlp([latent_dim, hidden_dim, hidden_dim, 1])
def forward(self, z):
return self.net(z).squeeze(-1)
class Decoder(nn.Module):
def __init__(self, latent_dim, state_dim, hidden_dim=128):
super().__init__()
self.net = mlp([latent_dim, hidden_dim, hidden_dim, state_dim])
def forward(self, z):
return self.net(z)
class WorldModel(nn.Module):
def __init__(self, state_dim, num_actions, latent_dim=64, hidden_dim=128):
super().__init__()
self.state_dim = state_dim
self.num_actions = num_actions
self.latent_dim = latent_dim
self.encoder = Encoder(state_dim, latent_dim, hidden_dim)
self.dynamics = DynamicsModel(latent_dim, num_actions, hidden_dim)
self.value = ValueHead(latent_dim, hidden_dim)
self.decoder = Decoder(latent_dim, state_dim, hidden_dim)
# Value-target normalization stats: NOT learned, set once by
# whichever value-training pass runs (see verifier.py's
# train_mc_value_onpolicy) from the actual label distribution it
# sees. Buffers (not plain attributes) so they save/load with the
# checkpoint automatically.
self.register_buffer("value_target_mean", torch.tensor(0.0))
self.register_buffer("value_target_std", torch.tensor(1.0))
def encode(self, state):
return self.encoder(state)
def imagine_step(self, z, action_idx):
return self.dynamics(z, action_idx)
def evaluate(self, z):
"""Always returns real-scale value estimates (remaining cost, same
units `imagine_step`'s predicted reward uses) -- the head internally
predicts a normalized target, denormalized here so no caller needs
to know normalization is happening."""
raw = self.value(z)
return raw * self.value_target_std + self.value_target_mean
def reconstruct(self, z):
return self.decoder(z)
|