Spaces:
Running on Zero
Running on Zero
| """ | |
| State representation module inspired by CUT3R. | |
| Maintains persistent state across time steps for multi-body dynamics. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from typing import Optional, Tuple | |
| import math | |
| class StateRepresentation(nn.Module): | |
| """ | |
| Persistent state representation for multi-body dynamics, following | |
| CUT3R's state management: | |
| - Register tokens stored in an nn.Embedding, as in CUT3R | |
| - Tokens laid out on a 3D grid whose coordinates drive RoPE | |
| - Rotary-only position encoding (no additive position embeddings) | |
| """ | |
| def __init__( | |
| self, | |
| num_tokens: int = 216, # 6x6x6 = 216 tokens on a 3D grid | |
| token_dim: int = 768, | |
| position_encoding: str = 'learnable_3d', # rope-only: use 3D grid positions | |
| rope_dim: Optional[int] = None | |
| ): | |
| super().__init__() | |
| self.num_tokens = num_tokens | |
| self.token_dim = token_dim | |
| self.position_encoding = position_encoding | |
| # Register tokens (no additive position embeddings; rope-only) | |
| self.register_tokens = nn.Embedding(num_tokens, token_dim) | |
| # Small random initialization, following CUT3R | |
| nn.init.normal_(self.register_tokens.weight, mean=0., std=0.02) | |
| # 3D grid positions for RoPE (rope-only) | |
| if position_encoding != 'learnable_3d': | |
| raise ValueError("StateRepresentation now supports only 'learnable_3d' (rope-only)") | |
| if num_tokens == 216: | |
| self.grid_3d = (6, 6, 6) | |
| else: | |
| self.grid_3d = self._compute_3d_grid(num_tokens) | |
| # Layer normalization | |
| self.norm = nn.LayerNorm(token_dim) | |
| def _compute_3d_grid(self, n): | |
| """Compute 3D grid dimensions for n tokens.""" | |
| # Try to make roughly cubic | |
| cube_root = int(round(n ** (1/3))) | |
| for x in range(cube_root, 0, -1): | |
| if n % x == 0: | |
| remaining = n // x | |
| sqrt_rem = int(round(remaining ** 0.5)) | |
| for y in range(sqrt_rem, 0, -1): | |
| if remaining % y == 0: | |
| z = remaining // y | |
| return (x, y, z) | |
| return (n, 1, 1) | |
| def get_state_and_positions(self, batch_size: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """ | |
| Get initial state tokens and 3D grid positions for RoPE (rope-only). | |
| Returns: | |
| state_tokens: (B, num_tokens, token_dim) | |
| positions: (B, num_tokens, 3) in [-1, 1] | |
| """ | |
| # Get state tokens using embedding | |
| indices = torch.arange(self.num_tokens, dtype=torch.long, device=device) | |
| state_tokens = self.register_tokens(indices) # (num_tokens, token_dim) | |
| state_tokens = state_tokens.unsqueeze(0).expand(batch_size, -1, -1) | |
| # Normalize tokens | |
| state_tokens = self.norm(state_tokens) | |
| # Build 3D grid positions normalized to [-1, 1] | |
| gx, gy, gz = self.grid_3d | |
| xs = torch.linspace(-1, 1, steps=gx, device=device, dtype=state_tokens.dtype) | |
| ys = torch.linspace(-1, 1, steps=gy, device=device, dtype=state_tokens.dtype) | |
| zs = torch.linspace(-1, 1, steps=gz, device=device, dtype=state_tokens.dtype) | |
| grid = torch.stack(torch.meshgrid(xs, ys, zs, indexing='ij'), dim=-1).view(-1, 3) # (num_tokens,3) | |
| if grid.shape[0] != self.num_tokens: | |
| # Tile or truncate to match the token count | |
| if grid.shape[0] < self.num_tokens: | |
| reps = (self.num_tokens + grid.shape[0] - 1) // grid.shape[0] | |
| grid = grid.repeat(reps, 1)[:self.num_tokens] | |
| else: | |
| grid = grid[:self.num_tokens] | |
| positions = grid.unsqueeze(0).expand(batch_size, -1, -1) # (B,N,3) | |
| return state_tokens, positions | |
| def get_initial_state(self, batch_size: int) -> torch.Tensor: | |
| """ | |
| Return the initial state tokens only (grid positions discarded). | |
| """ | |
| device = self.register_tokens.weight.device | |
| state, _ = self.get_state_and_positions(batch_size, device) | |
| return state | |
| def forward(self, state: Optional[torch.Tensor] = None, | |
| batch_size: Optional[int] = None) -> torch.Tensor: | |
| """ | |
| Get state representation (rope-only). | |
| Args: | |
| state: Current state if available (B, num_tokens, token_dim) | |
| batch_size: Batch size if state is None | |
| Returns: | |
| State tokens (normalized) | |
| """ | |
| if state is None: | |
| if batch_size is None: | |
| raise ValueError("Either state or batch_size must be provided") | |
| return self.get_initial_state(batch_size) | |
| # Normalize existing state | |
| state = self.norm(state) | |
| return state | |
| class StateUpdateModule(nn.Module): | |
| """ | |
| Module for updating state based on new observations. | |
| Implements the State-Update mechanism from CUT3R. | |
| """ | |
| def __init__( | |
| self, | |
| state_dim: int = 512, | |
| observation_dim: int = 512, | |
| num_heads: int = 8, | |
| dropout: float = 0.1 | |
| ): | |
| super().__init__() | |
| self.state_dim = state_dim | |
| self.observation_dim = observation_dim | |
| # Project observations to state dimension if needed | |
| if observation_dim != state_dim: | |
| self.obs_projection = nn.Linear(observation_dim, state_dim) | |
| else: | |
| self.obs_projection = nn.Identity() | |
| # Multi-head attention for state update | |
| self.cross_attention = nn.MultiheadAttention( | |
| embed_dim=state_dim, | |
| num_heads=num_heads, | |
| dropout=dropout, | |
| batch_first=True | |
| ) | |
| # Gating mechanism for state update | |
| self.gate = nn.Sequential( | |
| nn.Linear(state_dim * 2, state_dim), | |
| nn.Sigmoid() | |
| ) | |
| # Feedforward network | |
| self.ffn = nn.Sequential( | |
| nn.Linear(state_dim, state_dim * 4), | |
| nn.GELU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(state_dim * 4, state_dim) | |
| ) | |
| # Layer normalization | |
| self.norm1 = nn.LayerNorm(state_dim) | |
| self.norm2 = nn.LayerNorm(state_dim) | |
| def forward( | |
| self, | |
| state: torch.Tensor, | |
| observation: torch.Tensor | |
| ) -> torch.Tensor: | |
| """ | |
| Update state based on new observation. | |
| Args: | |
| state: Current state (B, num_tokens, state_dim) | |
| observation: New observation features (B, obs_len, obs_dim) | |
| Returns: | |
| Updated state (B, num_tokens, state_dim) | |
| """ | |
| # Project observation to state dimension | |
| obs_proj = self.obs_projection(observation) | |
| # Cross-attention: state attends to observation | |
| attn_output, _ = self.cross_attention( | |
| query=state, | |
| key=obs_proj, | |
| value=obs_proj | |
| ) | |
| # Gated update mechanism | |
| gate_input = torch.cat([state, attn_output], dim=-1) | |
| gate = self.gate(gate_input) | |
| # Apply gated update | |
| state = state + gate * self.norm1(attn_output) | |
| # Feedforward network with residual connection | |
| state = state + self.ffn(self.norm2(state)) | |
| return state | |
| class StateReadoutModule(nn.Module): | |
| """ | |
| Module for reading information from state. | |
| Implements the State-Readout mechanism from CUT3R. | |
| """ | |
| def __init__( | |
| self, | |
| state_dim: int = 512, | |
| query_dim: int = 512, | |
| output_dim: int = 512, | |
| num_heads: int = 8, | |
| dropout: float = 0.1 | |
| ): | |
| super().__init__() | |
| self.state_dim = state_dim | |
| self.query_dim = query_dim | |
| self.output_dim = output_dim | |
| # Project query to state dimension if needed | |
| if query_dim != state_dim: | |
| self.query_projection = nn.Linear(query_dim, state_dim) | |
| else: | |
| self.query_projection = nn.Identity() | |
| # Multi-head attention for state readout | |
| self.cross_attention = nn.MultiheadAttention( | |
| embed_dim=state_dim, | |
| num_heads=num_heads, | |
| dropout=dropout, | |
| batch_first=True | |
| ) | |
| # Output projection | |
| self.output_projection = nn.Sequential( | |
| nn.Linear(state_dim, output_dim), | |
| nn.LayerNorm(output_dim) | |
| ) | |
| # Feedforward network | |
| self.ffn = nn.Sequential( | |
| nn.Linear(state_dim, state_dim * 4), | |
| nn.GELU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(state_dim * 4, state_dim) | |
| ) | |
| # Layer normalization | |
| self.norm1 = nn.LayerNorm(state_dim) | |
| self.norm2 = nn.LayerNorm(state_dim) | |
| def forward( | |
| self, | |
| state: torch.Tensor, | |
| query: torch.Tensor | |
| ) -> torch.Tensor: | |
| """ | |
| Read information from state based on query. | |
| Args: | |
| state: Current state (B, num_tokens, state_dim) | |
| query: Query features (B, query_len, query_dim) | |
| Returns: | |
| Readout features (B, query_len, output_dim) | |
| """ | |
| # Project query to state dimension | |
| query_proj = self.query_projection(query) | |
| # Cross-attention: query attends to state | |
| attn_output, attention_weights = self.cross_attention( | |
| query=query_proj, | |
| key=state, | |
| value=state | |
| ) | |
| # Residual connection and normalization | |
| readout = query_proj + self.norm1(attn_output) | |
| # Feedforward network with residual connection | |
| readout = readout + self.ffn(self.norm2(readout)) | |
| # Project to output dimension | |
| output = self.output_projection(readout) | |
| return output | |
| class PersistentStateManager(nn.Module): | |
| """ | |
| Complete state management system combining initialization, | |
| update, and readout mechanisms. | |
| """ | |
| def __init__( | |
| self, | |
| num_tokens: int = 256, | |
| token_dim: int = 512, | |
| observation_dim: int = 512, | |
| output_dim: int = 512, | |
| num_heads: int = 8, | |
| dropout: float = 0.1, | |
| learnable_init: bool = True | |
| ): | |
| super().__init__() | |
| # State representation | |
| self.state_repr = StateRepresentation( | |
| num_tokens=num_tokens, | |
| token_dim=token_dim | |
| ) | |
| # State update module | |
| self.state_update = StateUpdateModule( | |
| state_dim=token_dim, | |
| observation_dim=observation_dim, | |
| num_heads=num_heads, | |
| dropout=dropout | |
| ) | |
| # State readout module | |
| self.state_readout = StateReadoutModule( | |
| state_dim=token_dim, | |
| query_dim=observation_dim, | |
| output_dim=output_dim, | |
| num_heads=num_heads, | |
| dropout=dropout | |
| ) | |
| # Current state storage (for sequential processing) | |
| self.register_buffer('current_state', None) | |
| def reset_state(self, batch_size: int): | |
| """Reset state to initial values.""" | |
| self.current_state = self.state_repr.get_initial_state(batch_size) | |
| def forward( | |
| self, | |
| observation: torch.Tensor, | |
| query: Optional[torch.Tensor] = None, | |
| state: Optional[torch.Tensor] = None, | |
| return_updated_state: bool = True | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """ | |
| Process observation and update/read from state. | |
| Args: | |
| observation: New observation features (B, obs_len, obs_dim) | |
| query: Query for readout (B, query_len, obs_dim). | |
| If None, use observation as query. | |
| state: External state if provided. If None, use internal state. | |
| return_updated_state: Whether to return the updated state | |
| Returns: | |
| readout: Features read from state (B, query_len, output_dim) | |
| updated_state: Updated state if return_updated_state=True | |
| """ | |
| batch_size = observation.shape[0] | |
| # Get current state | |
| if state is not None: | |
| current_state = state | |
| elif self.current_state is not None: | |
| current_state = self.current_state | |
| else: | |
| current_state = self.state_repr.get_initial_state(batch_size) | |
| # Update state with new observation | |
| updated_state = self.state_update(current_state, observation) | |
| # Store updated state internally | |
| self.current_state = updated_state | |
| # Read from updated state | |
| if query is None: | |
| query = observation | |
| readout = self.state_readout(updated_state, query) | |
| if return_updated_state: | |
| return readout, updated_state | |
| else: | |
| return readout | |
| if __name__ == "__main__": | |
| """Test the state management modules.""" | |
| print("Testing StateRepresentation...") | |
| state_repr = StateRepresentation(num_tokens=256, token_dim=512) | |
| initial_state = state_repr.get_initial_state(batch_size=4) | |
| print(f"Initial state shape: {initial_state.shape}") # (4, 256, 512) | |
| print("\nTesting StateUpdateModule...") | |
| state_update = StateUpdateModule(state_dim=512, observation_dim=256) | |
| observation = torch.randn(4, 10, 256) # 10 observation tokens | |
| updated_state = state_update(initial_state, observation) | |
| print(f"Updated state shape: {updated_state.shape}") # (4, 256, 512) | |
| print("\nTesting StateReadoutModule...") | |
| state_readout = StateReadoutModule(state_dim=512, query_dim=256, output_dim=512) | |
| query = torch.randn(4, 5, 256) # 5 query tokens | |
| readout = state_readout(updated_state, query) | |
| print(f"Readout shape: {readout.shape}") # (4, 5, 512) | |
| print("\nTesting PersistentStateManager...") | |
| manager = PersistentStateManager( | |
| num_tokens=256, | |
| token_dim=512, | |
| observation_dim=256, | |
| output_dim=512 | |
| ) | |
| # Process sequence | |
| for t in range(3): | |
| obs = torch.randn(4, 10, 256) | |
| readout, state = manager(obs, return_updated_state=True) | |
| print(f"Time {t}: Readout shape {readout.shape}, State shape {state.shape}") | |
| print("\nState management test successful!") | |