| """A2C2 correction head architecture for cached BEHAVIOR/OpenPI features.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| import math |
|
|
| import torch |
| from torch import Tensor, nn |
|
|
|
|
| @dataclass(frozen=True) |
| class A2C2CorrectionHeadConfig: |
| state_dim: int = 256 |
| action_dim: int = 23 |
| action_horizon: int = 32 |
| base_policy_z_dim: int = 2048 |
| use_base_policy_z: bool = True |
| time_dim: int = 2 |
|
|
| dim_model: int = 512 |
| n_heads: int = 8 |
| n_encoder_layers: int = 6 |
| dim_feedforward: int = 2048 |
| dropout: float = 0.1 |
| mlp_hidden_dim: int = 1024 |
|
|
|
|
| def _sinusoidal_positions(length: int, dim: int) -> Tensor: |
| if dim % 2 != 0: |
| raise ValueError("dim must be even for sinusoidal positional encoding.") |
|
|
| position = torch.arange(length, dtype=torch.float32).unsqueeze(1) |
| div_term = torch.exp(torch.arange(0, dim, 2, dtype=torch.float32) * (-math.log(10000.0) / dim)) |
| pe = torch.zeros(length, dim, dtype=torch.float32) |
| pe[:, 0::2] = torch.sin(position * div_term) |
| pe[:, 1::2] = torch.cos(position * div_term) |
| return pe |
|
|
|
|
| class A2C2CorrectionHead(nn.Module): |
| """Transformer + MLP correction head following the A2C2 residual design.""" |
|
|
| def __init__(self, config: A2C2CorrectionHeadConfig | None = None) -> None: |
| super().__init__() |
| self.config = config or A2C2CorrectionHeadConfig() |
| cfg = self.config |
|
|
| self.cls_token = nn.Parameter(torch.zeros(1, 1, cfg.dim_model)) |
| self.type_embedding = nn.Parameter(torch.zeros(6, cfg.dim_model)) |
|
|
| self.state_proj = nn.Linear(cfg.state_dim, cfg.dim_model) |
| if cfg.use_base_policy_z: |
| self.z_proj = nn.Linear(cfg.base_policy_z_dim, cfg.dim_model) |
| self.time_proj = nn.Linear(cfg.time_dim, cfg.dim_model) |
| self.action_proj = nn.Linear(cfg.action_dim, cfg.dim_model) |
|
|
| chunk_pos = _sinusoidal_positions(cfg.action_horizon, cfg.dim_model) |
| self.register_buffer("chunk_pos_embedding", chunk_pos, persistent=False) |
|
|
| encoder_layer = nn.TransformerEncoderLayer( |
| d_model=cfg.dim_model, |
| nhead=cfg.n_heads, |
| dim_feedforward=cfg.dim_feedforward, |
| dropout=cfg.dropout, |
| activation="gelu", |
| batch_first=True, |
| norm_first=True, |
| ) |
| self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=cfg.n_encoder_layers) |
| self.encoder_norm = nn.LayerNorm(cfg.dim_model) |
|
|
| head_input_token_count = 5 if cfg.use_base_policy_z else 4 |
| head_input_dim = cfg.dim_model * head_input_token_count + cfg.action_dim |
| self.residual_head = nn.Sequential( |
| nn.Linear(head_input_dim, cfg.mlp_hidden_dim), |
| nn.GELU(), |
| nn.Dropout(cfg.dropout), |
| nn.Linear(cfg.mlp_hidden_dim, cfg.mlp_hidden_dim), |
| nn.GELU(), |
| nn.Dropout(cfg.dropout), |
| nn.Linear(cfg.mlp_hidden_dim, cfg.action_dim), |
| ) |
|
|
| self._reset_parameters() |
|
|
| @staticmethod |
| def make_time_feature(chunk_index: Tensor, horizon: int) -> Tensor: |
| """Create [sin, cos] phase features from chunk indices.""" |
|
|
| idx = chunk_index.to(dtype=torch.float32) |
| denom = max(horizon - 1, 1) |
| phase = 2.0 * math.pi * idx / denom |
| return torch.stack([torch.sin(phase), torch.cos(phase)], dim=-1) |
|
|
| def forward( |
| self, |
| observation_state: Tensor, |
| selected_base_action: Tensor, |
| base_action_chunk: Tensor, |
| base_policy_z: Tensor, |
| time_feature: Tensor, |
| valid_action_mask: Tensor | None = None, |
| ) -> Tensor: |
| """Predict residual action delta. |
| |
| Args: |
| observation_state: [B, state_dim] |
| selected_base_action: [B, action_dim], the current action being corrected. |
| base_action_chunk: [B, H, action_dim] |
| base_policy_z: [B, z_dim] |
| time_feature: [B, 2] |
| valid_action_mask: optional bool tensor [B, H], True for valid chunk |
| entries. Invalid chunk entries are ignored by transformer attention. |
| |
| Returns: |
| Tensor [B, action_dim], the predicted residual delta. |
| """ |
|
|
| cfg = self.config |
| batch_size = observation_state.shape[0] |
| device = observation_state.device |
| dtype = observation_state.dtype |
|
|
| self._validate_inputs(observation_state, selected_base_action, base_action_chunk, base_policy_z, time_feature) |
|
|
| cls = self.cls_token.to(device=device, dtype=dtype).expand(batch_size, -1, -1) |
| cls = cls + self.type_embedding[0].to(device=device, dtype=dtype) |
|
|
| state_token = self.state_proj(observation_state).unsqueeze(1) |
| state_token = state_token + self.type_embedding[1].to(device=device, dtype=dtype) |
|
|
| time_token = self.time_proj(time_feature).unsqueeze(1) |
| time_token = time_token + self.type_embedding[3].to(device=device, dtype=dtype) |
|
|
| selected_action_token = self.action_proj(selected_base_action).unsqueeze(1) |
| selected_action_token = selected_action_token + self.type_embedding[4].to(device=device, dtype=dtype) |
|
|
| chunk_tokens = self.action_proj(base_action_chunk) |
| chunk_pos = self.chunk_pos_embedding[: base_action_chunk.shape[1]].to(device=device, dtype=dtype) |
| chunk_tokens = chunk_tokens + chunk_pos.unsqueeze(0) |
| chunk_tokens = chunk_tokens + self.type_embedding[5].to(device=device, dtype=dtype) |
|
|
| prefix_tokens = [cls, state_token] |
| if cfg.use_base_policy_z: |
| z_token = self.z_proj(base_policy_z).unsqueeze(1) |
| z_token = z_token + self.type_embedding[2].to(device=device, dtype=dtype) |
| prefix_tokens.append(z_token) |
| prefix_tokens.extend([time_token, selected_action_token]) |
|
|
| tokens = torch.cat([*prefix_tokens, chunk_tokens], dim=1) |
|
|
| padding_mask = None |
| if valid_action_mask is not None: |
| valid_action_mask = valid_action_mask.to(device=device, dtype=torch.bool) |
| prefix_mask = torch.zeros(batch_size, len(prefix_tokens), device=device, dtype=torch.bool) |
| padding_mask = torch.cat([prefix_mask, ~valid_action_mask], dim=1) |
|
|
| encoded = self.encoder(tokens, src_key_padding_mask=padding_mask) |
| encoded = self.encoder_norm(encoded) |
|
|
| cls_state = encoded[:, 0] |
| state_state = encoded[:, 1] |
| if cfg.use_base_policy_z: |
| z_state = encoded[:, 2] |
| time_state = encoded[:, 3] |
| selected_action_state = encoded[:, 4] |
| head_states = [cls_state, state_state, z_state, time_state, selected_action_state] |
| else: |
| time_state = encoded[:, 2] |
| selected_action_state = encoded[:, 3] |
| head_states = [cls_state, state_state, time_state, selected_action_state] |
|
|
| head_input = torch.cat( |
| [*head_states, selected_base_action], |
| dim=-1, |
| ) |
| return self.residual_head(head_input) |
|
|
| def _validate_inputs( |
| self, |
| observation_state: Tensor, |
| selected_base_action: Tensor, |
| base_action_chunk: Tensor, |
| base_policy_z: Tensor, |
| time_feature: Tensor, |
| ) -> None: |
| cfg = self.config |
| if observation_state.ndim != 2 or observation_state.shape[-1] != cfg.state_dim: |
| raise ValueError(f"observation_state must have shape [B, {cfg.state_dim}].") |
| if selected_base_action.ndim != 2 or selected_base_action.shape[-1] != cfg.action_dim: |
| raise ValueError(f"selected_base_action must have shape [B, {cfg.action_dim}].") |
| if base_action_chunk.ndim != 3 or base_action_chunk.shape[-1] != cfg.action_dim: |
| raise ValueError(f"base_action_chunk must have shape [B, H, {cfg.action_dim}].") |
| if base_action_chunk.shape[1] > cfg.action_horizon: |
| raise ValueError(f"base_action_chunk horizon cannot exceed {cfg.action_horizon}.") |
| if cfg.use_base_policy_z and (base_policy_z.ndim != 2 or base_policy_z.shape[-1] != cfg.base_policy_z_dim): |
| raise ValueError(f"base_policy_z must have shape [B, {cfg.base_policy_z_dim}].") |
| if time_feature.ndim != 2 or time_feature.shape[-1] != cfg.time_dim: |
| raise ValueError(f"time_feature must have shape [B, {cfg.time_dim}].") |
|
|
| def _reset_parameters(self) -> None: |
| nn.init.trunc_normal_(self.cls_token, std=0.02) |
| nn.init.trunc_normal_(self.type_embedding, std=0.02) |
| for module in self.modules(): |
| if isinstance(module, nn.Linear): |
| nn.init.xavier_uniform_(module.weight) |
| if module.bias is not None: |
| nn.init.zeros_(module.bias) |
|
|
