| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
| from torch import Tensor, nn |
|
|
|
|
| @dataclass |
| class ChunkDecoderConfig: |
| hidden_dim: int = 512 |
| num_heads: int = 8 |
| num_layers: int = 4 |
| ff_dim: int = 2048 |
| dropout: float = 0.1 |
| chunk_size: int = 8 |
| action_dim: int = 14 |
| arm_action_dim: int = 7 |
| num_candidates: int = 8 |
| num_phases: int = 5 |
| num_arm_roles: int = 4 |
| num_proposal_modes: int = 6 |
| planner_top_k: int = 4 |
|
|
|
|
| class ACTBimanualChunkDecoder(nn.Module): |
| def __init__(self, config: ChunkDecoderConfig) -> None: |
| super().__init__() |
| self.config = config |
| decoder_layer = nn.TransformerDecoderLayer( |
| d_model=config.hidden_dim, |
| nhead=config.num_heads, |
| dim_feedforward=config.ff_dim, |
| dropout=config.dropout, |
| batch_first=True, |
| norm_first=True, |
| ) |
| self.revealer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=config.num_layers) |
| actor_layer = nn.TransformerDecoderLayer( |
| d_model=config.hidden_dim, |
| nhead=config.num_heads, |
| dim_feedforward=config.ff_dim, |
| dropout=config.dropout, |
| batch_first=True, |
| norm_first=True, |
| ) |
| self.actor_decoder = nn.TransformerDecoder(actor_layer, num_layers=config.num_layers) |
| self.query_embed = nn.Embedding(config.chunk_size, config.hidden_dim) |
| self.actor_role_bias = nn.Parameter(torch.zeros(1, config.chunk_size, config.hidden_dim)) |
| self.revealer_mean = nn.Linear(config.hidden_dim, config.arm_action_dim) |
| self.revealer_log_std = nn.Linear(config.hidden_dim, config.arm_action_dim) |
| self.actor_mean = nn.Linear(config.hidden_dim, config.action_dim - config.arm_action_dim) |
| self.actor_log_std = nn.Linear(config.hidden_dim, config.action_dim - config.arm_action_dim) |
| self.coordination = nn.Sequential( |
| nn.LayerNorm(config.hidden_dim * 3), |
| nn.Linear(config.hidden_dim * 3, config.hidden_dim), |
| nn.GELU(), |
| nn.Linear(config.hidden_dim, config.hidden_dim), |
| ) |
| self.proposal_score = nn.Sequential( |
| nn.LayerNorm(config.hidden_dim * 3), |
| nn.Linear(config.hidden_dim * 3, 1), |
| ) |
|
|
| def _deterministic_candidate_noise( |
| self, |
| action_mean: Tensor, |
| num_candidates: int, |
| ) -> Tensor: |
| batch_size, chunk_size, action_dim = action_mean.shape |
| noise = torch.zeros( |
| batch_size, |
| num_candidates, |
| chunk_size, |
| action_dim, |
| device=action_mean.device, |
| dtype=action_mean.dtype, |
| ) |
| if num_candidates <= 1: |
| return noise |
|
|
| candidate_index = torch.arange(1, num_candidates, device=action_mean.device, dtype=action_mean.dtype).view( |
| num_candidates - 1, 1, 1 |
| ) |
| step_index = torch.arange(chunk_size, device=action_mean.device, dtype=action_mean.dtype).view(1, chunk_size, 1) |
| dim_index = torch.arange(action_dim, device=action_mean.device, dtype=action_mean.dtype).view(1, 1, action_dim) |
|
|
| base = torch.sin(candidate_index * 0.73 + step_index * 0.37 + dim_index * 0.19) |
| base = base + torch.cos(candidate_index * 1.11 + step_index * 0.17 + dim_index * 0.41) |
| base = base / base.square().mean(dim=(1, 2), keepdim=True).sqrt().clamp_min(1e-6) |
| noise[:, 1:] = base.unsqueeze(0).expand(batch_size, -1, -1, -1) |
| return noise |
|
|
| def forward( |
| self, |
| scene_tokens: Tensor, |
| reveal_tokens: Tensor | None = None, |
| memory_token: Tensor | None = None, |
| ) -> dict[str, Tensor]: |
| batch_size = scene_tokens.shape[0] |
| query = self.query_embed.weight.unsqueeze(0).expand(batch_size, -1, -1) |
| decoder_memory = scene_tokens |
| if reveal_tokens is not None: |
| decoder_memory = torch.cat([decoder_memory, reveal_tokens], dim=1) |
| if memory_token is not None: |
| decoder_memory = torch.cat([decoder_memory, memory_token], dim=1) |
|
|
| revealer_tokens = self.revealer_decoder(query, decoder_memory) |
| actor_query = query + self.actor_role_bias |
| actor_tokens = self.actor_decoder(actor_query, torch.cat([decoder_memory, revealer_tokens], dim=1)) |
| if reveal_tokens is not None: |
| reveal_context = reveal_tokens.mean(dim=1, keepdim=True).expand(-1, self.config.chunk_size, -1) |
| else: |
| reveal_context = scene_tokens.mean(dim=1, keepdim=True).expand(-1, self.config.chunk_size, -1) |
| coordination_input = torch.cat([revealer_tokens, actor_tokens, reveal_context], dim=-1) |
| coordination = torch.tanh(self.coordination(coordination_input)) |
| revealer_tokens = revealer_tokens + coordination |
| actor_tokens = actor_tokens + coordination |
| action_mean = torch.cat([self.revealer_mean(revealer_tokens), self.actor_mean(actor_tokens)], dim=-1) |
| action_log_std = torch.cat( |
| [ |
| self.revealer_log_std(revealer_tokens), |
| self.actor_log_std(actor_tokens), |
| ], |
| dim=-1, |
| ).clamp(min=-5.0, max=2.0) |
| proposal_features = torch.cat( |
| [ |
| revealer_tokens.mean(dim=1), |
| actor_tokens.mean(dim=1), |
| coordination.mean(dim=1), |
| ], |
| dim=-1, |
| ) |
| return { |
| "decoded_tokens": torch.cat([revealer_tokens, actor_tokens], dim=-1), |
| "revealer_tokens": revealer_tokens, |
| "actor_tokens": actor_tokens, |
| "coordination_tokens": coordination, |
| "action_mean": action_mean, |
| "action_log_std": action_log_std, |
| "proposal_score": self.proposal_score(proposal_features).squeeze(-1), |
| } |
|
|
| def sample_candidates(self, action_mean: Tensor, action_log_std: Tensor, num_candidates: int | None = None) -> Tensor: |
| num_candidates = num_candidates or self.config.num_candidates |
| if num_candidates <= 1: |
| return action_mean.unsqueeze(1) |
| std = action_log_std.exp() |
| if self.training: |
| noise = torch.randn( |
| action_mean.size(0), |
| num_candidates, |
| action_mean.size(1), |
| action_mean.size(2), |
| device=action_mean.device, |
| dtype=action_mean.dtype, |
| ) |
| else: |
| noise = self._deterministic_candidate_noise(action_mean, num_candidates) |
| candidates = action_mean.unsqueeze(1) + noise * std.unsqueeze(1) |
| candidates[:, 0] = action_mean |
| return candidates |
|
|
|
|
| class InteractionChunkDecoder(nn.Module): |
| def __init__(self, config: ChunkDecoderConfig) -> None: |
| super().__init__() |
| self.config = config |
| decoder_layer = nn.TransformerDecoderLayer( |
| d_model=config.hidden_dim, |
| nhead=config.num_heads, |
| dim_feedforward=config.ff_dim, |
| dropout=config.dropout, |
| batch_first=True, |
| norm_first=True, |
| ) |
| self.right_decoder = nn.TransformerDecoder(decoder_layer, num_layers=config.num_layers) |
| left_layer = nn.TransformerDecoderLayer( |
| d_model=config.hidden_dim, |
| nhead=config.num_heads, |
| dim_feedforward=config.ff_dim, |
| dropout=config.dropout, |
| batch_first=True, |
| norm_first=True, |
| ) |
| self.left_decoder = nn.TransformerDecoder(left_layer, num_layers=config.num_layers) |
| self.query_embed = nn.Embedding(config.chunk_size, config.hidden_dim) |
| self.proposal_queries = nn.Embedding(config.num_candidates, config.hidden_dim) |
| self.arm_identity = nn.Embedding(2, config.hidden_dim) |
| self.phase_adapter = nn.Linear(config.num_phases, config.hidden_dim) |
| self.role_adapter = nn.Linear(config.num_arm_roles, config.hidden_dim) |
| self.context_proj = nn.Sequential( |
| nn.LayerNorm(config.hidden_dim), |
| nn.Linear(config.hidden_dim, config.hidden_dim), |
| nn.GELU(), |
| ) |
| self.coordination = nn.Sequential( |
| nn.LayerNorm(config.hidden_dim * 3), |
| nn.Linear(config.hidden_dim * 3, config.hidden_dim), |
| nn.GELU(), |
| nn.Linear(config.hidden_dim, config.hidden_dim), |
| ) |
| self.right_mean = nn.Linear(config.hidden_dim, config.arm_action_dim) |
| self.right_log_std = nn.Linear(config.hidden_dim, config.arm_action_dim) |
| self.left_mean = nn.Linear(config.hidden_dim, config.action_dim - config.arm_action_dim) |
| self.left_log_std = nn.Linear(config.hidden_dim, config.action_dim - config.arm_action_dim) |
| self.proposal_score = nn.Sequential( |
| nn.LayerNorm(config.hidden_dim * 3), |
| nn.Linear(config.hidden_dim * 3, config.hidden_dim), |
| nn.GELU(), |
| nn.Linear(config.hidden_dim, 1), |
| ) |
|
|
| def _conditioning( |
| self, |
| interaction_state: dict[str, Tensor] | None, |
| batch_size: int, |
| device: torch.device, |
| dtype: torch.dtype, |
| ) -> tuple[Tensor, Tensor, Tensor | None]: |
| if interaction_state is None: |
| zero_phase = torch.zeros(batch_size, self.config.hidden_dim, device=device, dtype=dtype) |
| zero_roles = torch.zeros(batch_size, 2, self.config.hidden_dim, device=device, dtype=dtype) |
| return zero_phase, zero_roles, None |
| phase_probs = interaction_state["phase_logits"].softmax(dim=-1).to(dtype=dtype) |
| arm_role_probs = interaction_state["arm_role_logits"].softmax(dim=-1).to(dtype=dtype) |
| phase_context = self.phase_adapter(phase_probs) |
| role_context = self.role_adapter(arm_role_probs) |
| return phase_context, role_context, interaction_state.get("interaction_tokens") |
|
|
| def _decode_from_queries( |
| self, |
| queries: Tensor, |
| decoder_memory: Tensor, |
| phase_context: Tensor, |
| role_context: Tensor, |
| interaction_context: Tensor, |
| ) -> dict[str, Tensor]: |
| phase_bias = phase_context.unsqueeze(1) |
| right_queries = ( |
| queries |
| + phase_bias |
| + role_context[:, 0].unsqueeze(1) |
| + self.arm_identity.weight[0].view(1, 1, -1).to(dtype=queries.dtype) |
| ) |
| left_queries = ( |
| queries |
| + phase_bias |
| + role_context[:, 1].unsqueeze(1) |
| + self.arm_identity.weight[1].view(1, 1, -1).to(dtype=queries.dtype) |
| ) |
| right_tokens = self.right_decoder(right_queries, decoder_memory) |
| left_tokens = self.left_decoder(left_queries, torch.cat([decoder_memory, right_tokens], dim=1)) |
| context = interaction_context.unsqueeze(1).expand(-1, queries.shape[1], -1) |
| coordination_input = torch.cat([right_tokens, left_tokens, context], dim=-1) |
| coordination = torch.tanh(self.coordination(coordination_input)) |
| right_tokens = right_tokens + coordination |
| left_tokens = left_tokens + coordination |
| action_mean = torch.cat([self.right_mean(right_tokens), self.left_mean(left_tokens)], dim=-1) |
| action_log_std = torch.cat( |
| [self.right_log_std(right_tokens), self.left_log_std(left_tokens)], |
| dim=-1, |
| ).clamp(min=-5.0, max=2.0) |
| pooled_features = torch.cat( |
| [right_tokens.mean(dim=1), left_tokens.mean(dim=1), coordination.mean(dim=1)], |
| dim=-1, |
| ) |
| return { |
| "right_tokens": right_tokens, |
| "left_tokens": left_tokens, |
| "coordination_tokens": coordination, |
| "decoded_tokens": torch.cat([right_tokens, left_tokens], dim=-1), |
| "action_mean": action_mean, |
| "action_log_std": action_log_std, |
| "proposal_score": self.proposal_score(pooled_features).squeeze(-1), |
| } |
|
|
| def forward( |
| self, |
| scene_tokens: Tensor, |
| interaction_state: dict[str, Tensor] | None = None, |
| memory_tokens: Tensor | None = None, |
| reveal_tokens: Tensor | None = None, |
| memory_token: Tensor | None = None, |
| ) -> dict[str, Tensor]: |
| if memory_tokens is None: |
| memory_tokens = memory_token |
| batch_size = scene_tokens.shape[0] |
| dtype = scene_tokens.dtype |
| phase_context, role_context, interaction_tokens = self._conditioning( |
| interaction_state=interaction_state, |
| batch_size=batch_size, |
| device=scene_tokens.device, |
| dtype=dtype, |
| ) |
|
|
| decoder_memory = scene_tokens |
| if interaction_tokens is not None: |
| decoder_memory = torch.cat([decoder_memory, interaction_tokens], dim=1) |
| elif reveal_tokens is not None: |
| decoder_memory = torch.cat([decoder_memory, reveal_tokens], dim=1) |
| if memory_tokens is not None: |
| decoder_memory = torch.cat([decoder_memory, memory_tokens], dim=1) |
|
|
| if interaction_tokens is not None and interaction_tokens.numel() > 0: |
| interaction_context = interaction_tokens.mean(dim=1) |
| elif reveal_tokens is not None and reveal_tokens.numel() > 0: |
| interaction_context = reveal_tokens.mean(dim=1) |
| else: |
| interaction_context = scene_tokens.mean(dim=1) |
| interaction_context = self.context_proj(interaction_context) |
|
|
| base_queries = self.query_embed.weight.unsqueeze(0).expand(batch_size, -1, -1) |
| decoded = self._decode_from_queries( |
| queries=base_queries, |
| decoder_memory=decoder_memory, |
| phase_context=phase_context, |
| role_context=role_context, |
| interaction_context=interaction_context, |
| ) |
|
|
| num_candidates = self.config.num_candidates |
| proposal_bias = self.proposal_queries.weight.view(1, num_candidates, 1, -1).expand( |
| batch_size, -1, self.config.chunk_size, -1 |
| ) |
| candidate_queries = base_queries.unsqueeze(1) + proposal_bias |
| flat_queries = candidate_queries.reshape(batch_size * num_candidates, self.config.chunk_size, self.config.hidden_dim) |
| flat_memory = decoder_memory.unsqueeze(1).expand(-1, num_candidates, -1, -1).reshape( |
| batch_size * num_candidates, decoder_memory.shape[1], decoder_memory.shape[2] |
| ) |
| flat_phase = phase_context.unsqueeze(1).expand(-1, num_candidates, -1).reshape( |
| batch_size * num_candidates, self.config.hidden_dim |
| ) |
| flat_roles = role_context.unsqueeze(1).expand(-1, num_candidates, -1, -1).reshape( |
| batch_size * num_candidates, 2, self.config.hidden_dim |
| ) |
| flat_context = interaction_context.unsqueeze(1).expand(-1, num_candidates, -1).reshape( |
| batch_size * num_candidates, self.config.hidden_dim |
| ) |
| candidate_decoded = self._decode_from_queries( |
| queries=flat_queries, |
| decoder_memory=flat_memory, |
| phase_context=flat_phase, |
| role_context=flat_roles, |
| interaction_context=flat_context, |
| ) |
|
|
| proposal_deltas = candidate_decoded["action_mean"].view( |
| batch_size, |
| num_candidates, |
| self.config.chunk_size, |
| self.config.action_dim, |
| ) |
| proposal_logits = candidate_decoded["proposal_score"].view(batch_size, num_candidates) |
| proposal_candidates = decoded["action_mean"].unsqueeze(1) + 0.35 * torch.tanh(proposal_deltas) |
| proposal_candidates[:, 0] = decoded["action_mean"] |
| proposal_logits[:, 0] = decoded["proposal_score"] |
| decoded["proposal_candidates"] = proposal_candidates |
| decoded["proposal_logits"] = proposal_logits |
| return decoded |
|
|
| def sample_candidates( |
| self, |
| action_mean: Tensor, |
| action_log_std: Tensor, |
| num_candidates: int | None = None, |
| proposal_candidates: Tensor | None = None, |
| ) -> Tensor: |
| if proposal_candidates is not None: |
| return proposal_candidates |
| num_candidates = num_candidates or self.config.num_candidates |
| if num_candidates <= 1: |
| return action_mean.unsqueeze(1) |
| noise = torch.randn( |
| action_mean.size(0), |
| num_candidates, |
| action_mean.size(1), |
| action_mean.size(2), |
| device=action_mean.device, |
| dtype=action_mean.dtype, |
| ) |
| candidates = action_mean.unsqueeze(1) + noise * action_log_std.exp().unsqueeze(1) |
| candidates[:, 0] = action_mean |
| return candidates |
|
|
|
|
| DEFAULT_PROPOSAL_MODES = ( |
| "widen_opening", |
| "maintain_opening", |
| "slide_occluder", |
| "lift_support_layer", |
| "stabilize_support", |
| "retrieve", |
| ) |
|
|
| TASK_PROPOSAL_MODES = { |
| "foliage": ( |
| "sweep_left", |
| "sweep_right", |
| "pin_canopy", |
| "widen_gap", |
| "maintain_gap", |
| "insert_actor", |
| "retrieve", |
| ), |
| "bag": ( |
| "pin_left_rim", |
| "pin_right_rim", |
| "widen_mouth", |
| "maintain_mouth", |
| "probe_inside", |
| "insert_actor", |
| "retrieve", |
| ), |
| "cloth": ( |
| "lift_edge", |
| "separate_layer", |
| "stabilize_fold", |
| "maintain_lift", |
| "insert_actor", |
| "retrieve", |
| ), |
| } |
|
|
| TASK_INDEX = {"foliage": 0, "bag": 1, "cloth": 2} |
|
|
|
|
| def infer_task_name_from_text(text: str | None) -> str: |
| if not text: |
| return "generic" |
| lowered = text.lower() |
| if any(token in lowered for token in ("foliage", "canopy", "leaf", "leaves", "snail")): |
| return "foliage" |
| if any(token in lowered for token in ("bag", "mouth", "rim", "aperture")): |
| return "bag" |
| if any(token in lowered for token in ("cloth", "fold", "layer", "suitcase", "garment")): |
| return "cloth" |
| return "generic" |
|
|
|
|
| def proposal_mode_vocab(task_name: str, num_modes: int) -> tuple[str, ...]: |
| if task_name == "generic": |
| base_vocab = tuple(DEFAULT_PROPOSAL_MODES) |
| else: |
| vocab = TASK_PROPOSAL_MODES[task_name] |
| if len(vocab) > num_modes: |
| if num_modes >= 6: |
| base_vocab = ( |
| vocab[0], |
| vocab[1], |
| vocab[2], |
| vocab[3], |
| vocab[-2], |
| vocab[-1], |
| )[:num_modes] |
| else: |
| base_vocab = vocab[:num_modes] |
| else: |
| base_vocab = vocab |
| if len(base_vocab) >= num_modes: |
| return tuple(base_vocab[:num_modes]) |
| if not base_vocab: |
| return tuple("retrieve" for _ in range(num_modes)) |
| padded = list(base_vocab) |
| while len(padded) < num_modes: |
| padded.append(base_vocab[-1]) |
| return tuple(padded) |
|
|
|
|
| def swap_arm_action_order(action_chunk: Tensor) -> Tensor: |
| midpoint = action_chunk.shape[-1] // 2 |
| return torch.cat([action_chunk[..., midpoint:], action_chunk[..., :midpoint]], dim=-1) |
|
|
|
|
| class SymmetricCoordinatedChunkDecoder(nn.Module): |
| def __init__(self, config: ChunkDecoderConfig) -> None: |
| super().__init__() |
| self.config = config |
| proposal_context_dim = config.action_dim + (config.hidden_dim * 2) |
| decoder_layer = nn.TransformerDecoderLayer( |
| d_model=config.hidden_dim, |
| nhead=config.num_heads, |
| dim_feedforward=config.ff_dim, |
| dropout=config.dropout, |
| batch_first=True, |
| norm_first=True, |
| ) |
| self.arm_decoder = nn.TransformerDecoder(decoder_layer, num_layers=config.num_layers) |
| self.query_embed = nn.Embedding(config.chunk_size, config.hidden_dim) |
| self.arm_identity = nn.Embedding(2, config.hidden_dim) |
| self.task_embedding = nn.Embedding(len(TASK_INDEX), config.hidden_dim) |
| self.phase_adapter = nn.Linear(config.num_phases, config.hidden_dim) |
| self.role_adapter = nn.Linear(config.num_arm_roles, config.hidden_dim) |
| self.context_proj = nn.Sequential( |
| nn.LayerNorm(config.hidden_dim), |
| nn.Linear(config.hidden_dim, config.hidden_dim), |
| nn.GELU(), |
| ) |
| self.coordination = nn.Sequential( |
| nn.LayerNorm(config.hidden_dim * 3), |
| nn.Linear(config.hidden_dim * 3, config.hidden_dim), |
| nn.GELU(), |
| nn.Linear(config.hidden_dim, config.hidden_dim), |
| ) |
| self.arm_head = nn.Sequential( |
| nn.LayerNorm(config.hidden_dim), |
| nn.Linear(config.hidden_dim, config.hidden_dim), |
| nn.GELU(), |
| ) |
| self.arm_mean = nn.Linear(config.hidden_dim, config.arm_action_dim) |
| self.arm_log_std = nn.Linear(config.hidden_dim, config.arm_action_dim) |
| self.proposal_mode_head = nn.Sequential( |
| nn.LayerNorm(proposal_context_dim), |
| nn.Linear(proposal_context_dim, config.hidden_dim), |
| nn.GELU(), |
| nn.Linear(config.hidden_dim, config.num_proposal_modes), |
| ) |
| self.proposal_mode_embeddings = nn.Embedding(config.num_proposal_modes, config.hidden_dim) |
| self.proposal_slot_embeddings = nn.Embedding(config.num_candidates, config.hidden_dim) |
| self.mode_residual_heads = nn.ModuleList( |
| [ |
| nn.Sequential( |
| nn.LayerNorm(proposal_context_dim), |
| nn.Linear(proposal_context_dim, config.hidden_dim), |
| nn.GELU(), |
| nn.Linear(config.hidden_dim, config.chunk_size * config.action_dim), |
| ) |
| for _ in range(config.num_proposal_modes) |
| ] |
| ) |
| self.slot_delta = nn.Sequential( |
| nn.LayerNorm(config.hidden_dim), |
| nn.Linear(config.hidden_dim, config.hidden_dim), |
| nn.GELU(), |
| nn.Linear(config.hidden_dim, config.chunk_size * config.action_dim), |
| ) |
| self.proposal_score = nn.Sequential( |
| nn.LayerNorm(proposal_context_dim + config.hidden_dim), |
| nn.Linear(proposal_context_dim + config.hidden_dim, config.hidden_dim), |
| nn.GELU(), |
| nn.Linear(config.hidden_dim, 1), |
| ) |
|
|
| def _conditioning( |
| self, |
| interaction_state: dict[str, Tensor] | None, |
| batch_size: int, |
| device: torch.device, |
| dtype: torch.dtype, |
| swap_roles: bool = False, |
| ) -> tuple[Tensor, Tensor, Tensor]: |
| if interaction_state is None: |
| zero_phase = torch.zeros(batch_size, self.config.hidden_dim, device=device, dtype=dtype) |
| zero_roles = torch.zeros(batch_size, 2, self.config.hidden_dim, device=device, dtype=dtype) |
| zero_context = torch.zeros(batch_size, self.config.hidden_dim, device=device, dtype=dtype) |
| return zero_phase, zero_roles, zero_context |
| phase_probs = interaction_state["phase_logits"].softmax(dim=-1).to(dtype=dtype) |
| arm_role_probs = interaction_state["arm_role_logits"].softmax(dim=-1).to(dtype=dtype) |
| if swap_roles: |
| arm_role_probs = arm_role_probs.flip(1) |
| phase_context = self.phase_adapter(phase_probs) |
| role_context = self.role_adapter(arm_role_probs) |
| if interaction_state.get("interaction_tokens") is not None: |
| interaction_context = interaction_state["interaction_tokens"].mean(dim=1) |
| else: |
| interaction_context = interaction_state["field_tokens"].mean(dim=1) |
| return phase_context, role_context, self.context_proj(interaction_context) |
|
|
| def _decode_arm_tokens( |
| self, |
| queries: Tensor, |
| decoder_memory: Tensor, |
| phase_context: Tensor, |
| role_context: Tensor, |
| interaction_context: Tensor, |
| swap_roles: bool = False, |
| ) -> tuple[Tensor, Tensor, Tensor]: |
| batch_size, chunk_size, _ = queries.shape |
| identity_order = torch.tensor([1, 0], device=queries.device) if swap_roles else torch.tensor([0, 1], device=queries.device) |
| arm_queries = queries.unsqueeze(1).expand(-1, 2, -1, -1) |
| arm_queries = arm_queries + phase_context.unsqueeze(1).unsqueeze(2) |
| arm_queries = arm_queries + role_context.unsqueeze(2) |
| arm_queries = arm_queries + self.arm_identity(identity_order).view(1, 2, 1, -1).to(dtype=queries.dtype) |
| flat_queries = arm_queries.reshape(batch_size * 2, chunk_size, self.config.hidden_dim) |
| flat_memory = decoder_memory.unsqueeze(1).expand(-1, 2, -1, -1).reshape( |
| batch_size * 2, |
| decoder_memory.shape[1], |
| decoder_memory.shape[2], |
| ) |
| decoded = self.arm_decoder(flat_queries, flat_memory).reshape(batch_size, 2, chunk_size, self.config.hidden_dim) |
| coordination_input = torch.cat( |
| [ |
| decoded[:, 0], |
| decoded[:, 1], |
| interaction_context.unsqueeze(1).expand(-1, chunk_size, -1), |
| ], |
| dim=-1, |
| ) |
| coordination = torch.tanh(self.coordination(coordination_input)) |
| decoded[:, 0] = decoded[:, 0] + coordination |
| decoded[:, 1] = decoded[:, 1] + coordination |
| decoded = self.arm_head(decoded) |
| arm_mean = self.arm_mean(decoded) |
| arm_log_std = self.arm_log_std(decoded).clamp(min=-5.0, max=2.0) |
| return arm_mean, arm_log_std, coordination |
|
|
| def _proposal_outputs( |
| self, |
| base_action: Tensor, |
| pooled_context: Tensor, |
| task_names: list[str], |
| ) -> tuple[Tensor, Tensor, Tensor, list[list[str]]]: |
| batch_size = pooled_context.shape[0] |
| mode_logits = self.proposal_mode_head(pooled_context) |
| mode_residuals = [] |
| for head in self.mode_residual_heads: |
| residual = head(pooled_context).view(batch_size, self.config.chunk_size, self.config.action_dim) |
| mode_residuals.append(residual) |
| mode_residuals = torch.stack(mode_residuals, dim=1) |
|
|
| mode_assignments = torch.arange(self.config.num_candidates, device=pooled_context.device) % self.config.num_proposal_modes |
| slot_embeddings = self.proposal_slot_embeddings.weight |
| slot_deltas = self.slot_delta(slot_embeddings).view( |
| self.config.num_candidates, |
| self.config.chunk_size, |
| self.config.action_dim, |
| ) |
| proposal_candidates = [] |
| proposal_logits = [] |
| proposal_mode_names = [ |
| [ |
| proposal_mode_vocab(task_name, self.config.num_proposal_modes)[int(mode_assignments[slot_idx])] |
| for slot_idx in range(self.config.num_candidates) |
| ] |
| for task_name in task_names |
| ] |
| for slot_idx in range(self.config.num_candidates): |
| mode_idx = int(mode_assignments[slot_idx]) |
| candidate = base_action + 0.35 * torch.tanh(mode_residuals[:, mode_idx]) + 0.05 * torch.tanh(slot_deltas[slot_idx]).unsqueeze(0) |
| proposal_candidates.append(candidate) |
| score_features = torch.cat( |
| [ |
| pooled_context, |
| self.proposal_mode_embeddings.weight[mode_idx].unsqueeze(0).expand(batch_size, -1) |
| + slot_embeddings[slot_idx].unsqueeze(0).expand(batch_size, -1), |
| ], |
| dim=-1, |
| ) |
| proposal_logits.append( |
| self.proposal_score(score_features).squeeze(-1) + mode_logits[:, mode_idx] |
| ) |
| stacked_candidates = torch.stack(proposal_candidates, dim=1) |
| stacked_logits = torch.stack(proposal_logits, dim=1) |
| stacked_candidates[:, 0] = base_action |
| return stacked_candidates, stacked_logits, mode_logits, proposal_mode_names |
|
|
| def forward( |
| self, |
| scene_tokens: Tensor, |
| interaction_state: dict[str, Tensor] | None = None, |
| memory_tokens: Tensor | None = None, |
| reveal_tokens: Tensor | None = None, |
| memory_token: Tensor | None = None, |
| compute_equivariance_probe: bool = False, |
| task_names: list[str] | None = None, |
| ) -> dict[str, Tensor]: |
| if memory_tokens is None: |
| memory_tokens = memory_token |
| batch_size = scene_tokens.shape[0] |
| dtype = scene_tokens.dtype |
| phase_context, role_context, interaction_context = self._conditioning( |
| interaction_state=interaction_state, |
| batch_size=batch_size, |
| device=scene_tokens.device, |
| dtype=dtype, |
| ) |
|
|
| decoder_memory = scene_tokens |
| interaction_tokens = interaction_state.get("interaction_tokens") if interaction_state is not None else None |
| if interaction_tokens is not None: |
| decoder_memory = torch.cat([decoder_memory, interaction_tokens], dim=1) |
| elif reveal_tokens is not None: |
| decoder_memory = torch.cat([decoder_memory, reveal_tokens], dim=1) |
| if memory_tokens is not None: |
| decoder_memory = torch.cat([decoder_memory, memory_tokens], dim=1) |
|
|
| canonical_task_names = [infer_task_name_from_text(name) for name in (task_names or ["generic"] * batch_size)] |
| task_ids = torch.as_tensor( |
| [TASK_INDEX[name] for name in canonical_task_names if name in TASK_INDEX], |
| device=scene_tokens.device, |
| dtype=torch.long, |
| ) |
| if task_ids.numel() != batch_size: |
| task_ids = torch.as_tensor( |
| [TASK_INDEX.get(name, 0) for name in canonical_task_names], |
| device=scene_tokens.device, |
| dtype=torch.long, |
| ) |
| interaction_context = interaction_context + self.task_embedding(task_ids) |
|
|
| base_queries = self.query_embed.weight.unsqueeze(0).expand(batch_size, -1, -1) |
| arm_mean, arm_log_std, coordination = self._decode_arm_tokens( |
| queries=base_queries, |
| decoder_memory=decoder_memory, |
| phase_context=phase_context, |
| role_context=role_context, |
| interaction_context=interaction_context, |
| ) |
| action_mean = torch.cat([arm_mean[:, 0], arm_mean[:, 1]], dim=-1) |
| action_log_std = torch.cat([arm_log_std[:, 0], arm_log_std[:, 1]], dim=-1) |
| pooled_context = torch.cat( |
| [ |
| arm_mean[:, 0].mean(dim=1), |
| arm_mean[:, 1].mean(dim=1), |
| coordination.mean(dim=1), |
| interaction_context, |
| ], |
| dim=-1, |
| ) |
| proposal_candidates, proposal_logits, proposal_mode_logits, proposal_mode_names = self._proposal_outputs( |
| action_mean, |
| pooled_context, |
| canonical_task_names, |
| ) |
|
|
| outputs = { |
| "decoded_tokens": torch.cat([arm_mean[:, 0], arm_mean[:, 1]], dim=-1), |
| "right_tokens": arm_mean[:, 0], |
| "left_tokens": arm_mean[:, 1], |
| "coordination_tokens": coordination, |
| "action_mean": action_mean, |
| "action_log_std": action_log_std, |
| "proposal_candidates": proposal_candidates, |
| "proposal_logits": proposal_logits, |
| "proposal_mode_logits": proposal_mode_logits, |
| "proposal_mode_assignments": torch.arange( |
| self.config.num_candidates, |
| device=scene_tokens.device, |
| ) % self.config.num_proposal_modes, |
| "proposal_mode_names": proposal_mode_names, |
| "proposal_task_names": canonical_task_names, |
| } |
| if compute_equivariance_probe: |
| swapped_phase, swapped_roles, swapped_context = self._conditioning( |
| interaction_state=interaction_state, |
| batch_size=batch_size, |
| device=scene_tokens.device, |
| dtype=dtype, |
| swap_roles=True, |
| ) |
| swapped_arm_mean, _, _ = self._decode_arm_tokens( |
| queries=base_queries, |
| decoder_memory=decoder_memory, |
| phase_context=swapped_phase, |
| role_context=swapped_roles, |
| interaction_context=swapped_context, |
| swap_roles=True, |
| ) |
| outputs["equivariance_probe_action_mean"] = torch.cat( |
| [swapped_arm_mean[:, 0], swapped_arm_mean[:, 1]], |
| dim=-1, |
| ) |
| outputs["equivariance_target_action_mean"] = swap_arm_action_order(action_mean) |
| return outputs |
|
|
| def sample_candidates( |
| self, |
| action_mean: Tensor, |
| action_log_std: Tensor, |
| num_candidates: int | None = None, |
| proposal_candidates: Tensor | None = None, |
| ) -> Tensor: |
| if proposal_candidates is not None: |
| return proposal_candidates |
| num_candidates = num_candidates or self.config.num_candidates |
| if num_candidates <= 1: |
| return action_mean.unsqueeze(1) |
| noise = torch.randn( |
| action_mean.size(0), |
| num_candidates, |
| action_mean.size(1), |
| action_mean.size(2), |
| device=action_mean.device, |
| dtype=action_mean.dtype, |
| ) |
| candidates = action_mean.unsqueeze(1) + noise * action_log_std.exp().unsqueeze(1) |
| candidates[:, 0] = action_mean |
| return candidates |
|
|
