ZibinDong/ActionCodec-bridge

1import math2from copy import deepcopy3from typing import List, Literal, Optional, Tuple, Union4 5import einops6import numpy as np7import torch8import torch.nn as nn9import torch.nn.functional as F10 11from .configuration_actioncodec import ActionCodecConfig12 13 14def apply_rotary_pos_emb(x: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor:15    original_dtype = x.dtype16 17    x = x.to(torch.float32)18    sin = sin.to(torch.float32)19    cos = cos.to(torch.float32)20 21    x1 = x[..., 0::2]22    x2 = x[..., 1::2]23 24    rotated_x1 = x1 * cos - x2 * sin25    rotated_x2 = x1 * sin + x2 * cos26 27    x_out = torch.empty_like(x)28    x_out[..., 0::2] = rotated_x129    x_out[..., 1::2] = rotated_x230 31    return x_out.to(original_dtype)32 33 34def attention_op(35    q: torch.Tensor,36    k: torch.Tensor,37    v: torch.Tensor,38    mask: torch.Tensor | None = None,39    is_causal: bool = False,40) -> torch.Tensor:41    """42 43    Args:44        q (torch.Tensor): (*b, h, l, d)45        k (torch.Tensor): (*b, k, s, d)46        v (torch.Tensor): (*b, k, s, d)47        mask (torch.Tensor | None, optional): (*b, l, s), where `True` indicates the element should take part in attention. Defaults to None.48        is_causal (bool, optional): Whether to apply causal mask. Defaults to False.49 50    Returns:51        torch.Tensor: (*b, h, l, d)52    """53    heads, kv_heads = q.shape[-3], k.shape[-3]54    if heads != kv_heads:55        assert heads % kv_heads == 0, f"q_heads must be divisible by kv_heads, but got {heads} and {kv_heads}"56        heads_per_kv_head = heads // kv_heads57        k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim=1), (k, v))58 59    if mask is not None:60        if mask.dim() == 3:61            mask = mask.unsqueeze(1)62            mask = mask.expand(mask.shape[0], heads, -1, -1)63 64    out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=is_causal)65    return out66 67 68class L2Norm(nn.Module):69    def forward(self, x: torch.Tensor):70        return F.normalize(x, p=2, dim=-1)71 72 73class Attention(nn.Module):74    """75    Args:76        hidden_size (int): Hidden size of the input tensor.77        num_heads (int): Number of attention heads.78        num_kv_heads (int, optional): Number of key/value heads. Defaults to None.79        qk_norm (Literal["l2", "ln", "none"], optional): Type of normalization to apply to query/key. Defaults to "none".80        bias (bool, optional): Whether to use bias in linear layers. Defaults to False.81 82    """83 84    def __init__(85        self,86        hidden_size: int,87        num_heads: int,88        num_kv_heads: int | None = None,89        qk_norm: Literal["l2", "ln", "none"] = "none",90        bias: bool = False,91        zero_init_output: bool = False,92    ):93        super().__init__()94        num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads95        self.dim = hidden_size // num_heads96        self.num_heads, self.num_kv_heads = num_heads, num_kv_heads97 98        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias)99        self.k_proj = nn.Linear(hidden_size, self.dim * num_kv_heads, bias=bias)100        self.v_proj = nn.Linear(hidden_size, self.dim * num_kv_heads, bias=bias)101        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=bias)102 103        if qk_norm == "l2":104            self.q_norm = L2Norm()105            self.k_norm = L2Norm()106        elif qk_norm == "ln":107            self.q_norm = nn.LayerNorm(self.dim, elementwise_affine=False)108            self.k_norm = nn.LayerNorm(self.dim, elementwise_affine=False)109        else:110            self.q_norm = nn.Identity()111            self.k_norm = nn.Identity()112 113        if zero_init_output:114            nn.init.zeros_(self.out_proj.weight)115            if self.out_proj.bias is not None:116                nn.init.zeros_(self.out_proj.bias)117 118    def forward(119        self,120        x: torch.Tensor,121        context: torch.Tensor | None = None,122        mask: torch.Tensor | None = None,123        rotary_pos_emb: Tuple[torch.Tensor, torch.Tensor] | None = None,124        is_causal: bool = False,125    ) -> torch.Tensor:126        context = x if context is None else context127 128        q = self.q_proj(x)129        k, v = self.k_proj(context), self.v_proj(context)130 131        q = einops.rearrange(q, "b l (h d) -> b h l d", h=self.num_heads)132        k = einops.rearrange(k, "b s (h d) -> b h s d", h=self.num_kv_heads)133        v = einops.rearrange(v, "b s (h d) -> b h s d", h=self.num_kv_heads)134 135        q, k = self.q_norm(q), self.k_norm(k)136 137        if rotary_pos_emb is not None:138            q, k = map(lambda t: apply_rotary_pos_emb(t, *rotary_pos_emb), (q, k))139 140        out = attention_op(q, k, v, mask=mask, is_causal=is_causal)141        out = einops.rearrange(out, "b h l d -> b l (h d)")142        out = self.out_proj(out)143 144        return out145 146 147class PositionalEmbedding(nn.Module):148    def __init__(149        self,150        dim: int,151        encoding_type: Literal["sincos", "fourier"] = "sincos",152        scale: float = 2.0,153    ):154        super().__init__()155        self.dim = dim156        self.encoding_type = encoding_type157 158        if encoding_type == "fourier":159            self.register_buffer("freqs", torch.randn(dim // 2) * scale, persistent=True)160        elif encoding_type == "sincos":161            pass162        else:163            raise ValueError(f"encoding_type must be 'sincos' or 'fourier', but got {encoding_type}")164 165    def _create_sincos_emb(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:166        position = torch.arange(seq_len, device=device, dtype=torch.float32).unsqueeze(1)167        div_term = torch.exp(168            torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) * -(math.log(10000.0) / self.dim)169        )170 171        pos_emb = torch.zeros(seq_len, self.dim, device=device, dtype=dtype)172        pos_emb[:, 0::2] = torch.sin(position * div_term).to(dtype)173        pos_emb[:, 1::2] = torch.cos(position * div_term).to(dtype)174 175        return pos_emb176 177    def _create_fourier_emb(self, timestamps: torch.Tensor, device: torch.device, dtype: torch.dtype) -> torch.Tensor:178        # Ensure freqs is on the correct device179        freqs = self.freqs.to(device)180        pos_emb = torch.einsum("b t, d -> b t d", timestamps, 2 * np.pi * freqs).to(device, torch.float32)181        pos_emb = torch.cat([pos_emb.cos(), pos_emb.sin()], dim=-1).to(dtype)182        return pos_emb183 184    def forward(185        self, x: torch.Tensor, freq: Optional[Union[float, torch.Tensor]] = None, dtype: torch.dtype = torch.float32186    ) -> torch.Tensor:187        b, t = x.shape[0], x.shape[1]188        device = x.device189 190        if self.encoding_type == "sincos":191            pos_emb = self._create_sincos_emb(t, device, dtype)192            pos_emb = pos_emb.unsqueeze(0).expand(b, -1, -1)193            return pos_emb * 0.1194 195        elif self.encoding_type == "fourier":196            if freq is None:197                raise ValueError(198                    "freq must be provided when encoding_type is 'fourier'. Please provide the sequence frequency."199                )200            if isinstance(freq, float):201                freq = torch.tensor(freq, dtype=dtype, device=device)[None].expand(b)202            timestamps = torch.einsum("t, b -> b t", torch.arange(t, dtype=dtype, device=device), 1 / freq)203            pos_emb = self._create_fourier_emb(timestamps, device, dtype)204            return pos_emb * 0.1205        else:206            raise ValueError(f"Unknown encoding_type: {self.encoding_type}")207 208 209class SinusoidalPositionalEmbedding(PositionalEmbedding):210    def __init__(self, dim: int):211        super().__init__(dim=dim, encoding_type="sincos")212 213    def forward(self, x: torch.Tensor, pos: Optional[torch.Tensor] = None) -> torch.Tensor:214        return super().forward(x, freq=None)215 216 217class FeedForward(nn.Module):218    def __init__(self, hidden_size: int, intermediate_size: int, bias: bool = False):219        super().__init__()220        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=bias)221        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=bias)222        self.act_fn = nn.GELU()223 224    def forward(self, x: torch.Tensor) -> torch.Tensor:225        down_proj = self.down_proj(self.act_fn(self.up_proj(x)))226        return down_proj227 228 229class LayerScale(nn.Module):230    def __init__(self, dim, init_val=1e-2):231        super().__init__()232        self.scale = nn.Parameter(torch.full([dim], init_val))233 234    def forward(self, x):235        return x * self.scale236 237 238class PerceiverTransformerBlock(nn.Module):239    def __init__(240        self,241        dim: int,242        num_heads: int,243        mlp_ratio: int = 4,244        dropout: float = 0.0,245        qk_norm: str = "ln",246        layer_scale: bool = True,247        zero_init_output: bool = False,248        add_self_attn: bool = False,249        add_causal_mask: bool = False,250    ):251        super().__init__()252        self.add_self_attn = add_self_attn253        self.add_causal_mask = add_causal_mask254 255        self.norm1 = nn.LayerNorm(dim, eps=1e-2)256        self.cross_attn = Attention(257            hidden_size=dim, num_heads=num_heads, qk_norm=qk_norm, bias=False, zero_init_output=zero_init_output258        )259 260        if add_self_attn:261            self.norm_self_attn = nn.LayerNorm(dim, eps=1e-2)262            self.self_attn = Attention(263                hidden_size=dim, num_heads=num_heads, qk_norm=qk_norm, bias=False, zero_init_output=zero_init_output264            )265        else:266            self.self_attn = None267 268        self.norm2 = nn.LayerNorm(dim, eps=1e-2)269        self.mlp = FeedForward(hidden_size=dim, intermediate_size=int(mlp_ratio * dim), bias=True)270        self.dropout = nn.Dropout(dropout)271 272        self.attn_scale = LayerScale(dim) if layer_scale else nn.Identity()273        self.mlp_scale = LayerScale(dim) if layer_scale else nn.Identity()274 275        if zero_init_output:276            nn.init.zeros_(self.mlp.down_proj.weight)277            if self.mlp.down_proj.bias is not None:278                nn.init.zeros_(self.mlp.down_proj.bias)279 280    def forward(281        self,282        x: torch.Tensor,283        context: torch.Tensor,284        context_mask: Optional[torch.Tensor] = None,285        rotary_pos_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,286    ) -> torch.Tensor:287        residual = x288        x = self.norm1(x)289        x = self.cross_attn(x=x, context=context, mask=context_mask, rotary_pos_emb=rotary_pos_emb, is_causal=False)290        x = self.dropout(x)291        x = self.attn_scale(x)292        x = x + residual293 294        if self.add_self_attn:295            residual = x296            x = self.norm_self_attn(x)297            x = self.self_attn(298                x=x,299                context=None,300                mask=None,301                rotary_pos_emb=rotary_pos_emb,302                is_causal=self.add_causal_mask,303            )304            x = self.dropout(x)305            x = self.attn_scale(x)306            x = x + residual307 308        residual = x309        x = self.norm2(x)310        x = self.mlp(x)311        x = self.dropout(x)312        x = self.mlp_scale(x)313        x = x + residual314 315        return x316 317 318class EmbodimentEmbedding(nn.Module):319    def __init__(self, embodiment_config: dict, out_len: int, out_dim: int) -> None:320        super().__init__()321        self.out_len, self.out_dim = out_len, out_dim322 323        self.embodiment_config = embodiment_config324        self.num_embodiments = len(self.embodiment_config)325 326        self.embedding = nn.Embedding(self.num_embodiments, out_dim * out_len)327 328    @torch.no_grad()329    def expand_embodiment(self, embodiment_config: dict):330        for k in embodiment_config.keys():331            assert k not in self.embodiment_config.keys()332        self.embodiment_config.update(embodiment_config)333        self.num_embodiments = len(self.embodiment_config)334 335        extra_embodiments = len(embodiment_config)336 337        old_weights = torch.clone(self.embedding.weight)338        self.embedding = nn.Embedding(self.num_embodiments, self.out_dim * self.out_len)339        self.embedding.weight.data[:-extra_embodiments] = old_weights340        return self341 342    def keys(self) -> list[str]:343        return list(self.embodiment_config.keys())344 345    def ids_to_keys(self, ids: torch.Tensor) -> List[str]:346        return [self.keys()[i] for i in ids]347 348    def keys_to_ids(self, keys: List[str]) -> torch.Tensor:349        return torch.tensor([self.keys().index(k) for k in keys])350 351    def forward(self, x: torch.Tensor) -> torch.Tensor:352        return einops.rearrange(self.embedding(x), "b (l d) -> b l d", d=self.out_dim)353 354 355class PerceiverEncoder(nn.Module):356    def __init__(self, config: ActionCodecConfig):357        super().__init__()358        self.config = config359        self.embodiment_config = deepcopy(config.embodiment_config)360 361        out_len = int(config.n_tokens // config.n_quantizers)362        dim = config.encoder_dim363 364        _action_dim, _freq, _duration = list(), list(), list()365        for k, v in self.embodiment_config.items():366            _action_dim.append(v["action_dim"])367            _freq.append(v["freq"])368            _duration.append(v["duration"])369        self.register_buffer("_action_dim", torch.tensor(_action_dim), persistent=False)370        self.register_buffer("_freq", torch.tensor(_freq), persistent=False)371        self.register_buffer("_duration", torch.tensor(_duration), persistent=False)372 373        self.max_action_dim = max(v["action_dim"] for v in self.embodiment_config.values())374        self.input_proj = nn.Linear(self.max_action_dim, dim)375 376        self.cls_tokens = EmbodimentEmbedding(self.embodiment_config, out_len, dim)377 378        self.pos_emb_q = PositionalEmbedding(dim, encoding_type="sincos")379        self.pos_emb_kv = PositionalEmbedding(dim, encoding_type=config.encoder_pos_encoding_type)380 381        self.layers = nn.ModuleList(382            [383                PerceiverTransformerBlock(384                    dim=dim,385                    num_heads=config.encoder_n_heads,386                    add_self_attn=config.encoder_add_self_attn,387                    add_causal_mask=config.encoder_add_causal_mask,388                )389                for _ in range(config.encoder_n_layers)390            ]391        )392 393        self.output_proj = nn.Linear(dim, config.z_dim)394        self._init_weights()395 396    def _init_weights(self):397        nn.init.trunc_normal_(self.input_proj.weight, std=0.02)398        if self.input_proj.bias is not None:399            nn.init.zeros_(self.input_proj.bias)400        nn.init.trunc_normal_(self.output_proj.weight, std=0.02)401        if self.output_proj.bias is not None:402            nn.init.zeros_(self.output_proj.bias)403 404        nn.init.trunc_normal_(self.cls_tokens.embedding.weight, std=0.02)405 406    @torch.no_grad()407    def expand_embodiment(self, embodiment_config: dict):408        self.cls_tokens.expand_embodiment(embodiment_config)409        self.embodiment_config = self.cls_tokens.embodiment_config410        _action_dim, _freq, _duration = list(), list(), list()411        for k, v in self.embodiment_config.items():412            _action_dim.append(v["action_dim"])413            _freq.append(v["freq"])414            _duration.append(v["duration"])415        self._action_dim = torch.tensor(_action_dim)416        self._freq = torch.tensor(_freq)417        self._duration = torch.tensor(_duration)418 419        max_action_dim = max(v["action_dim"] for v in self.embodiment_config.values())420        if max_action_dim > self.max_action_dim:421            old_weights = torch.clone(self.input_proj.weight)422            old_bias = torch.clone(self.input_proj.bias)423            self.input_proj = nn.Linear(max_action_dim, self.config.encoder_dim)424            self.input_proj.weight.data[:, : self.max_action_dim] = old_weights425            self.input_proj.bias.data = old_bias426            self.max_action_dim = max_action_dim427 428        return self429 430    def forward(431        self,432        x: torch.Tensor,433        embodiment_ids: torch.Tensor | int,434        padding_mask: Optional[torch.Tensor] = None,435    ) -> torch.Tensor:436        """Encode action sequences into latent representations.437 438        Args:439            x (torch.Tensor): Action sequences to encode. Shape: (b, seq_len, max_action_dim).440                Assumes that the action dimension is zero-padded to the max action dimension.441                `seq_len` is supposed to be `int(duration * freq)` for each embodiment and padded to the max sequence length.442            embodiment_ids (torch.Tensor | int): Embodiment IDs. Shape: (b,).443                If int, the same embodiment ID is repeated for all sequences in the batch.444                It specifies the embodiment to encode.445            padding_mask (Optional[torch.Tensor], optional): Padding mask, where `False` values indicate padding. Shape: (b, seq_len). Defaults to None.446                It is used to mask the padding tokens on `seq_len` dimension.447 448        Returns:449            torch.Tensor: Encoded latent representations. Shape: (b, n_tokens_per_quantizer, z_dim).450        """451        b, seq_len, _ = x.shape452 453        x = self.input_proj(x)454 455        if isinstance(embodiment_ids, int):456            embodiment_ids = torch.tensor([embodiment_ids], dtype=torch.long, device=x.device).repeat(b)457 458        cls_tokens = self.cls_tokens(embodiment_ids)459 460        freqs = self._freq[embodiment_ids].to(x.device, x.dtype)461 462        pos_emb_q = self.pos_emb_q(cls_tokens)463        pos_emb_kv = self.pos_emb_kv(x, freqs)464 465        cls_tokens = cls_tokens + pos_emb_q466        x = x + pos_emb_kv467 468        if padding_mask is not None:469            padding_mask = padding_mask.unsqueeze(1).expand(-1, cls_tokens.shape[1], -1)470 471        for layer in self.layers:472            cls_tokens = layer(x=cls_tokens, context=x, context_mask=padding_mask)473 474        return self.output_proj(cls_tokens)475 476 477class PerceiverDecoder(nn.Module):478    def __init__(self, config: ActionCodecConfig):479        super().__init__()480        self.config = config481        self.embodiment_config = deepcopy(config.embodiment_config)482 483        dim = config.decoder_dim484 485        _action_dim, _freq, _duration = list(), list(), list()486        for k, v in self.embodiment_config.items():487            _action_dim.append(v["action_dim"])488            _freq.append(v["freq"])489            _duration.append(v["duration"])490        self.register_buffer("_action_dim", torch.tensor(_action_dim), persistent=False)491        self.register_buffer("_freq", torch.tensor(_freq), persistent=False)492        self.register_buffer("_duration", torch.tensor(_duration), persistent=False)493 494        self.max_action_dim = max(v["action_dim"] for v in self.embodiment_config.values())495        self.input_proj = nn.Linear(config.z_dim, dim)496 497        self.cls_tokens = EmbodimentEmbedding(self.embodiment_config, config.decoder_cls_size, dim)498 499        self.pos_emb_q = PositionalEmbedding(dim, encoding_type=config.decoder_pos_encoding_type)500        self.pos_emb_kv = PositionalEmbedding(dim, encoding_type="sincos")501 502        self.layers = nn.ModuleList(503            [504                PerceiverTransformerBlock(505                    dim=dim,506                    num_heads=config.decoder_n_heads,507                    add_self_attn=config.decoder_add_self_attn,508                    add_causal_mask=config.decoder_add_causal_mask,509                )510                for _ in range(config.decoder_n_layers)511            ]512        )513 514        self.output_proj = nn.Linear(dim, self.max_action_dim)515        self._init_weights()516 517    def _init_weights(self):518        nn.init.trunc_normal_(self.input_proj.weight, std=0.02)519        if self.input_proj.bias is not None:520            nn.init.zeros_(self.input_proj.bias)521        nn.init.trunc_normal_(self.output_proj.weight, std=0.02)522        if self.output_proj.bias is not None:523            nn.init.zeros_(self.output_proj.bias)524        nn.init.trunc_normal_(self.cls_tokens.embedding.weight, std=0.02)525 526    @torch.no_grad()527    def expand_embodiment(self, embodiment_config: dict):528        self.cls_tokens.expand_embodiment(embodiment_config)529        self.embodiment_config = self.cls_tokens.embodiment_config530 531        _action_dim, _freq, _duration = list(), list(), list()532        for k, v in self.embodiment_config.items():533            _action_dim.append(v["action_dim"])534            _freq.append(v["freq"])535            _duration.append(v["duration"])536        self._action_dim = torch.tensor(_action_dim)537        self._freq = torch.tensor(_freq)538        self._duration = torch.tensor(_duration)539 540        max_action_dim = max(v["action_dim"] for v in self.embodiment_config.values())541 542        if max_action_dim > self.max_action_dim:543            old_weights = torch.clone(self.output_proj.weight)544            old_bias = torch.clone(self.output_proj.bias)545 546            self.output_proj = nn.Linear(self.config.decoder_dim, max_action_dim)547 548            self.output_proj.weight.data[: self.max_action_dim, :] = old_weights549            self.output_proj.bias.data[: self.max_action_dim] = old_bias550 551            self.max_action_dim = max_action_dim552 553        return self554 555    def forward(556        self, x: torch.Tensor, embodiment_ids: torch.Tensor | int, durations: torch.Tensor | None = None557    ) -> torch.Tensor:558        """Decode latent representations into action sequences.559 560        Args:561            x (torch.Tensor): Latent representations to decode. Shape: (b, n_tokens_per_quantizer, z_dim).562            embodiment_ids (torch.Tensor | int): Embodiment IDs. Shape: (b,).563                If int, the same embodiment ID is repeated for all sequences in the batch.564                It specifies the embodiment to decode.565            durations (torch.Tensor | None, optional): Duration of each action sequence. Shape: (b,).566                If `None`, the duration is inferred from the default values in `embodiment_config`.567 568        Returns:569            torch.Tensor: Decoded action sequences. Shape: (b, seq_len, max_action_dim).570                Assumes that the action dimension is zero-padded to the max action dimension.571                `seq_len` is supposed to be `int(duration * freq)` for each embodiment and padded to the max sequence length.572        """573        b, seq_len, _ = x.shape574        x = self.input_proj(x)575 576        if isinstance(embodiment_ids, int):577            embodiment_ids = torch.tensor([embodiment_ids], dtype=torch.long, device=x.device).repeat(b)578 579        cls_tokens = self.cls_tokens(embodiment_ids)580 581        freqs = self._freq[embodiment_ids]582        if freqs.device != x.device:583            freqs = freqs.to(x.device)584 585        durations = self._duration[embodiment_ids] if durations is None else durations586        if isinstance(durations, torch.Tensor) and durations.device != x.device:587            durations = durations.to(x.device)588 589        action_horizons = (durations * freqs).long()590        max_horizon = action_horizons.max().item()591        padding_mask = torch.arange(max_horizon, device=x.device).expand(b, -1) < action_horizons.unsqueeze(1)592 593        if self.config.decoder_cls_size == 1:594            cls_tokens = cls_tokens.repeat(1, max_horizon, 1)595 596        pos_emb_q = self.pos_emb_q(cls_tokens, freqs)597        pos_emb_kv = self.pos_emb_kv(x)598 599        cls_tokens = cls_tokens + pos_emb_q600        x = x + pos_emb_kv601 602        for layer in self.layers:603            cls_tokens = layer(x=cls_tokens, context=x)604 605        output = self.output_proj(cls_tokens)606 607        return output, padding_mask608 609 610if __name__ == "__main__":611    # ------------------------------------------612    # 1. Initialization613    # ------------------------------------------614    print("=== Test 1: Initialization ===")615 616    # Define initial config with two smaller robots617    initial_embodiment_config = {618        "robot_small_7d": {"action_dim": 7, "freq": 20, "duration": 1, "description": "Original Robot"},619        "robot_tiny_3d": {"action_dim": 3, "freq": 10, "duration": 2, "description": "Tiny Robot"},620    }621 622    config = ActionCodecConfig(embodiment_config=initial_embodiment_config)623 624    # Set seed for reproducibility625    torch.manual_seed(42)626 627    encoder = PerceiverEncoder(config)628    decoder = PerceiverDecoder(config)629 630    encoder.eval()631    decoder.eval()632    print("✅ Models initialized successfully.")633 634    # ------------------------------------------635    # 2. Baseline Inference (Before Expansion)636    # ------------------------------------------637    print("\n=== Test 2: Baseline Inference (Before Expansion) ===")638 639    # Simulate Robot 1 (7-dim) data640    # Max action dim currently is 7.641    batch_size = 1642    seq_len = 20  # 20Hz * 1s643 644    # Input: (1, 20, 7)645    input_action_v0 = torch.randn(batch_size, seq_len, 7)646    emb_id_v0 = torch.tensor([0], dtype=torch.long)  # ID 0 -> robot_small_7d647 648    with torch.no_grad():649        z_ref = encoder(input_action_v0, emb_id_v0)650        rec_action_ref, _ = decoder(z_ref, emb_id_v0)651 652    print(f"Reference Latent Shape: {z_ref.shape}")653    print(f"Reference Recon Shape: {rec_action_ref.shape}")654 655    # ------------------------------------------656    # 3. Model Expansion (Add New Embodiment)657    # ------------------------------------------658    print("\n=== Test 3: Model Expansion ===")659 660    # Add a larger robot: 10-dim, high frequency661    new_embodiment_config = {662        "robot_large_10d": {"action_dim": 10, "freq": 30, "duration": 1, "description": "New Large Robot"}663    }664 665    print(f"Expanding from Max Dim {encoder.max_action_dim} to 10...")666    encoder.expand_embodiment(new_embodiment_config)667    decoder.expand_embodiment(new_embodiment_config)668 669    # Verify buffer updates670    assert encoder._action_dim[-1] == 10671    assert encoder.max_action_dim == 10672    assert decoder.max_action_dim == 10673    print(f"✅ Expansion successful. New Encoder Input Dim: {encoder.input_proj.weight.shape[1]}")674    print(f"✅ New Decoder Output Dim: {decoder.output_proj.weight.shape[0]}")675 676    # ------------------------------------------677    # 4. Encoder Invariance Check678    # ------------------------------------------679    print("\n=== Test 4: Encoder Invariance Check ===")680 681    # Pad old data (7 dims) to new max dim (10 dims) with ZEROS.682    input_action_padded = torch.zeros(batch_size, seq_len, 10)683    input_action_padded[:, :, :7] = input_action_v0684 685    with torch.no_grad():686        z_new = encoder(input_action_padded, emb_id_v0)687 688    # Compare latents689    diff_z = (z_ref - z_new).abs().max().item()690    print(f"Latent Difference (Max Abs): {diff_z:.8f}")691 692    if diff_z < 1e-6:693        print("✅ PASS: Encoder produces identical latents for old data.")694    else:695        print("❌ FAIL: Encoder outputs changed after expansion!")696 697    # ------------------------------------------698    # 5. Decoder Invariance Check699    # ------------------------------------------700    print("\n=== Test 5: Decoder Invariance Check ===")701 702    with torch.no_grad():703        # Feed old latent to expanded decoder704        rec_action_new_full, _ = decoder(z_ref, emb_id_v0)705 706    # Output shape should be (1, 20, 10)707    print(f"Expanded Decoder Output Shape: {rec_action_new_full.shape}")708 709    # Slice first 7 dims, should match reference710    rec_action_new_sliced = rec_action_new_full[:, :, :7]711 712    diff_rec = (rec_action_ref - rec_action_new_sliced).abs().max().item()713    print(f"Reconstruction Difference (Max Abs on valid dims): {diff_rec:.8f}")714 715    if diff_rec < 1e-6:716        print("✅ PASS: Decoder produces identical action values for valid dimensions.")717    else:718        print("❌ FAIL: Decoder outputs changed!")719 720    # Check phantom dimensions (7-9)721    # For old embodiment, these are driven by random weights and should be random722    new_dims_mean = rec_action_new_full[:, :, 7:].abs().mean().item()723    print(f"Values in new phantom dimensions (should be random garbage): {new_dims_mean:.4f}")724 725    # ------------------------------------------726    # 6. New Embodiment Inference727    # ------------------------------------------728    print("\n=== Test 6: New Embodiment Inference ===")729 730    # ID 2 -> robot_large_10d731    emb_id_new = torch.tensor([2], dtype=torch.long)732    seq_len_new = 30  # 30Hz * 1s733 734    input_action_new = torch.randn(1, seq_len_new, 10)735 736    with torch.no_grad():737        z_large = encoder(input_action_new, emb_id_new)738        rec_large, mask_large = decoder(z_large, emb_id_new)739 740    print(f"New Embodiment Output Shape: {rec_large.shape}")741 742    if rec_large.shape == (1, 30, 10):743        print("✅ PASS: New embodiment handled correctly with full dimensions.")744    else:745        print(f"❌ FAIL: Expected (1, 30, 10), got {rec_large.shape}")746 747    # ------------------------------------------748    # 7. Mixed Batch Processing (Masking)749    # ------------------------------------------750    print("\n=== Test 7: Mixed Batch Processing ===")751 752    # Batch size 2: [Robot 0 (20Hz, 7dim), Robot 2 (30Hz, 10dim)]753    mixed_emb_ids = torch.tensor([0, 2], dtype=torch.long)754 755    # Max seq len is 30. Max action dim is 10.756    batch_input = torch.zeros(2, 30, 10)757 758    # Fill data759    # Batch 0: Length 20, Dim 7 valid760    batch_input[0, :20, :7] = torch.randn(20, 7)761    # Batch 1: Length 30, Dim 10 valid762    batch_input[1, :30, :10] = torch.randn(30, 10)763 764    # Encoder Mask: True = Valid765    enc_padding_mask = torch.zeros(2, 30, dtype=torch.bool)766    enc_padding_mask[0, :20] = True767    enc_padding_mask[1, :30] = True768 769    print("Running mixed batch...")770    with torch.no_grad():771        z_mixed = encoder(batch_input, mixed_emb_ids, padding_mask=enc_padding_mask)772        rec_mixed, dec_padding_mask = decoder(z_mixed, mixed_emb_ids)773 774    print(f"Mixed Reconstruction Shape: {rec_mixed.shape}")  # Should be (2, 30, 10)775 776    # Verify Decoder Generated Mask777    valid_len_0 = dec_padding_mask[0].sum().item()778    valid_len_1 = dec_padding_mask[1].sum().item()779 780    print(f"Decoder Mask Valid Lengths: Batch 0={valid_len_0}, Batch 1={valid_len_1}")781 782    if valid_len_0 == 20 and valid_len_1 == 30:783        print("✅ PASS: Decoder correctly generated masks based on frequency and duration.")784    else:785        print("❌ FAIL: Decoder masks are incorrect.")786 787    print("\n✨ All Tests Completed ✨")788