GR00T-WholeBodyControl/gear_sonic/trl/modules/universal_token_modules.py

from __future__ import annotations

from loguru import logger
import numpy as np
import torch
from torch import nn

from gear_sonic.trl.utils import common


def set_fuzzy_config_params(config, candidate_keys, value):
    """Set matching config keys to a given value.

    Scans ``config`` for any key that appears in ``candidate_keys`` and
    overwrites its value.  Used to wire dynamically-computed dimensions
    (e.g. ``input_dim``, ``output_dim``) into encoder / decoder param
    dicts before instantiation.

    Args:
        config: Mutable mapping of parameter names to values.
        candidate_keys: Iterable of key names to look for in ``config``.
        value: Replacement value to assign to each matched key.

    Returns:
        The same ``config`` mapping with matching keys updated in-place.
    """
    for key in candidate_keys:
        if key in config:
            config[key] = value
    return config


class UniversalTokenModule(nn.Module):
    """SONIC-style action transform module (ATM) with FSQ token bottleneck.

    Implements the encoder → FSQ quantizer → decoder pipeline described in the
    SONIC paper.  Multiple named encoders (e.g. ``g1``, ``smpl``, ``teleop``)
    convert different flavors of tokenizer observations into a shared latent
    space.  An optional Finite Scalar Quantizer (FSQ) discretises the latent,
    producing a compact token representation.  One or more decoders then map
    the tokens (plus proprioception) back to joint-space actions.

    Token flow::

        tokenizer_obs ──► encoder(s) ──► [+ additive encoders]
                                      ──► [latent_residual (pre_quantization)]
                                      ──► FSQ quantizer
                                      ──► [+ latent_residual (post_quantization)]
                                      ──► decoder(s)
                                      ──► action_mean

    Latent residual modes allow an external HOI policy to inject corrections
    into the latent space without retraining the base ATM:

    * ``"post_quantization"`` - residual added to FSQ-quantized tokens (default).
    * ``"pre_quantization"`` - residual added *before* FSQ; the sum gets quantized.
    * ``"pre_quantization_replace"`` - latent is replaced entirely by the residual.

    Attributes:
        encoders: ``nn.ModuleDict`` of named encoder networks.
        quantizer: FSQ quantizer (``None`` when quantization is disabled).
        decoders: ``nn.ModuleDict`` of named decoder networks.
        aux_loss_func: ``nn.ModuleDict`` of auxiliary loss callables.
        token_dim: Dimensionality of a single token (equals ``num_fsq_levels``).
        max_num_tokens: Number of tokens produced per timestep.
        token_total_dim: ``token_dim * max_num_tokens`` - flat token size.
    """

    def __init__(
        self,
        env_config,
        algo_config,
        obs_dim_dict=None,
        module_dim_dict=None,
        proprioception_features=[],
        num_fsq_levels=5,
        fsq_level_list=16,
        max_num_tokens=None,
        down_t=2,
        num_future_frames=1,
        quantizer=None,
        encoders=None,
        decoders=None,
        aux_loss_func={},
        aux_loss_coef={},
        encoder_sample_probs=None,
        reencode_smpl_g1_recon=False,
        meta_action_dim=None,  # For hierarchical policies with split actions
        body_action_dim=None,  # For separate body decoder (e.g., 29 body joints)
        hand_action_dim=None,  # For separate hand decoder (e.g., 14 hand joints)
        freeze_encoders=False,
        freeze_decoders=False,
        freeze_quantizer=False,
        stiff_compliance_threshold=0.01,  # Threshold for stiff mode filtering
        optimize_encoders_ratio_for_CHIP=False,  # CHIP compliance training optimization
        active_encoders=None,  # Optional list of encoder names to activate (None = all)
        active_decoders=None,  # Optional list of decoder names to activate (None = all)
        **kwargs,  # noqa: ARG002
    ):
        """Initialise encoders, FSQ quantizer, decoders, and auxiliary losses.

        Args:
            env_config: Environment configuration object exposing
                ``obs.group_obs_dims``, ``obs.group_obs_names``, and
                ``robot.actions_dim``.
            algo_config: Algorithm configuration (stored but not directly
                accessed during construction).
            obs_dim_dict: Mapping from observation name to dimension.  Falls
                back to ``env_config.robot.algo_obs_dim_dict`` when ``None``.
            module_dim_dict: Optional extra dimension overrides for named
                intermediate features.
            proprioception_features: List of observation keys whose concatenation
                forms the proprioception input fed to every decoder.
            num_fsq_levels: Number of FSQ levels (equals ``token_dim``).
            fsq_level_list: Per-level codebook size.  An ``int`` is broadcast
                to ``[fsq_level_list] * num_fsq_levels``.
            max_num_tokens: Explicit token count per timestep.  When ``None``
                this is derived as ``max(1, num_future_frames // 2**down_t)``.
            down_t: Temporal downsampling factor for token count derivation.
            num_future_frames: Number of future motion frames in the tokenizer
                input window.
            quantizer: Hydra-instantiable config for the FSQ quantizer.
                ``None`` disables quantization (identity passthrough).
            encoders: Dict of encoder configs (Hydra DictConfig).  Each entry
                specifies ``inputs``, ``outputs``, ``params``, and optional
                ``additive_to`` / ``sub_encoders`` / ``mask`` / ``freeze``.
            decoders: Dict of decoder configs (Hydra DictConfig).  Each entry
                specifies ``inputs``, ``outputs``, ``conds``, ``params``, and
                ``has_temporal_dim``.
            aux_loss_func: Dict of auxiliary loss configs to instantiate as
                ``nn.Module`` callables (stored in ``self.aux_loss_func``).
            aux_loss_coef: Dict mapping loss name → scalar coefficient passed
                back to the trainer.
            encoder_sample_probs: Ordered dict mapping encoder name → sampling
                probability; defines the column order of ``encoder_index`` in
                tokenizer observations.
            reencode_smpl_g1_recon: When ``True``, re-encode the G1 kinematic
                decoder output to compute cycle-consistency auxiliary losses.
            meta_action_dim: Action dimensionality for hierarchical policies
                where the top-level action differs from the full joint count.
            body_action_dim: Dimensionality of body joints for split
                body / hand decoding (default 29).
            hand_action_dim: Dimensionality of hand joints for split
                body / hand decoding (default 14).
            freeze_encoders: Freeze all encoder parameters after init.
            freeze_decoders: Freeze all decoder parameters after init.
            freeze_quantizer: Freeze quantizer parameters after init.
            stiff_compliance_threshold: Absolute compliance value below which
                an environment is treated as "stiff" for auxiliary loss gating.
            optimize_encoders_ratio_for_CHIP: When ``True``, enforce one-hot
                encoder selection (CHIP training mode) instead of the default
                multi-hot selection where SMPL-native envs activate both SMPL
                and G1 encoders.
            active_encoders: Optional list of encoder names to instantiate and
                run.  All encoders are active when ``None``.
            active_decoders: Optional list of decoder names to run during
                forward.  All decoders are run when ``None``.
            **kwargs: Absorbed for forward-compatibility.
        """
        super().__init__()

        # Store freeze flags
        self.freeze_encoders = freeze_encoders
        self.freeze_decoders = freeze_decoders
        self.stiff_compliance_threshold = stiff_compliance_threshold
        self.freeze_quantizer = freeze_quantizer
        self.optimize_encoders_ratio_for_CHIP = optimize_encoders_ratio_for_CHIP

        if self.optimize_encoders_ratio_for_CHIP:
            logger.info(
                "[CHIP Mode] optimize_encoders_ratio_for_CHIP=True: "
                "Native encoder selection is now one-hot (no G1 auto-activation for SMPL). "
                "G1 latents for SMPL-native envs computed only in aux losses when compliance=0."
            )

        # Cache for last encoded tokens (for external access, e.g., by callbacks)
        # These are populated during forward() and can be read afterwards
        self._last_encoded_tokens = None  # dict[encoder_name -> Tensor]
        self._last_encoded_latents = None  # dict[encoder_name -> Tensor] (pre-quantization)

        self.env_config = env_config
        self.algo_config = algo_config
        self.encoders_cfg = encoders
        self.decoders_cfg = decoders
        self.encoder_sample_probs = encoder_sample_probs
        self.reencode_smpl_g1_recon = reencode_smpl_g1_recon
        self.tokenizer_obs_dims = self.env_config.obs.group_obs_dims["tokenizer"]
        self.tokenizer_obs_names = self.env_config.obs.group_obs_names["tokenizer"]
        # The tokenizer observation schema is fixed after module construction.
        tokenizer_obs_specs = []
        tokenizer_obs_offset = 0
        for name in self.tokenizer_obs_names:
            dims = tuple(self.tokenizer_obs_dims[name])
            flat_dim = int(np.prod(dims))
            tokenizer_obs_specs.append(
                (name, tokenizer_obs_offset, tokenizer_obs_offset + flat_dim, dims)
            )
            tokenizer_obs_offset += flat_dim
        self.tokenizer_obs_specs = tuple(tokenizer_obs_specs)
        self.tokenizer_obs_total_dim = tokenizer_obs_offset
        self.actions_dim = self.env_config.robot.actions_dim
        self.meta_action_dim = (
            meta_action_dim  # Can be different from actions_dim for hierarchical policies
        )

        if obs_dim_dict is None:
            obs_dim_dict = getattr(env_config.robot, "algo_obs_dim_dict", {})

        if module_dim_dict is None:
            module_dim_dict = {}

        self.obs_dim_dict = obs_dim_dict
        self.module_dim_dict = module_dim_dict
        self.proprioception_features = proprioception_features
        self.num_future_frames = num_future_frames

        # Initialize auxiliary loss functions (nn.ModuleDict so sub-module buffers
        # are part of the model state_dict and visible to DDP)
        self.aux_loss_func = nn.ModuleDict()
        for name, loss_func in aux_loss_func.items():
            self.aux_loss_func[name] = common.custom_instantiate(loss_func, _resolve=False)
        self.aux_loss_coef = aux_loss_coef

        module_input_candidate_params = ["input_dim"]

        module_output_candidate_params = ["output_dim"]

        """
        Initialize quantizer
        """
        if isinstance(fsq_level_list, int):
            fsq_level_list = [fsq_level_list] * num_fsq_levels
        if quantizer is not None:
            self.quantizer = common.custom_instantiate(
                quantizer, levels=fsq_level_list, _resolve=False
            )
        else:
            self.quantizer = None
        self.num_fsq_levels = num_fsq_levels
        self.fsq_level_list = fsq_level_list

        # quantizer levels determine the embedding dimension
        self.token_dim = self.num_fsq_levels
        self.down_t = down_t  # default
        if max_num_tokens is not None:
            self.max_num_tokens = max_num_tokens
        else:
            self.max_num_tokens = max(1, self.num_future_frames // (2**self.down_t))
        self.token_total_dim = self.token_dim * self.max_num_tokens
        logger.info(
            f"Motion Encoder and Quantizer initialized with embedding dim: {self.token_total_dim} (num_tokens={self.max_num_tokens}, token_dim={self.token_dim})"  # noqa: E501
        )

        """
        Initialize motion encoders dynamically
        """
        self.encoders = nn.ModuleDict()
        self.sub_encoders = {}
        self.encoder_input_features = {}
        self.encoder_mask_features = {}
        self.encoders_to_iterate = []
        self.additive_encoders = {}  # Map: base_encoder_name -> list of additive encoder names

        # First pass: identify additive encoders
        for encoder_name, encoder_config in self.encoders_cfg.items():
            additive_to = encoder_config.get("additive_to", None)
            if additive_to is not None:
                if additive_to not in self.additive_encoders:
                    self.additive_encoders[additive_to] = []
                self.additive_encoders[additive_to].append(encoder_name)

        # New config structure - initialize all encoders from config
        for encoder_name, encoder_config in self.encoders_cfg.items():
            # Skip non-active encoders early (avoids KeyError on missing obs dims)
            if active_encoders is not None and encoder_name not in active_encoders:
                continue

            sub_encoder_only = encoder_config.get("sub_encoder_only", False)
            is_additive = encoder_config.get("additive_to", None) is not None
            # Don't iterate additive encoders independently - they're called within their base encoder
            if not sub_encoder_only and not is_additive:
                self.encoders_to_iterate.append(encoder_name)
            self.sub_encoders[encoder_name] = encoder_config.get("sub_encoders", [])

            input_features = encoder_config.get("inputs", [])
            output_features = encoder_config.get("outputs", [])
            input_feature_dim = sum([self.tokenizer_obs_dims[key][-1] for key in input_features])
            if len(output_features) > 0:
                output_feature_dim = sum(
                    [self.tokenizer_obs_dims[key][-1] for key in output_features]
                )
            else:
                output_feature_dim = self.token_dim
            self.encoder_input_features[encoder_name] = input_features
            self.encoder_mask_features[encoder_name] = list(encoder_config.get("mask", []))

            if len(self.sub_encoders[encoder_name]) > 0:
                continue

            encoder_params = encoder_config["params"].copy()
            # Set dynamic parameters
            set_fuzzy_config_params(
                encoder_params, module_input_candidate_params, input_feature_dim
            )
            set_fuzzy_config_params(
                encoder_params, module_output_candidate_params, output_feature_dim
            )

            # Instantiate encoder
            encoder = common.custom_instantiate(encoder_params, _resolve=False)
            self.encoders[encoder_name] = encoder

            # Per-encoder freeze support
            if encoder_config.get("freeze", False):
                for param in encoder.parameters():
                    param.requires_grad = False
                logger.info(f"Froze encoder: {encoder_name}")

            logger.info(f"Initialized {encoder_name} encoder with input features: {input_features}")

        # Log additive encoder relationships
        if self.additive_encoders:
            for base_encoder, additive_list in self.additive_encoders.items():
                logger.info(
                    f"Additive encoders for '{base_encoder}': {additive_list} (outputs will be summed)"
                )

        """
        Initialize motion decoders dynamically
        """
        self.decoders = nn.ModuleDict()
        self.decoder_input_features = {}
        self.decoder_output_features = {}
        self.decoder_output_feature_dims = {}
        self.decoder_cond_features = {}  # For root-disentangled decoders (external_cond)
        self.decoder_mask_features = {}

        # Support custom meta_action_dim for hierarchical policies
        meta_action_dim = getattr(self, "meta_action_dim", None) or self.actions_dim

        # Compute proprioception dim from actual features (not hardcoded actor_obs)
        proprioception_dim = sum(
            self.obs_dim_dict[key]
            for key in self.proprioception_features
            if key in self.obs_dim_dict
        ) or self.obs_dim_dict.get("actor_obs", 0)

        decoder_feature_dims_map = {
            "proprioception": proprioception_dim,
            "token": self.token_dim,
            "token_flattened": self.token_dim * self.max_num_tokens,
            "action": self.actions_dim,
            "meta_action": meta_action_dim,  # For hierarchical policies
            "hand_action": hand_action_dim or 14,  # 14 hand joints default
            "body_action": body_action_dim or 29,  # 29 body joints default
        }
        for k, v in self.tokenizer_obs_dims.items():
            if k not in decoder_feature_dims_map:
                decoder_feature_dims_map[k] = v[-1]
        self.decoder_feature_dims_map = decoder_feature_dims_map

        for decoder_name, decoder_config in self.decoders_cfg.items():
            # Skip non-active decoders early (avoids KeyError on missing obs dims)
            if active_decoders is not None and decoder_name not in active_decoders:
                continue

            decoder_params = decoder_config["params"].copy()
            input_features = decoder_config.get("inputs", [])
            output_features = decoder_config.get("outputs", [])
            cond_features = decoder_config.get("conds", [])
            has_temporal_dim = decoder_config.has_temporal_dim
            input_feature_dim = 0
            for key in input_features:
                input_feature_dim += self.decoder_feature_dims_map[key]
            output_feature_dim = 0
            self.decoder_output_feature_dims[decoder_name] = {}
            for key in output_features:
                if key in decoder_feature_dims_map:
                    feature_dim = decoder_feature_dims_map[key]
                else:
                    assert has_temporal_dim
                    feature_dim = self.tokenizer_obs_dims[key][-1]
                self.decoder_output_feature_dims[decoder_name][key] = feature_dim
                output_feature_dim += feature_dim
            set_fuzzy_config_params(
                decoder_params, module_input_candidate_params, input_feature_dim
            )
            set_fuzzy_config_params(
                decoder_params, module_output_candidate_params, output_feature_dim
            )

            # For decoders with conds: compute external_cond_dim from cond feature dimensions.
            if cond_features:
                external_cond_dim = sum(self.tokenizer_obs_dims[key][-1] for key in cond_features)
                decoder_params["external_cond_dim"] = external_cond_dim
                logger.info(
                    f"Decoder '{decoder_name}' has conds={cond_features} "
                    f"with external_cond_dim={external_cond_dim}"
                )
            self.decoder_cond_features[decoder_name] = cond_features
            self.decoder_mask_features[decoder_name] = list(decoder_config.get("mask", []))

            # Instantiate decoder
            decoder = common.custom_instantiate(decoder_params, _resolve=False)
            self.decoders[decoder_name] = decoder
            self.decoder_input_features[decoder_name] = input_features
            self.decoder_output_features[decoder_name] = output_features

            logger.info(
                f"Initialized {decoder_name} decoder with input features: {input_features} and output features: {output_features}"  # noqa: E501
            )

        # Filter active encoders/decoders if specified (for kinematic-only training)
        if active_encoders is not None:
            self.encoders_to_iterate = [e for e in self.encoders_to_iterate if e in active_encoders]
            logger.info(f"Active encoders filtered to: {self.encoders_to_iterate}")
        if active_decoders is not None:
            # Keep all decoders in ModuleDict (for checkpoint compat), but only iterate active ones
            self._active_decoders = set(active_decoders)
            logger.info(f"Active decoders filtered to: {list(self._active_decoders)}")
        else:
            self._active_decoders = None

        # Apply freeze logic
        if self.freeze_encoders:
            for encoder in self.encoders.values():
                for param in encoder.parameters():
                    param.requires_grad = False
            logger.info(f"Froze encoders: {list(self.encoders.keys())}")

        if self.freeze_decoders:
            for decoder in self.decoders.values():
                for param in decoder.parameters():
                    param.requires_grad = False
            logger.info(f"Froze decoders: {list(self.decoders.keys())}")

        if self.freeze_quantizer and self.quantizer is not None:
            for param in self.quantizer.parameters():
                param.requires_grad = False
            logger.info("Froze quantizer")

        # Variable frame support: enabled when any encoder or decoder has mask observations
        self.variable_frames_enabled = any(
            feats
            for feats in (
                *self.encoder_mask_features.values(),
                *self.decoder_mask_features.values(),
            )
        )
        if self.variable_frames_enabled:
            # Cache frame indices for mask creation (avoids allocation every forward)
            self.register_buffer(
                "_frame_indices",
                torch.arange(self.num_future_frames).unsqueeze(0),
                persistent=False,
            )

    def _create_frame_and_token_masks(self, tokenizer_obs):
        """Create frame and token masks from num_frames observation.

        Returns (frame_mask, token_mask) or (None, None) if disabled.
            frame_mask: [B*seq, max_frames] bool, True=valid
            token_mask: [B*seq, max_tokens] bool, True=valid
        """
        if "command_num_frames" not in tokenizer_obs:
            return None, None

        num_frames = tokenizer_obs["command_num_frames"]  # [B, seq, 1]
        num_frames_flat = num_frames.reshape(-1, 1)  # [B*seq, 1]
        frame_mask = self._frame_indices < num_frames_flat  # [B*seq, max_frames]

        frames_per_token = 2**self.down_t
        token_frame_mask = frame_mask.reshape(-1, self.max_num_tokens, frames_per_token)
        token_mask = token_frame_mask.all(dim=-1)  # [B*seq, max_tokens]

        return frame_mask, token_mask

    def parse_tokenizer_obs(self, input_data):
        """Split the flat tokenizer observation tensor into a named dict.

        The tokenizer observation is stored as a single concatenated vector in
        ``input_data["tokenizer"]``.  This method reshapes each slice back to
        its original multi-dimensional form using the dimension metadata from
        ``env_config``.

        Args:
            input_data: Dict containing ``"tokenizer"`` key with shape
                ``(..., total_tokenizer_dim)``.

        Returns:
            Dict mapping observation name to tensor with shape
            ``(*batch_dims, *obs_dims)`` for each registered tokenizer
            observation.
        """
        tokenizer_obs = input_data["tokenizer"]
        tokenizer_obs_dict = {}
        for name, start, end, dims in self.tokenizer_obs_specs:
            tokenizer_obs_dict[name] = tokenizer_obs[..., start:end].reshape(
                tokenizer_obs.shape[:-1] + dims
            )
        assert self.tokenizer_obs_total_dim == tokenizer_obs.shape[-1], (
            f"{self.tokenizer_obs_total_dim=}, {tokenizer_obs.shape[-1]=}"
        )
        return tokenizer_obs_dict

    def create_encoder_masks(self, tokenizer_obs):
        """Create encoder masks for each encoder.

        When optimize_encoders_ratio_for_CHIP=False (legacy):
            - encoder_index is multi-hot: SMPL-native envs have both SMPL=1 and G1=1
            - encoder_masks["g1"] includes G1-native AND SMPL-native envs
            - encoder_masks["g1_has_smpl"] identifies which G1-masked envs are SMPL-native

        When optimize_encoders_ratio_for_CHIP=True (CHIP mode):
            - encoder_index is one-hot: each env has exactly one native encoder
            - encoder_masks are disjoint (no overlap)
            - encoder_masks["g1_has_smpl"] is empty (G1-native envs don't have SMPL)
            - G1 latents for SMPL-native envs are computed separately in aux losses
        """
        if len(self.encoders) == 1:
            return {list(self.encoders.keys())[0]: None}

        encoder_masks = {}
        encoder_index = tokenizer_obs["encoder_index"]

        for i, encoder_name in enumerate(self.encoder_sample_probs.keys()):
            encoder_masks[encoder_name] = encoder_index[..., i].bool().flatten()

        # Re-organized creation of intersection encoder masks.
        # NOTE: No need to use compliance encoders in this version -- directly use the normal encoders

        # For encoder pairs, create combined masks reflecting shared samples.
        # The pairs and their intersection masks:
        encoder_mask_combinations = [
            # (key1, key2, intersection_mask_names)
            ("g1", "smpl", [("g1_has_smpl", "smpl", "g1")]),
            (
                "teleop",
                "smpl",
                [("teleop_has_smpl", "smpl", "teleop"), ("smpl_has_teleop", "teleop", "smpl")],
            ),
            ("g1", "teleop", [("g1_has_teleop", "teleop", "g1")]),
            ("teleop", "g1", [("teleop_has_g1", "g1", "teleop")]),
            ("g1", "soma", [("g1_has_soma", "soma", "g1")]),
        ]
        for key1, key2, mask_defs in encoder_mask_combinations:
            if key1 in encoder_masks and key2 in encoder_masks:
                for mask_name, mask_src, mask_cond in mask_defs:
                    encoder_masks[mask_name] = encoder_masks[mask_src][encoder_masks[mask_cond]]
        return encoder_masks

    def assemble_all_tokens(self, encoded_tokens, encoder_masks, batch_size, seq_len):
        """Scatter per-encoder tokens into a single batch-aligned tensor.

        Each encoder processes only the environments assigned to it (via
        ``encoder_masks``).  This method writes each encoder's output back to
        the correct rows of a zero-initialised buffer, then reshapes to
        ``(batch, seq, num_tokens, token_dim)``.

        Args:
            encoded_tokens: Dict mapping encoder name to tensor of shape
                ``(num_masked, num_tokens, token_dim)``.
            encoder_masks: Dict mapping encoder name to a boolean index tensor
                of length ``batch*seq``.  ``None`` means all samples.
            batch_size: Batch dimension ``B``.
            seq_len: Sequence dimension ``S``.

        Returns:
            Tensor of shape ``(B, S, num_tokens, token_dim)`` with each row
            filled by its assigned encoder's output.
        """
        first_token = encoded_tokens[list(encoded_tokens.keys())[0]]
        if len(encoded_tokens) == 1:
            return first_token.view(batch_size, seq_len, *first_token.shape[1:])

        all_tokens = torch.zeros(
            (batch_size * seq_len,) + first_token.shape[1:],
            dtype=first_token.dtype,
            device=first_token.device,
        )
        for encoder_name in encoded_tokens.keys():
            all_tokens[encoder_masks[encoder_name]] = encoded_tokens[encoder_name]
        all_tokens = all_tokens.view(batch_size, seq_len, *all_tokens.shape[1:])
        return all_tokens

    def _encode_single(self, encoder_name, tokenizer_obs, encoder_mask=None, frame_mask=None):
        """Encode using a single encoder without quantization or additive composition.
        This is the core encoding logic used by both regular and additive encoders.

        Args:
            encoder_name: Name of the encoder to use
            tokenizer_obs: Dictionary of tokenizer observations
            encoder_mask: Optional mask for selecting specific samples
            frame_mask: Optional [B*seq, max_frames] bool mask for variable frame support

        Returns:
            latent: The encoded latent representation
        """  # noqa: D205
        encoder = sub_encoders = None
        if encoder_name in self.encoders:
            encoder = self.encoders[encoder_name]
        else:
            sub_encoders = [
                self.encoders[sub_encoder_name]
                for sub_encoder_name in self.sub_encoders[encoder_name]
            ]

        input_features = self.encoder_input_features[encoder_name]

        obs_list = [tokenizer_obs[key] for key in input_features]
        # Only handle dimension broadcasting for the "smpl" encoder_name and for obs with last dim 9 or 12
        if "smpl" in encoder_name:
            max_ndim = max(obs.ndim for obs in obs_list)
            # Check if any obs needs broadcasting: (last dim is 9 or 12 and ndim < max_ndim)
            need_broadcast = any(
                (obs.shape[-1] in (9, 12) and obs.ndim < max_ndim) for obs in obs_list
            )
            if need_broadcast:
                aligned_obs_list = []
                # Use the first obs of max_ndim as the temporal reference
                target_obs = next(o for o in obs_list if o.ndim == max_ndim)
                temporal_dim = target_obs.shape[-2]
                for obs in obs_list:
                    if obs.shape[-1] in (9, 12) and obs.ndim < max_ndim:
                        # Add missing dimension(s)
                        while obs.ndim < max_ndim:
                            obs = obs.unsqueeze(-2)  # noqa: PLW2901
                        # Expand temporal dimension
                        expand_shape = list(obs.shape)
                        expand_shape[-2] = temporal_dim
                        obs = obs.expand(*expand_shape)  # noqa: PLW2901
                    aligned_obs_list.append(obs)
                encoder_input = torch.cat(aligned_obs_list, dim=-1)
            else:
                encoder_input = torch.cat(obs_list, dim=-1)
        else:
            encoder_input = torch.cat(obs_list, dim=-1)
        encoder_input = encoder_input.view(-1, *encoder_input.shape[2:])

        # Apply encoder_mask to both input and frame_mask
        frame_mask_enc = None
        if encoder_mask is not None:
            encoder_input = encoder_input[encoder_mask]
            if frame_mask is not None:
                frame_mask_enc = frame_mask[encoder_mask]
        elif frame_mask is not None:
            frame_mask_enc = frame_mask

        # encode using the provided encoder (with gradients enabled for end-to-end training)
        if encoder is not None:
            if frame_mask_enc is not None and self.encoder_mask_features.get(encoder_name):
                latent = encoder(encoder_input, frame_mask=frame_mask_enc)
            else:
                latent = encoder(encoder_input)
        else:
            latent = encoder_input
            for sub_encoder in sub_encoders:
                latent = sub_encoder(latent)

        return latent

    def encode(
        self, encoder_name, tokenizer_obs, encoder_mask=None, no_quantization=False, frame_mask=None
    ):
        """Encode using the specified encoder, including any additive encoders.

        Args:
            encoder_name: Name of the base encoder to use
            tokenizer_obs: Dictionary of tokenizer observations
            encoder_mask: Optional mask for selecting specific samples
            no_quantization: If True, skip quantization and return raw latent
            frame_mask: Optional [B*seq, max_frames] bool mask for variable frame support

        Returns:
            If no_quantization: latent tensor
            Otherwise: (encoded_tokens, latent) tuple
        """
        # Get base encoder latent
        latent = self._encode_single(
            encoder_name, tokenizer_obs, encoder_mask, frame_mask=frame_mask
        )

        # Add contributions from additive encoders
        if encoder_name in self.additive_encoders:
            for additive_encoder_name in self.additive_encoders[encoder_name]:
                additive_latent = self._encode_single(
                    additive_encoder_name, tokenizer_obs, encoder_mask, frame_mask=frame_mask
                )
                latent = latent + additive_latent

        if no_quantization:
            return latent

        # quantize using quantizer
        if self.quantizer is not None:
            quantized_codes, _ = self.quantizer(latent)
            encoded_tokens = quantized_codes.contiguous()
        else:
            encoded_tokens = latent
        return encoded_tokens, latent

    def decode(self, decoder_name, decode_input_dict, token_mask=None):
        """Run a single named decoder and split its output by feature.

        Concatenates the decoder's declared input features from
        ``decode_input_dict``, passes them through the decoder network
        (optionally supplying ``external_cond`` and ``token_mask``), then
        slices the output back into a per-feature dict.

        Args:
            decoder_name: Key into ``self.decoders`` and associated metadata.
            decode_input_dict: Dict containing at least the keys declared in
                the decoder's ``inputs`` and ``conds`` config entries.
                Common keys include ``"token"``, ``"token_flattened"``, and
                ``"proprioception"``.
            token_mask: Optional ``(B*seq, max_tokens)`` bool mask for
                variable-length token sequences.  Only forwarded to decoders
                that declare mask features.

        Returns:
            Dict mapping output feature name to the corresponding slice of the
            decoder output tensor, e.g.
            ``{"action": Tensor(..., action_dim)}``.
        """
        decoder = self.decoders[decoder_name]
        input_features = self.decoder_input_features[decoder_name]
        output_feature_dims = self.decoder_output_feature_dims[decoder_name]
        cond_features = self.decoder_cond_features.get(decoder_name, [])
        decoder_input = torch.cat([decode_input_dict[key] for key in input_features], dim=-1)

        # Build optional kwargs for the decoder call
        kwargs = {}
        if cond_features:
            kwargs["external_cond"] = torch.cat(
                [decode_input_dict[key] for key in cond_features], dim=-1
            )
        if token_mask is not None and self.decoder_mask_features.get(decoder_name):
            kwargs["token_mask"] = token_mask
        output = decoder(decoder_input, **kwargs)

        # parse output
        output_dict = {}
        index = 0
        for key, dim in output_feature_dims.items():
            output_dict[key] = output[..., index : index + dim]
            index += dim
        assert index == output.shape[-1], f"{index=}, {output.shape[-1]=}"

        return output_dict

    def forward(  # noqa: D417
        self,
        input_data,
        compute_aux_loss=False,
        return_dict=False,
        latent_residual=None,
        latent_residual_mode="post_quantization",
        **kwargs,  # noqa: ARG002
    ):
        """Run the full encode → quantize → decode pipeline.

        Parses tokenizer observations, routes each environment to its assigned
        encoder, optionally applies an external latent residual, then decodes
        to joint-space actions.  When ``compute_aux_loss=True``, also computes
        all registered auxiliary losses (e.g. G1-SMPL alignment, cycle
        consistency) and returns a rich result dict.

        Args:
            input_data: Dict of named observation tensors, all with leading
                shape ``(B, S, ...)``.  Must contain at minimum ``"actor_obs"``
                (for batch/seq inference) and ``"tokenizer"`` (flat tokenizer
                obs of shape ``(B, S, total_tokenizer_dim)``).
            compute_aux_loss: When ``True``, evaluate all auxiliary loss
                functions and include them in the returned dict.
            return_dict: When ``True`` and ``compute_aux_loss=False``, return
                the full output dict instead of just ``action_mean``.
            latent_residual: Optional additive correction in latent token
                space.  Shape ``(B, token_total_dim)`` where
                ``token_total_dim = max_num_tokens * token_dim``.  Allows HOI
                policies to steer motion without modifying the base ATM.
            latent_residual_mode: Controls *when* ``latent_residual`` is
                applied:

                * ``"post_quantization"`` - add after FSQ (default, residual
                  stays continuous).
                * ``"pre_quantization"`` - add before FSQ (residual gets
                  quantized together with the encoder latent).
                * ``"pre_quantization_replace"`` - replace the encoder latent
                  entirely with ``latent_residual`` before quantization.
            **kwargs: Passed through; currently unused.

        Returns:
            When ``compute_aux_loss=True`` or ``return_dict=True``: a dict
            with keys:

            * ``"action_mean"`` - joint targets, shape ``(B, S, action_dim)``.
            * ``"aux_losses"`` - dict of scalar loss tensors (empty when
              ``compute_aux_loss=False``).
            * ``"aux_loss_coef"`` - per-loss coefficient dict from config.
            * ``"decoded_outputs"`` - raw decoder output dicts keyed by
              decoder name.
            * ``"tokenizer_obs"`` - parsed tokenizer observation dict.
            * ``"encoder_masks"`` - per-encoder boolean index tensors.
            * ``"encoded_tokens"`` - post-quantization tokens per encoder.
            * ``"encoded_latents"`` - pre-quantization latents per encoder.
            * ``"encoders_cfg"`` / ``"decoders_cfg"`` - config references.

            When ``compute_aux_loss=False`` and ``return_dict=False``:
            ``action_mean`` tensor directly.
        """
        # parse tokenizer obs
        batch_size, seq_len = input_data["actor_obs"].shape[:2]
        tokenizer_obs = self.parse_tokenizer_obs(input_data)
        proprioception_input = torch.cat(
            [input_data[key] for key in self.proprioception_features], dim=-1
        )

        # Reshape residual if provided: (batch, token_total_dim) -> (batch, 1, max_num_tokens, token_dim)
        residual_reshaped = None
        if latent_residual is not None:
            residual_reshaped = latent_residual.view(
                batch_size, 1, self.max_num_tokens, self.token_dim
            )

        # Variable frame masks
        frame_mask, token_mask = self._create_frame_and_token_masks(tokenizer_obs)

        # encode motion using all available encoders
        encoder_masks = self.create_encoder_masks(tokenizer_obs)
        encoded_tokens = {}
        encoded_latents = {}
        if latent_residual is not None and latent_residual_mode in [
            "pre_quantization",
            "pre_quantization_replace",
        ]:
            # PRE-QUANTIZATION MODES: add or replace residual before FSQ
            # Reshape residual to (batch*seq, num_tokens, token_dim) for masking
            residual_flat = latent_residual.view(
                batch_size * seq_len, self.max_num_tokens, self.token_dim
            )

            for encoder_name in self.encoders_to_iterate:
                encoder_mask = encoder_masks[encoder_name]

                # Get raw latent (including additive encoders) without quantization
                latent = self.encode(
                    encoder_name,
                    tokenizer_obs,
                    encoder_mask,
                    no_quantization=True,
                    frame_mask=frame_mask,
                )
                # Apply same mask to residual before adding
                if encoder_mask is not None:
                    masked_residual = residual_flat[encoder_mask]
                else:
                    masked_residual = residual_flat

                # Apply residual before quantization (only if we have samples)
                if latent.shape[0] > 0:
                    if latent_residual_mode == "pre_quantization":
                        # Add residual to encoder latent
                        latent = latent + masked_residual
                    elif latent_residual_mode == "pre_quantization_replace":
                        # Replace encoder latent with residual (zero out encoder)
                        latent = masked_residual
                    else:
                        raise ValueError(f"Unknown latent_residual_mode: {latent_residual_mode}")

                # Now quantize
                if self.quantizer is not None:
                    quantized_codes, _ = self.quantizer(latent)
                    encoded_tokens[encoder_name] = quantized_codes.contiguous()
                else:
                    encoded_tokens[encoder_name] = latent
                encoded_latents[encoder_name] = latent
        else:
            # STANDARD MODE: encode normally
            for encoder_name in self.encoders_to_iterate:
                encoded_tokens[encoder_name], encoded_latents[encoder_name] = self.encode(
                    encoder_name,
                    tokenizer_obs,
                    encoder_masks[encoder_name],
                    frame_mask=frame_mask,
                )
        all_tokens = self.assemble_all_tokens(encoded_tokens, encoder_masks, batch_size, seq_len)

        # POST-QUANTIZATION MODE: add residual after FSQ tokens (default)
        if latent_residual is not None and latent_residual_mode == "post_quantization":
            all_tokens = all_tokens + residual_reshaped

        # Cache tokens for external access (e.g., by callbacks)
        self._last_encoded_tokens = {k: v.detach().cpu() for k, v in encoded_tokens.items()}
        self._last_encoded_latents = {k: v.detach().cpu() for k, v in encoded_latents.items()}

        # Cache flattened full latent on device for reward computation (token smoothness)
        # all_tokens is the post-quantization token that gets sent to decoder
        # Shape: (batch, seq, num_tokens, token_dim) -> (batch, seq, latent_dim)
        self._last_full_latent_flat = all_tokens.detach().view(*all_tokens.shape[:-2], -1)

        # decode action and motion
        decode_input_dict = {
            "token": all_tokens,
            "token_flattened": all_tokens.view(*all_tokens.shape[:-2], -1),
            "proprioception": proprioception_input,
        }
        decode_input_dict.update(tokenizer_obs)

        decoded_outputs = {}
        decoders_to_run = (
            self._active_decoders if self._active_decoders is not None else self.decoders.keys()
        )
        for decoder_name in decoders_to_run:
            decoded_outputs[decoder_name] = self.decode(
                decoder_name, decode_input_dict, token_mask=token_mask
            )

        # Support "body_action", "meta_action", and "action" outputs (g1_dyn may not exist in kinematic-only mode)
        if "g1_dyn" in decoded_outputs:
            g1_dyn_out = decoded_outputs["g1_dyn"]
            if "body_action" in g1_dyn_out:
                action_mean = g1_dyn_out["body_action"]
            elif "meta_action" in g1_dyn_out:
                action_mean = g1_dyn_out["meta_action"]
            else:
                action_mean = g1_dyn_out["action"]

            # Concatenate hand decoder output if present
            if "hand_dyn" in decoded_outputs and "hand_action" in decoded_outputs["hand_dyn"]:
                hand_action = decoded_outputs["hand_dyn"]["hand_action"]
                action_mean = torch.cat([action_mean, hand_action], dim=-1)
        else:
            action_mean = None

        # compute aux losses
        if compute_aux_loss:
            aux_losses = {}

            # Initialize all paired latents
            reencoded_smpl_g1_latents = None
            paired_g1_smpl_latents = None
            paired_compliance_latents = None
            original_g1_latents_for_reencode = None

            # Determine encoder names
            smpl_encoder_name = "smpl" if "smpl" in self.encoders_to_iterate else None
            teleop_encoder_name = "teleop" if "teleop" in self.encoders_to_iterate else None

            # =========================================================================
            # STIFF-MODE OPTIMIZATION: Check for stiff samples FIRST before computing
            # expensive G1 latents. G1-SMPL loss and cycle consistency loss ONLY apply
            # in stiff mode (compliance ≈ 0). Skip computation if all samples are compliant.
            # =========================================================================
            has_stiff_samples = False
            compliance_values = None
            smpl_latents = None
            smpl_mask = None

            if smpl_encoder_name is not None and "g1" in encoded_latents:
                smpl_mask = encoder_masks.get(smpl_encoder_name)

                if smpl_mask is not None and smpl_mask.sum() > 0:
                    smpl_latents = encoded_latents[smpl_encoder_name]

                    # Extract compliance values for these envs FIRST
                    if "compliance" in tokenizer_obs:
                        compliance_flat = tokenizer_obs["compliance"].view(
                            -1, tokenizer_obs["compliance"].shape[-1]
                        )
                        compliance_values = compliance_flat[smpl_mask]

                        # Check if ANY samples are stiff
                        is_stiff = (compliance_values.abs() < self.stiff_compliance_threshold).all(
                            dim=-1
                        )
                        has_stiff_samples = is_stiff.any().item()
                    else:
                        # No compliance info means all samples are "stiff" (default behavior)
                        has_stiff_samples = True

            # Only compute stiff-mode losses if there are stiff samples
            if has_stiff_samples and smpl_mask is not None:
                # Compute G1 latents for the same envs as SMPL (for proper pairing)
                # OPTIMIZATION: Detach here since all downstream losses use g1 as
                # fixed target (detach_g1_target=True). This avoids redundant gradient
                # computation through the g1 encoder for these latents.
                g1_latents_for_smpl = self.encode(
                    encoder_name="g1",
                    tokenizer_obs=tokenizer_obs,
                    encoder_mask=smpl_mask,
                    no_quantization=True,
                ).detach()

                paired_g1_smpl_latents = {
                    "g1": g1_latents_for_smpl,
                    "smpl": smpl_latents,
                    "compliance": compliance_values,
                }

                # Store original G1 latents for cycle consistency loss (already detached)
                original_g1_latents_for_reencode = g1_latents_for_smpl

                # Cycle consistency: re-encode decoded G1 motion (only if stiff samples exist)
                if self.reencode_smpl_g1_recon:
                    reencoded_smpl_g1_latents = self.encode(
                        encoder_name="g1",
                        tokenizer_obs=decoded_outputs["g1_kin"],
                        encoder_mask=smpl_mask,
                        no_quantization=True,
                    )

            # Compute paired teleop-SMPL latents for TeleopSmplComplianceLatentLoss
            # NOTE: This applies to ALL compliance values (not just stiff), so no stiff check
            if (
                teleop_encoder_name is not None
                and smpl_latents is not None
                and smpl_mask is not None
            ):
                # Compute teleop latents for the same envs as SMPL
                # OPTIMIZATION: Detach here since TeleopSmplComplianceLatentLoss uses
                # teleop as teacher (detach_teleop_target=True). Only SMPL learns.
                teleop_latents_for_smpl = self.encode(
                    encoder_name=teleop_encoder_name,
                    tokenizer_obs=tokenizer_obs,
                    encoder_mask=smpl_mask,
                    no_quantization=True,
                ).detach()

                paired_compliance_latents = {
                    "teleop": teleop_latents_for_smpl,
                    "smpl": smpl_latents,
                }

            loss_inputs = {
                "input_data": input_data,
                "tokenizer_obs": tokenizer_obs,
                "encoded_tokens": encoded_tokens,
                "encoded_latents": encoded_latents,
                "encoder_masks": encoder_masks,
                "decoded_outputs": decoded_outputs,
                "action_mean": action_mean,
                "encoders_cfg": self.encoders_cfg,
                "frame_mask": (
                    frame_mask.view(batch_size, seq_len, -1) if frame_mask is not None else None
                ),
                "token_mask": (
                    token_mask.view(batch_size, seq_len, -1) if token_mask is not None else None
                ),
                "decoders_cfg": self.decoders_cfg,
                "reencoded_smpl_g1_latents": reencoded_smpl_g1_latents,
                # New fields for compliance-aware losses
                "paired_g1_smpl_latents": paired_g1_smpl_latents,
                "paired_compliance_latents": paired_compliance_latents,
                "original_g1_latents_for_reencode": original_g1_latents_for_reencode,
            }

            for loss_name, loss_func in self.aux_loss_func.items():
                aux_losses[loss_name] = loss_func(loss_inputs)

            output = {
                "action_mean": action_mean,
                "aux_losses": aux_losses,
                "aux_loss_coef": self.aux_loss_coef,
                # Exposed for recon trainer (existing RL callers ignore these keys)
                "decoded_outputs": decoded_outputs,
                "tokenizer_obs": tokenizer_obs,
                "encoder_masks": encoder_masks,
                "encoded_tokens": encoded_tokens,
                "encoded_latents": encoded_latents,
                "encoders_cfg": self.encoders_cfg,
                "decoders_cfg": self.decoders_cfg,
            }
        elif return_dict:
            output = {
                "action_mean": action_mean,
                "aux_losses": {},
                "aux_loss_coef": self.aux_loss_coef,
                "decoded_outputs": decoded_outputs,
                "tokenizer_obs": tokenizer_obs,
                "encoder_masks": encoder_masks,
                "encoded_tokens": encoded_tokens,
                "encoded_latents": encoded_latents,
                "encoders_cfg": self.encoders_cfg,
                "decoders_cfg": self.decoders_cfg,
            }
        else:
            output = action_mean
        return output

    def get_token_info(self):
        """Return a summary of the FSQ token configuration.

        Returns:
            Dict with keys:

            * ``"token_dim"`` - dimensionality of one token.
            * ``"total_dim"`` - ``token_dim * max_num_tokens``.
            * ``"num_levels"`` - number of FSQ levels.
            * ``"level_list"`` - per-level codebook sizes, or ``None`` if the
              quantizer does not expose a ``levels`` attribute.
            * ``"model_available"`` - always ``True``.
        """
        return {
            "token_dim": self.token_dim,
            "total_dim": self.token_total_dim,
            "num_levels": self.num_fsq_levels,
            "level_list": (
                list(self.quantizer.levels) if hasattr(self.quantizer, "levels") else None
            ),
            "model_available": True,
        }

    def forward_with_external_tokens(  # noqa: D417
        self, input_data: dict, external_tokens: torch.Tensor, **kwargs  # noqa: ARG002
    ) -> torch.Tensor:
        """Forward pass with externally provided tokens (bypasses encoder).

        This is used when an external model (e.g., kinematic diffusion) provides
        pre-computed FSQ tokens, allowing the encoder to be bypassed while still
        using the decoder for action generation.

        Args:
            input_data: Dict with 'actor_obs' for proprioception
            external_tokens: Pre-computed FSQ tokens from kinematic diffusion
                            Shape: (B, 2, 32) or (B, seq, 2, 32)

        Returns:
            action_mean: (B, action_dim) or (B, seq, action_dim)
        """
        # Get proprioception input
        proprioception_input = torch.cat(
            [input_data[key] for key in self.proprioception_features], dim=-1
        )

        # Handle different input shapes
        if external_tokens.dim() == 3:
            # (B, 2, 32) -> add seq dim -> (B, 1, 2, 32)
            external_tokens = external_tokens.unsqueeze(1)

        # external_tokens: (B, seq, 2, 32) or (B, seq, num_tokens, token_dim)
        batch_size = external_tokens.shape[0]
        seq_len = external_tokens.shape[1]

        # Flatten tokens for decoder: (B, seq, 2, 32) -> (B, seq, 64)
        token_flattened = external_tokens.view(batch_size, seq_len, -1)

        # Build decode input dict
        decode_input_dict = {
            "token": external_tokens,
            "token_flattened": token_flattened,
            "proprioception": proprioception_input,
        }

        # Decode actions using g1_dyn decoder
        g1_dyn_out = self.decode("g1_dyn", decode_input_dict)
        if "body_action" in g1_dyn_out:
            action_mean = g1_dyn_out["body_action"]
        elif "meta_action" in g1_dyn_out:
            action_mean = g1_dyn_out["meta_action"]
        else:
            action_mean = g1_dyn_out["action"]

        # Concatenate hand decoder output if present
        if "hand_dyn" in self.decoders:
            hand_dyn_out = self.decode("hand_dyn", decode_input_dict)
            if "hand_action" in hand_dyn_out:
                action_mean = torch.cat([action_mean, hand_dyn_out["hand_action"]], dim=-1)

        # Squeeze seq dim if it was added
        if action_mean.shape[1] == 1:
            action_mean = action_mean.squeeze(1)

        return action_mean

    def get_example_input(self, encoder_name, batch_size=1, device="cpu"):
        """Generate example input for ONNX export with specific encoder.

        Args:
            encoder_name: Name of the encoder to use
            batch_size: Batch size for the example input
            device: Device to create tensors on

        Returns:
            Tensor with shape (batch_size, feature_dim) including encoder features and proprioception
        """
        if encoder_name not in self.encoders:
            raise ValueError(
                f"Encoder '{encoder_name}' not found. Available encoders: {list(self.encoders.keys())}"
            )

        # Calculate total input feature dimension for the specified encoder
        encoder_input_features = self.encoder_input_features[encoder_name]
        total_feature_dim = 0

        for feature_name in encoder_input_features:
            if feature_name in self.tokenizer_obs_dims:
                feature_dims = self.tokenizer_obs_dims[feature_name]
                total_feature_dim += torch.prod(torch.tensor(feature_dims)).item()

        # Add proprioception dimension
        proprioception_dim = self.obs_dim_dict.get("actor_obs", 0)
        total_feature_dim += proprioception_dim

        # Create a single 2D tensor (B, F) including both encoder features and proprioception
        example_input = torch.randn(batch_size, total_feature_dim, device=device)

        return example_input

    def get_all_example_input(self, batch_size=1, device="cpu"):
        """Generate example inputs covering all tokenizer observations.

        Constructs random tensors whose shapes match the full tokenizer
        observation space plus proprioception.  Intended for ONNX tracing
        or shape debugging.

        Args:
            batch_size: Number of examples in the batch dimension.
            device: Target device for the returned tensors (e.g. ``"cpu"``
                or ``"cuda"``).

        Returns:
            Tuple of:

            * ``example_input`` - flat tensor of shape
              ``(batch_size, total_feature_dim)`` where
              ``total_feature_dim = tokenizer_feature_dim + proprioception_dim``.
            * ``example_dict`` - dict with keys ``"tokenizer"`` and
              ``"proprioception"``, each a random tensor with the
              corresponding shape.
        """
        # Calculate total input feature dimension for the specified encoder
        total_feature_dim = 0
        for feature_name in self.tokenizer_obs_names:
            feature_dims = self.tokenizer_obs_dims[feature_name]
            print(  # noqa: T201
                f"{feature_name}: {feature_dims}, start: {total_feature_dim} end: {total_feature_dim + torch.prod(torch.tensor(feature_dims)).item()}"  # noqa: E501
            )
            total_feature_dim += torch.prod(torch.tensor(feature_dims)).item()
        tokenizer_feature_dim = total_feature_dim

        # Add proprioception dimension
        proprioception_dim = self.obs_dim_dict.get("actor_obs", 0)
        total_feature_dim += proprioception_dim

        print(f"tokenizer_feature_dim: {tokenizer_feature_dim}")  # noqa: T201
        print(f"proprioception_dim: {proprioception_dim}")  # noqa: T201

        # Create a single 2D tensor (B, F) including both encoder features and proprioception
        example_input = torch.randn(batch_size, total_feature_dim, device=device)

        example_dict = {
            "tokenizer": torch.randn(batch_size, tokenizer_feature_dim, device=device),
            "proprioception": torch.randn(batch_size, proprioception_dim, device=device),
        }

        return example_input, example_dict