STP / models_stp.py

Update the checkpoint filename in the header comment

986db04 verified 17 days ago

17.1 kB

# Copyright (c) 2024 Jiange Yang, Nanjing University / Microsoft Research
# All rights reserved.
#
# STP: Spatiotemporal Predictive Pre-training for Robotic Motor Control
# IJCV 2026: https://link.springer.com/article/10.1007/s11263-026-02948-3
#
# This file is a clean re-implementation of the pre-training model. It is
# parameter-name compatible with the released checkpoint
# `0.95_aug0.8.pth` (434 tensors, verified strict=True).
#
# Reference implementations:
#   MAE  : https://github.com/facebookresearch/mae
#   timm : https://github.com/rwightman/pytorch-image-models
# --------------------------------------------------------
from functools import partial

import torch
import torch.nn as nn

# `timm==0.3.2` (the version this model was trained with) imports
# `container_abcs` from `torch._six`, which PyTorch removed in 2.0. Installing
# the shim before importing timm keeps that pinned version working on modern
# PyTorch without editing site-packages.
import util.torch_six_compat  # noqa: F401  (must precede the timm import)

from timm.models.vision_transformer import Block, PatchEmbed

from util.pos_embed import get_2d_sincos_pos_embed


class CrossAttention(nn.Module):
    """Cross attention with a *frozen* key/value context.
    In STP's temporal decoder the current-frame feature `context` acts as
    key/value and is never updated, and the future-frame query cannot write
    back into the current-frame representation space. This asymmetry is what
    keeps the encoder's image representation free of future-frame leakage.
    """

    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5

        self.q = nn.Linear(dim, dim, bias=qkv_bias)
        self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)

        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x, context):
        B, N, C = x.shape
        M = context.shape[1]

        q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
        kv = self.kv(context).reshape(B, M, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        k, v = kv[0], kv[1]

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class DecoderBlock(nn.Module):
    """Temporal decoder block: self-attn -> cross-attn(context) -> FFN."""

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.self_attn = _SelfAttention(dim, num_heads, qkv_bias, attn_drop, drop)
        self.cross_attn = CrossAttention(dim, num_heads, qkv_bias, attn_drop, drop)
        self.query_norm = norm_layer(dim)
        self.context_norm = norm_layer(dim)
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = _Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

    def forward(self, x, context):
        x = x + self.self_attn(self.norm1(x))
        x = x + self.cross_attn(self.query_norm(x), self.context_norm(context))
        x = x + self.mlp(self.norm2(x))
        return x


class _SelfAttention(nn.Module):
    """Plain multi-head self attention. Named so the released checkpoint's
    `self_attn.qkv.*` / `self_attn.proj.*` keys load exactly."""

    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class _Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None,
                 act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class MaskedAutoencoderViT(nn.Module):
    """STP: one shared ViT encoder + decoupled spatial / temporal decoders.
    Pre-training samples a frame pair (current frame I_c, future frame I_f)
    from a clip and applies *asymmetric* masking:
        spatial  branch: I_c masked at mask_ratio_current (75%) -> spatial decoder
                         predicts the removed patches of I_c.
        temporal branch: I_f masked at mask_ratio_future  (95%) -> temporal decoder
                         predicts the removed patches of I_f, conditioned on Z_c.
    Only the encoder is kept for downstream policy learning; both decoders are
    discarded. The encoder is a plain image ViT with no cross-frame attention.
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3,
                 embed_dim=1024, depth=24, num_heads=16,
                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
                 mlp_ratio=4., norm_layer=nn.LayerNorm, norm_pix_loss=False,
                 mask_ratio_current=0.75, mask_ratio_future=0.95):
        super().__init__()

        # ------------------------------------------------------------------
        # encoder
        # ------------------------------------------------------------------
        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim),
                                      requires_grad=False)  # fixed sin-cos

        self.blocks = nn.ModuleList([
            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)
            for _ in range(depth)])
        self.norm = norm_layer(embed_dim)

        # ------------------------------------------------------------------
        # dual decoders
        # ------------------------------------------------------------------
        self.decoder_embed_1 = nn.Linear(embed_dim, decoder_embed_dim, bias=True)
        self.decoder_embed_2 = nn.Linear(embed_dim, decoder_embed_dim, bias=True)

        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))

        self.decoder_pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, decoder_embed_dim),
                                              requires_grad=False)  # fixed sin-cos

        # spatial decoder: standard transformer encoder blocks
        self.decoder_blocks_1 = nn.ModuleList([
            Block(decoder_embed_dim, decoder_num_heads, mlp_ratio,
                  qkv_bias=True, norm_layer=norm_layer)
            for _ in range(decoder_depth)])

        # temporal decoder: self-attention + cross-attention over the current frame
        self.decoder_blocks_2 = nn.ModuleList([
            DecoderBlock(decoder_embed_dim, decoder_num_heads, mlp_ratio,
                         qkv_bias=True, norm_layer=norm_layer)
            for _ in range(decoder_depth)])

        self.decoder_norm = norm_layer(decoder_embed_dim)
        self.decoder_pred_1 = nn.Linear(decoder_embed_dim, patch_size ** 2 * in_chans, bias=True)
        self.decoder_pred_2 = nn.Linear(decoder_embed_dim, patch_size ** 2 * in_chans, bias=True)

        self.norm_pix_loss = norm_pix_loss
        self.mask_ratio_current = mask_ratio_current
        self.mask_ratio_future = mask_ratio_future

        self.initialize_weights()

    # ----------------------------------------------------------------------
    # init
    # ----------------------------------------------------------------------
    def initialize_weights(self):
        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1],
                                            int(self.patch_embed.num_patches ** .5),
                                            cls_token=True)
        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))

        decoder_pos_embed = get_2d_sincos_pos_embed(self.decoder_pos_embed.shape[-1],
                                                    int(self.patch_embed.num_patches ** .5),
                                                    cls_token=True)
        self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))

        w = self.patch_embed.proj.weight.data
        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
        torch.nn.init.normal_(self.cls_token, std=.02)
        torch.nn.init.normal_(self.mask_token, std=.02)
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            torch.nn.init.xavier_uniform_(m.weight)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    # ----------------------------------------------------------------------
    # helpers
    # ----------------------------------------------------------------------
    def patchify(self, imgs):
        """imgs: (N, 3, H, W) -> (N, L, patch_size**2 * 3)"""
        p = self.patch_embed.patch_size[0]
        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0
        h = w = imgs.shape[2] // p
        x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p))
        x = torch.einsum('nchpwq->nhwpqc', x)
        return x.reshape(shape=(imgs.shape[0], h * w, p ** 2 * 3))

    def unpatchify(self, x):
        """x: (N, L, patch_size**2 * 3) -> (N, 3, H, W)"""
        p = self.patch_embed.patch_size[0]
        h = w = int(x.shape[1] ** .5)
        assert h * w == x.shape[1]
        x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))
        x = torch.einsum('nhwpqc->nchpwq', x)
        return x.reshape(shape=(x.shape[0], 3, h * p, h * p))

    def random_masking(self, x, mask_ratio):
        """Per-sample random masking. Returns kept tokens, binary mask (1=removed), restore idx."""
        N, L, D = x.shape
        len_keep = int(L * (1 - mask_ratio))

        noise = torch.rand(N, L, device=x.device)
        ids_shuffle = torch.argsort(noise, dim=1)
        ids_restore = torch.argsort(ids_shuffle, dim=1)

        ids_keep = ids_shuffle[:, :len_keep]
        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))

        mask = torch.ones([N, L], device=x.device)
        mask[:, :len_keep] = 0
        mask = torch.gather(mask, dim=1, index=ids_restore)
        return x_masked, mask, ids_restore

    # ----------------------------------------------------------------------
    # encoder / decoders
    # ----------------------------------------------------------------------
    def forward_encoder(self, x, mask_ratio):
        x = self.patch_embed(x)
        x = x + self.pos_embed[:, 1:, :]
        x, mask, ids_restore = self.random_masking(x, mask_ratio)

        cls_token = self.cls_token + self.pos_embed[:, :1, :]
        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)

        for blk in self.blocks:
            x = blk(x)
        x = self.norm(x)
        return x, mask, ids_restore

    def _assemble(self, x, ids_restore, decoder_embed):
        """Project, re-insert mask tokens in the original order, add decoder pos embed."""
        x = decoder_embed(x)
        mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1)
        x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1)          # drop cls
        x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]))
        x = torch.cat([x[:, :1, :], x_], dim=1)                    # restore cls
        return x + self.decoder_pos_embed

    def forward_decoder_spatial(self, x, ids_restore):
        x = self._assemble(x, ids_restore, self.decoder_embed_1)
        for blk in self.decoder_blocks_1:
            x = blk(x)
        x = self.decoder_norm(x)
        return self.decoder_pred_1(x)[:, 1:, :]

    def forward_decoder_temporal(self, x, context, ids_restore):
        x = self._assemble(x, ids_restore, self.decoder_embed_2)
        context = self.decoder_embed_2(context)                    # Z_c stays fixed
        for blk in self.decoder_blocks_2:
            x = blk(x, context)
        x = self.decoder_norm(x)
        return self.decoder_pred_2(x)[:, 1:, :]

    def forward_loss(self, imgs, pred, mask):
        target = self.patchify(imgs)
        if self.norm_pix_loss:
            mean = target.mean(dim=-1, keepdim=True)
            var = target.var(dim=-1, keepdim=True)
            target = (target - mean) / (var + 1.e-6) ** .5

        loss = (pred - target) ** 2
        loss = loss.mean(dim=-1)                                   # [N, L]
        return (loss * mask).sum() / mask.sum()                    # only removed patches

    # ----------------------------------------------------------------------
    # forward (pre-training)
    # ----------------------------------------------------------------------
    def forward(self, current_frame, future_frame):
        """One pre-training step.
        current_frame: (N, 3, H, W) the frame at time t
        future_frame : (N, 3, H, W) the frame at time t + interval
        returns (loss, pred_spatial, pred_temporal, mask_spatial)
        """
        latent_c, mask_c, ids_c = self.forward_encoder(current_frame, self.mask_ratio_current)
        latent_f, mask_f, ids_f = self.forward_encoder(future_frame, self.mask_ratio_future)

        pred_c = self.forward_decoder_spatial(latent_c, ids_c)
        pred_f = self.forward_decoder_temporal(latent_f, latent_c, ids_f)

        loss = self.forward_loss(current_frame, pred_c, mask_c) \
             + self.forward_loss(future_frame, pred_f, mask_f)
        return loss, pred_c, pred_f, mask_c

    # ----------------------------------------------------------------------
    # frozen-encoder feature extraction (this is what downstream policies use)
    # ----------------------------------------------------------------------
    @torch.no_grad()
    def forward_features(self, imgs):
        """Full (unmasked) image -> [CLS] representation, shape (N, embed_dim).
        This is the interface used by downstream policy learning: the encoder
        is frozen and the [CLS] token is taken as the global visual feature.
        """
        x = self.patch_embed(imgs) + self.pos_embed[:, 1:, :]

        cls_token = self.cls_token + self.pos_embed[:, :1, :]
        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)

        for blk in self.blocks:
            x = blk(x)
        x = self.norm(x)
        return x[:, 0]


def mae_vit_base_patch16_dec512d8b(**kwargs):
    return MaskedAutoencoderViT(
        patch_size=16, embed_dim=768, depth=12, num_heads=12,
        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)


def mae_vit_large_patch16_dec512d8b(**kwargs):
    return MaskedAutoencoderViT(
        patch_size=16, embed_dim=1024, depth=24, num_heads=16,
        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)


def mae_vit_huge_patch14_dec512d8b(**kwargs):
    return MaskedAutoencoderViT(
        patch_size=14, embed_dim=1280, depth=32, num_heads=16,
        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)


# recommended archs
mae_vit_base_patch16 = mae_vit_base_patch16_dec512d8b   # ViT-B, released checkpoint
mae_vit_large_patch16 = mae_vit_large_patch16_dec512d8b
mae_vit_huge_patch14 = mae_vit_huge_patch14_dec512d8b

	# Copyright (c) 2024 Jiange Yang, Nanjing University / Microsoft Research
	# All rights reserved.
	#
	# STP: Spatiotemporal Predictive Pre-training for Robotic Motor Control
	# IJCV 2026: https://link.springer.com/article/10.1007/s11263-026-02948-3
	#
	# This file is a clean re-implementation of the pre-training model. It is
	# parameter-name compatible with the released checkpoint
	# `0.95_aug0.8.pth` (434 tensors, verified strict=True).
	#
	# Reference implementations:
	# MAE : https://github.com/facebookresearch/mae
	# timm : https://github.com/rwightman/pytorch-image-models
	# --------------------------------------------------------
	from functools import partial

	import torch
	import torch.nn as nn

	# `timm==0.3.2` (the version this model was trained with) imports
	# `container_abcs` from `torch._six`, which PyTorch removed in 2.0. Installing
	# the shim before importing timm keeps that pinned version working on modern
	# PyTorch without editing site-packages.
	import util.torch_six_compat # noqa: F401 (must precede the timm import)

	from timm.models.vision_transformer import Block, PatchEmbed

	from util.pos_embed import get_2d_sincos_pos_embed


	class CrossAttention(nn.Module):
	"""Cross attention with a frozen key/value context.

	In STP's temporal decoder the current-frame feature `context` acts as
	key/value and is never updated, and the future-frame query cannot write
	back into the current-frame representation space. This asymmetry is what
	keeps the encoder's image representation free of future-frame leakage.
	"""

	def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
	super().__init__()
	self.num_heads = num_heads
	head_dim = dim // num_heads
	self.scale = head_dim ** -0.5

	self.q = nn.Linear(dim, dim, bias=qkv_bias)
	self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)

	self.attn_drop = nn.Dropout(attn_drop)
	self.proj = nn.Linear(dim, dim)
	self.proj_drop = nn.Dropout(proj_drop)

	def forward(self, x, context):
	B, N, C = x.shape
	M = context.shape[1]

	q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
	kv = self.kv(context).reshape(B, M, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
	k, v = kv[0], kv[1]

	attn = (q @ k.transpose(-2, -1)) * self.scale
	attn = attn.softmax(dim=-1)
	attn = self.attn_drop(attn)

	x = (attn @ v).transpose(1, 2).reshape(B, N, C)
	x = self.proj(x)
	x = self.proj_drop(x)
	return x


	class DecoderBlock(nn.Module):
	"""Temporal decoder block: self-attn -> cross-attn(context) -> FFN."""

	def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
	act_layer=nn.GELU, norm_layer=nn.LayerNorm):
	super().__init__()
	self.norm1 = norm_layer(dim)
	self.self_attn = _SelfAttention(dim, num_heads, qkv_bias, attn_drop, drop)
	self.cross_attn = CrossAttention(dim, num_heads, qkv_bias, attn_drop, drop)
	self.query_norm = norm_layer(dim)
	self.context_norm = norm_layer(dim)
	self.norm2 = norm_layer(dim)
	mlp_hidden_dim = int(dim * mlp_ratio)
	self.mlp = _Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

	def forward(self, x, context):
	x = x + self.self_attn(self.norm1(x))
	x = x + self.cross_attn(self.query_norm(x), self.context_norm(context))
	x = x + self.mlp(self.norm2(x))
	return x


	class _SelfAttention(nn.Module):
	"""Plain multi-head self attention. Named so the released checkpoint's
	`self_attn.qkv.` / `self_attn.proj.` keys load exactly."""

	def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
	super().__init__()
	self.num_heads = num_heads
	head_dim = dim // num_heads
	self.scale = head_dim ** -0.5

	self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
	self.attn_drop = nn.Dropout(attn_drop)
	self.proj = nn.Linear(dim, dim)
	self.proj_drop = nn.Dropout(proj_drop)

	def forward(self, x):
	B, N, C = x.shape
	qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
	q, k, v = qkv.unbind(0)

	attn = (q @ k.transpose(-2, -1)) * self.scale
	attn = attn.softmax(dim=-1)
	attn = self.attn_drop(attn)

	x = (attn @ v).transpose(1, 2).reshape(B, N, C)
	x = self.proj(x)
	x = self.proj_drop(x)
	return x


	class _Mlp(nn.Module):
	def __init__(self, in_features, hidden_features=None, out_features=None,
	act_layer=nn.GELU, drop=0.):
	super().__init__()
	out_features = out_features or in_features
	hidden_features = hidden_features or in_features
	self.fc1 = nn.Linear(in_features, hidden_features)
	self.act = act_layer()
	self.fc2 = nn.Linear(hidden_features, out_features)
	self.drop = nn.Dropout(drop)

	def forward(self, x):
	x = self.fc1(x)
	x = self.act(x)
	x = self.fc2(x)
	x = self.drop(x)
	return x


	class MaskedAutoencoderViT(nn.Module):
	"""STP: one shared ViT encoder + decoupled spatial / temporal decoders.

	Pre-training samples a frame pair (current frame I_c, future frame I_f)
	from a clip and applies asymmetric masking:

	spatial branch: I_c masked at mask_ratio_current (75%) -> spatial decoder
	predicts the removed patches of I_c.
	temporal branch: I_f masked at mask_ratio_future (95%) -> temporal decoder
	predicts the removed patches of I_f, conditioned on Z_c.

	Only the encoder is kept for downstream policy learning; both decoders are
	discarded. The encoder is a plain image ViT with no cross-frame attention.
	"""

	def __init__(self, img_size=224, patch_size=16, in_chans=3,
	embed_dim=1024, depth=24, num_heads=16,
	decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
	mlp_ratio=4., norm_layer=nn.LayerNorm, norm_pix_loss=False,
	mask_ratio_current=0.75, mask_ratio_future=0.95):
	super().__init__()

	# ------------------------------------------------------------------
	# encoder
	# ------------------------------------------------------------------
	self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
	num_patches = self.patch_embed.num_patches

	self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
	self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim),
	requires_grad=False) # fixed sin-cos

	self.blocks = nn.ModuleList([
	Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)
	for _ in range(depth)])
	self.norm = norm_layer(embed_dim)

	# ------------------------------------------------------------------
	# dual decoders
	# ------------------------------------------------------------------
	self.decoder_embed_1 = nn.Linear(embed_dim, decoder_embed_dim, bias=True)
	self.decoder_embed_2 = nn.Linear(embed_dim, decoder_embed_dim, bias=True)

	self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))

	self.decoder_pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, decoder_embed_dim),
	requires_grad=False) # fixed sin-cos

	# spatial decoder: standard transformer encoder blocks
	self.decoder_blocks_1 = nn.ModuleList([
	Block(decoder_embed_dim, decoder_num_heads, mlp_ratio,
	qkv_bias=True, norm_layer=norm_layer)
	for _ in range(decoder_depth)])

	# temporal decoder: self-attention + cross-attention over the current frame
	self.decoder_blocks_2 = nn.ModuleList([
	DecoderBlock(decoder_embed_dim, decoder_num_heads, mlp_ratio,
	qkv_bias=True, norm_layer=norm_layer)
	for _ in range(decoder_depth)])

	self.decoder_norm = norm_layer(decoder_embed_dim)
	self.decoder_pred_1 = nn.Linear(decoder_embed_dim, patch_size ** 2 * in_chans, bias=True)
	self.decoder_pred_2 = nn.Linear(decoder_embed_dim, patch_size ** 2 * in_chans, bias=True)

	self.norm_pix_loss = norm_pix_loss
	self.mask_ratio_current = mask_ratio_current
	self.mask_ratio_future = mask_ratio_future

	self.initialize_weights()

	# ----------------------------------------------------------------------
	# init
	# ----------------------------------------------------------------------
	def initialize_weights(self):
	pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1],
	int(self.patch_embed.num_patches ** .5),
	cls_token=True)
	self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))

	decoder_pos_embed = get_2d_sincos_pos_embed(self.decoder_pos_embed.shape[-1],
	int(self.patch_embed.num_patches ** .5),
	cls_token=True)
	self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))

	w = self.patch_embed.proj.weight.data
	torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
	torch.nn.init.normal_(self.cls_token, std=.02)
	torch.nn.init.normal_(self.mask_token, std=.02)
	self.apply(self._init_weights)

	def _init_weights(self, m):
	if isinstance(m, nn.Linear):
	torch.nn.init.xavier_uniform_(m.weight)
	if m.bias is not None:
	nn.init.constant_(m.bias, 0)
	elif isinstance(m, nn.LayerNorm):
	nn.init.constant_(m.bias, 0)
	nn.init.constant_(m.weight, 1.0)

	# ----------------------------------------------------------------------
	# helpers
	# ----------------------------------------------------------------------
	def patchify(self, imgs):
	"""imgs: (N, 3, H, W) -> (N, L, patch_size*2 3)"""
	p = self.patch_embed.patch_size[0]
	assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0
	h = w = imgs.shape[2] // p
	x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p))
	x = torch.einsum('nchpwq->nhwpqc', x)
	return x.reshape(shape=(imgs.shape[0], h * w, p ** 2 * 3))

	def unpatchify(self, x):
	"""x: (N, L, patch_size*2 3) -> (N, 3, H, W)"""
	p = self.patch_embed.patch_size[0]
	h = w = int(x.shape[1] ** .5)
	assert h * w == x.shape[1]
	x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))
	x = torch.einsum('nhwpqc->nchpwq', x)
	return x.reshape(shape=(x.shape[0], 3, h * p, h * p))

	def random_masking(self, x, mask_ratio):
	"""Per-sample random masking. Returns kept tokens, binary mask (1=removed), restore idx."""
	N, L, D = x.shape
	len_keep = int(L * (1 - mask_ratio))

	noise = torch.rand(N, L, device=x.device)
	ids_shuffle = torch.argsort(noise, dim=1)
	ids_restore = torch.argsort(ids_shuffle, dim=1)

	ids_keep = ids_shuffle[:, :len_keep]
	x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))

	mask = torch.ones([N, L], device=x.device)
	mask[:, :len_keep] = 0
	mask = torch.gather(mask, dim=1, index=ids_restore)
	return x_masked, mask, ids_restore

	# ----------------------------------------------------------------------
	# encoder / decoders
	# ----------------------------------------------------------------------
	def forward_encoder(self, x, mask_ratio):
	x = self.patch_embed(x)
	x = x + self.pos_embed[:, 1:, :]
	x, mask, ids_restore = self.random_masking(x, mask_ratio)

	cls_token = self.cls_token + self.pos_embed[:, :1, :]
	cls_tokens = cls_token.expand(x.shape[0], -1, -1)
	x = torch.cat((cls_tokens, x), dim=1)

	for blk in self.blocks:
	x = blk(x)
	x = self.norm(x)
	return x, mask, ids_restore

	def _assemble(self, x, ids_restore, decoder_embed):
	"""Project, re-insert mask tokens in the original order, add decoder pos embed."""
	x = decoder_embed(x)
	mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1)
	x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1) # drop cls
	x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]))
	x = torch.cat([x[:, :1, :], x_], dim=1) # restore cls
	return x + self.decoder_pos_embed

	def forward_decoder_spatial(self, x, ids_restore):
	x = self._assemble(x, ids_restore, self.decoder_embed_1)
	for blk in self.decoder_blocks_1:
	x = blk(x)
	x = self.decoder_norm(x)
	return self.decoder_pred_1(x)[:, 1:, :]

	def forward_decoder_temporal(self, x, context, ids_restore):
	x = self._assemble(x, ids_restore, self.decoder_embed_2)
	context = self.decoder_embed_2(context) # Z_c stays fixed
	for blk in self.decoder_blocks_2:
	x = blk(x, context)
	x = self.decoder_norm(x)
	return self.decoder_pred_2(x)[:, 1:, :]

	def forward_loss(self, imgs, pred, mask):
	target = self.patchify(imgs)
	if self.norm_pix_loss:
	mean = target.mean(dim=-1, keepdim=True)
	var = target.var(dim=-1, keepdim=True)
	target = (target - mean) / (var + 1.e-6) ** .5

	loss = (pred - target) ** 2
	loss = loss.mean(dim=-1) # [N, L]
	return (loss * mask).sum() / mask.sum() # only removed patches

	# ----------------------------------------------------------------------
	# forward (pre-training)
	# ----------------------------------------------------------------------
	def forward(self, current_frame, future_frame):
	"""One pre-training step.

	current_frame: (N, 3, H, W) the frame at time t
	future_frame : (N, 3, H, W) the frame at time t + interval

	returns (loss, pred_spatial, pred_temporal, mask_spatial)
	"""
	latent_c, mask_c, ids_c = self.forward_encoder(current_frame, self.mask_ratio_current)
	latent_f, mask_f, ids_f = self.forward_encoder(future_frame, self.mask_ratio_future)

	pred_c = self.forward_decoder_spatial(latent_c, ids_c)
	pred_f = self.forward_decoder_temporal(latent_f, latent_c, ids_f)

	loss = self.forward_loss(current_frame, pred_c, mask_c) \
	+ self.forward_loss(future_frame, pred_f, mask_f)
	return loss, pred_c, pred_f, mask_c

	# ----------------------------------------------------------------------
	# frozen-encoder feature extraction (this is what downstream policies use)
	# ----------------------------------------------------------------------
	@torch.no_grad()
	def forward_features(self, imgs):
	"""Full (unmasked) image -> [CLS] representation, shape (N, embed_dim).

	This is the interface used by downstream policy learning: the encoder
	is frozen and the [CLS] token is taken as the global visual feature.
	"""
	x = self.patch_embed(imgs) + self.pos_embed[:, 1:, :]

	cls_token = self.cls_token + self.pos_embed[:, :1, :]
	cls_tokens = cls_token.expand(x.shape[0], -1, -1)
	x = torch.cat((cls_tokens, x), dim=1)

	for blk in self.blocks:
	x = blk(x)
	x = self.norm(x)
	return x[:, 0]


	def mae_vit_base_patch16_dec512d8b(**kwargs):
	return MaskedAutoencoderViT(
	patch_size=16, embed_dim=768, depth=12, num_heads=12,
	decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
	mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)


	def mae_vit_large_patch16_dec512d8b(**kwargs):
	return MaskedAutoencoderViT(
	patch_size=16, embed_dim=1024, depth=24, num_heads=16,
	decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
	mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)


	def mae_vit_huge_patch14_dec512d8b(**kwargs):
	return MaskedAutoencoderViT(
	patch_size=14, embed_dim=1280, depth=32, num_heads=16,
	decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
	mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)


	# recommended archs
	mae_vit_base_patch16 = mae_vit_base_patch16_dec512d8b # ViT-B, released checkpoint
	mae_vit_large_patch16 = mae_vit_large_patch16_dec512d8b
	mae_vit_huge_patch14 = mae_vit_huge_patch14_dec512d8b