| from __future__ import annotations |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
|
|
| class LongMemoryState(nn.Module): |
| """Chunk-level long-memory reader/writer driven by clean observation latents.""" |
|
|
| def __init__( |
| self, |
| *, |
| latent_channels: int, |
| memory_dim: int, |
| num_tokens: int, |
| num_heads: int, |
| video_hidden_size: int, |
| action_hidden_size: int, |
| memory_noise_std: float = 0.0, |
| memory_dropout_prob: float = 0.0, |
| ) -> None: |
| super().__init__() |
| self.memory_dim = int(memory_dim) |
| self.num_tokens = int(num_tokens) |
| self.memory_noise_std = float(memory_noise_std) |
| self.memory_dropout_prob = float(memory_dropout_prob) |
| self.init_memory = nn.Parameter(torch.zeros(self.num_tokens, self.memory_dim)) |
| nn.init.trunc_normal_(self.init_memory, std=0.02) |
|
|
| self.obs_proj = nn.Linear(int(latent_channels), self.memory_dim) |
| self.obs_norm = nn.LayerNorm(self.memory_dim) |
| self.mem_norm = nn.LayerNorm(self.memory_dim) |
| self.write_attn = nn.MultiheadAttention( |
| embed_dim=self.memory_dim, |
| num_heads=int(num_heads), |
| batch_first=True, |
| ) |
| self.write_ffn = nn.Sequential( |
| nn.LayerNorm(self.memory_dim), |
| nn.Linear(self.memory_dim, self.memory_dim * 4), |
| nn.GELU(), |
| nn.Linear(self.memory_dim * 4, self.memory_dim), |
| ) |
| self.write_gate = nn.Linear(self.memory_dim * 2, self.memory_dim) |
|
|
| self.read_norm = nn.LayerNorm(self.memory_dim) |
| self.to_video_bias = nn.Linear(self.memory_dim, int(video_hidden_size), bias=False) |
| self.to_action_bias = nn.Linear(self.memory_dim, int(action_hidden_size), bias=False) |
| nn.init.zeros_(self.to_video_bias.weight) |
| nn.init.zeros_(self.to_action_bias.weight) |
|
|
| def init_state(self, batch_size: int, *, device: torch.device, dtype: torch.dtype) -> torch.Tensor: |
| return self.init_memory.unsqueeze(0).expand(int(batch_size), -1, -1).to(device=device, dtype=dtype) |
|
|
| def encode_observation_tokens(self, latents: torch.Tensor) -> torch.Tensor: |
| if latents.ndim != 5: |
| raise ValueError(f"latents must be [B, C, T, H, W], got {tuple(latents.shape)}") |
| tokens = latents.permute(0, 2, 3, 4, 1).reshape(int(latents.shape[0]), -1, int(latents.shape[1])) |
| tokens = self.obs_proj(tokens) |
| return self.obs_norm(tokens) |
|
|
| def update(self, memory: torch.Tensor, observation_latents: torch.Tensor) -> torch.Tensor: |
| obs_tokens = self.encode_observation_tokens(observation_latents) |
| norm_memory = self.mem_norm(memory) |
| attn_update, _ = self.write_attn(norm_memory, obs_tokens, obs_tokens, need_weights=False) |
| proposal = self.write_ffn(attn_update) |
| gate = torch.sigmoid(self.write_gate(torch.cat([memory, proposal], dim=-1))) |
| updated = memory + gate * proposal |
| if self.training and self.memory_noise_std > 0.0: |
| updated = updated + torch.randn_like(updated) * float(self.memory_noise_std) |
| if self.training and self.memory_dropout_prob > 0.0: |
| updated = F.dropout(updated, p=float(self.memory_dropout_prob), training=True) |
| return updated |
|
|
| def summarize(self, memory: torch.Tensor) -> dict[str, torch.Tensor]: |
| summary = self.read_norm(memory).mean(dim=1) |
| return { |
| "summary": summary, |
| "video_bias": self.to_video_bias(summary), |
| "action_bias": self.to_action_bias(summary), |
| } |
|
|
