| """RSSM (Recurrent State-Space Model) for Grid-JEPA.""" |
| from typing import Dict, Tuple |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.distributions import OneHotCategorical |
|
|
| class RSSM(nn.Module): |
| def __init__(self, embed_dim=384, latent_dim=32, latent_classes=32, hidden_dim=256, |
| action_dim=64, num_actions=6, obs_dim=384, reward_bins=41): |
| super().__init__() |
| self.latent_dim = latent_dim |
| self.latent_classes = latent_classes |
| self.hidden_dim = hidden_dim |
| self.action_embed = nn.Embedding(num_actions, action_dim) |
| self.seq_model = nn.GRUCell(latent_dim * latent_classes + action_dim, hidden_dim) |
| self.encoder = nn.Sequential( |
| nn.Linear(obs_dim + hidden_dim, hidden_dim), nn.ELU(), |
| nn.Linear(hidden_dim, latent_dim * latent_classes), |
| ) |
| self.dynamics = nn.Sequential( |
| nn.Linear(hidden_dim, hidden_dim), nn.ELU(), |
| nn.Linear(hidden_dim, latent_dim * latent_classes), |
| ) |
| self.reward_predictor = nn.Sequential( |
| nn.Linear(hidden_dim + latent_dim * latent_classes, hidden_dim), nn.ELU(), |
| nn.Linear(hidden_dim, reward_bins), |
| ) |
| self.continue_predictor = nn.Sequential( |
| nn.Linear(hidden_dim + latent_dim * latent_classes, hidden_dim), nn.ELU(), |
| nn.Linear(hidden_dim, 1), |
| ) |
|
|
| def init_state(self, B, device): |
| return torch.zeros(B, self.hidden_dim, device=device), torch.zeros(B, self.latent_dim * self.latent_classes, device=device) |
|
|
| def _sample_st(self, logits): |
| dist = OneHotCategorical(logits=logits) |
| sample = dist.sample() |
| return sample + logits - logits.detach() |
|
|
| def observe(self, x_t, a_tm1, h_tm1, z_tm1): |
| B = x_t.shape[0] |
| a_emb = self.action_embed(a_tm1) |
| gru_in = torch.cat([z_tm1, a_emb], dim=-1) |
| h_t = self.seq_model(gru_in, h_tm1) |
| z_pred_logits = self.dynamics(h_t).reshape(B, self.latent_dim, self.latent_classes) |
| enc_in = torch.cat([x_t, h_t], dim=-1) |
| z_logits = self.encoder(enc_in).reshape(B, self.latent_dim, self.latent_classes) |
| z_onehot = self._sample_st(z_logits) |
| z_t = z_onehot.reshape(B, -1) |
| return h_t, z_t, z_logits, z_pred_logits |
|
|
| def imagine(self, h_t, z_t, a_t): |
| B = h_t.shape[0] |
| a_emb = self.action_embed(a_t) |
| gru_in = torch.cat([z_t, a_emb], dim=-1) |
| h_tp1 = self.seq_model(gru_in, h_t) |
| z_pred_logits = self.dynamics(h_tp1).reshape(B, self.latent_dim, self.latent_classes) |
| z_onehot = self._sample_st(z_pred_logits) |
| z_tp1 = z_onehot.reshape(B, -1) |
| return h_tp1, z_tp1, z_pred_logits |
|
|
| def predict_reward(self, h, z): |
| x = torch.cat([h, z], dim=-1) |
| return self.reward_predictor(x) |
|
|
| def predict_continue(self, h, z): |
| x = torch.cat([h, z], dim=-1) |
| return self.continue_predictor(x) |
|
|
| def rollout(self, h_0, z_0, actions): |
| B, H = actions.shape |
| h_states, z_states, rewards, continues = [], [], [], [] |
| h_t, z_t = h_0, z_0 |
| for t in range(H): |
| h_t, z_t, _ = self.imagine(h_t, z_t, actions[:, t]) |
| h_states.append(h_t) |
| z_states.append(z_t) |
| rewards.append(self.predict_reward(h_t, z_t)) |
| continues.append(self.predict_continue(h_t, z_t)) |
| return { |
| "h_states": torch.stack(h_states, 1), |
| "z_states": torch.stack(z_states, 1), |
| "rewards": torch.stack(rewards, 1), |
| "continues": torch.stack(continues, 1), |
| } |
|
|
| if __name__ == "__main__": |
| B, obs = 2, 128 |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| rssm = RSSM(obs_dim=obs).to(device) |
| x = torch.randn(B, obs, device=device) |
| a = torch.randint(0, 6, (B,), device=device) |
| h, z = rssm.init_state(B, device) |
| h, z, z_logits, z_pred = rssm.observe(x, a, h, z) |
| print(f"h: {h.shape}, z: {z.shape}, z_logits: {z_logits.shape}") |
| h2, z2, zpl = rssm.imagine(h, z, a) |
| print(f"h2: {h2.shape}, z2: {z2.shape}") |
| print(f"r: {rssm.predict_reward(h, z).shape}, c: {rssm.predict_continue(h, z).shape}") |
| actions = torch.randint(0, 6, (B, 10), device=device) |
| out = rssm.rollout(h, z, actions) |
| print(f"Rollout: h={out['h_states'].shape}, z={out['z_states'].shape}, r={out['rewards'].shape}, c={out['continues'].shape}") |
|
|
