dhb-xr 0.2.1__py3-none-any.whl

dhb_xr/tokenization/hierarchical.py

@@ -0,0 +1,359 @@
+"""
+Hierarchical tokenization for variable-rate DHB compression.
+
+Extends RVQ with multi-level structure for lossy compression:
+- Coarse tokens: Global trajectory structure
+- Fine tokens: Local details and refinements
+- Configurable depth for compression/quality tradeoff
+
+Inspired by:
+- BEAST (NeurIPS 2025): B-spline encoded action sequences
+- VQ-VLA: Multi-level vector quantization
+"""
+
+try:
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    from dhb_xr.tokenization.vqvae import VectorQuantizer
+    from dhb_xr.tokenization.causal_encoder import CausalConv1dEncoder
+    HAS_TORCH = True
+except ImportError:
+    HAS_TORCH = False
+
+
+if HAS_TORCH:
+
+    class HierarchicalTokenizer(nn.Module):
+        """
+        Hierarchical RVQ with variable-rate output.
+        
+        Provides coarse-to-fine tokenization:
+        - Level 0: Low-frequency global structure (high compression)
+        - Level 1-N: Residual details (configurable refinement)
+        
+        For inference, can truncate to fewer levels for faster/coarser output.
+        
+        Example:
+            >>> tokenizer = HierarchicalTokenizer(
+            ...     invariant_dim=8, latent_dim=32, codebook_size=256, num_levels=4
+            ... )
+            >>> tokens, recon = tokenizer(invariants)
+            >>> 
+            >>> # Coarse only (4x fewer tokens)
+            >>> tokens_coarse, recon_coarse = tokenizer(invariants, max_level=1)
+        """
+        
+        def __init__(
+            self,
+            invariant_dim: int,
+            latent_dim: int,
+            codebook_size: int,
+            num_levels: int = 4,
+            temporal_downsample: int = 2,
+            num_layers: int = 2,
+        ):
+            """
+            Args:
+                invariant_dim: DHB invariant dimension (typically 8)
+                latent_dim: Latent embedding dimension
+                codebook_size: VQ codebook size per level
+                num_levels: Number of hierarchy levels
+                temporal_downsample: Downsample factor between levels
+                num_layers: Conv layers per encoder/decoder
+            """
+            super().__init__()
+            
+            self.invariant_dim = invariant_dim
+            self.latent_dim = latent_dim
+            self.codebook_size = codebook_size
+            self.num_levels = num_levels
+            self.temporal_downsample = temporal_downsample
+            
+            # Per-level encoders (progressively downsample)
+            self.encoders = nn.ModuleList()
+            self.vqs = nn.ModuleList()
+            self.decoders = nn.ModuleList()
+            
+            for level in range(num_levels):
+                # Encoder: downsample temporally at each level
+                if level == 0:
+                    enc = CausalConv1dEncoder(
+                        invariant_dim, latent_dim, latent_dim, num_layers
+                    )
+                else:
+                    enc = nn.Sequential(
+                        CausalConv1dEncoder(latent_dim, latent_dim, latent_dim, num_layers),
+                        TemporalDownsample(temporal_downsample),
+                    )
+                self.encoders.append(enc)
+                
+                # VQ at each level
+                self.vqs.append(VectorQuantizer(codebook_size, latent_dim))
+                
+                # Decoder: upsample to match previous level
+                if level == 0:
+                    dec = CausalConv1dEncoder(
+                        latent_dim, latent_dim, invariant_dim, num_layers
+                    )
+                else:
+                    dec = nn.Sequential(
+                        TemporalUpsample(temporal_downsample),
+                        CausalConv1dEncoder(latent_dim, latent_dim, latent_dim, num_layers),
+                    )
+                self.decoders.append(dec)
+            
+            # Final projection back to invariant space
+            self.final_proj = nn.Linear(latent_dim, invariant_dim)
+        
+        def forward(
+            self, 
+            invariants: torch.Tensor, 
+            max_level: int = None,
+            return_all_levels: bool = False,
+        ) -> tuple:
+            """
+            Hierarchical encoding and decoding.
+            
+            Args:
+                invariants: (B, T, invariant_dim) input
+                max_level: Stop at this level (None = all levels)
+                return_all_levels: Return tokens/recon at each level
+                
+            Returns:
+                all_tokens: List of (B, T_l) tokens per level
+                reconstructed: (B, T, invariant_dim) reconstruction
+                level_info: Optional dict with per-level details
+            """
+            B, T, C = invariants.shape
+            max_level = max_level or self.num_levels
+            
+            all_tokens = []
+            all_z = []
+            all_z_q = []
+            level_info = {}
+            
+            # Encode through hierarchy
+            x = invariants
+            for level in range(max_level):
+                z = self.encoders[level](x if level == 0 else z_residual)
+                indices, z_q_st, z_q = self.vqs[level](z)
+                
+                all_tokens.append(indices)
+                all_z.append(z)
+                all_z_q.append(z_q)
+                
+                if level < max_level - 1:
+                    z_residual = z - z_q.detach()
+                
+                level_info[f"level_{level}"] = {
+                    "shape": tuple(z.shape),
+                    "tokens": indices.shape[-1],
+                }
+            
+            # Decode through hierarchy (reverse order)
+            reconstructed = torch.zeros_like(invariants)
+            for level in reversed(range(max_level)):
+                dec_out = self.decoders[level](all_z_q[level])
+                
+                # Match temporal dimension
+                if dec_out.shape[1] > reconstructed.shape[1]:
+                    dec_out = dec_out[:, :reconstructed.shape[1], :]
+                elif dec_out.shape[1] < reconstructed.shape[1]:
+                    # Upsample to match
+                    dec_out = F.interpolate(
+                        dec_out.transpose(1, 2), 
+                        size=reconstructed.shape[1],
+                        mode='linear',
+                        align_corners=True
+                    ).transpose(1, 2)
+                
+                if level == 0:
+                    reconstructed = dec_out
+                else:
+                    reconstructed = reconstructed + self.final_proj(dec_out)
+            
+            if return_all_levels:
+                return all_tokens, reconstructed, level_info
+            return all_tokens, reconstructed
+        
+        def loss(
+            self,
+            invariants: torch.Tensor,
+            reconstructed: torch.Tensor,
+            all_z: list,
+            all_z_q: list,
+            beta: float = 0.25,
+            level_weights: list = None,
+        ) -> torch.Tensor:
+            """
+            Compute hierarchical loss with per-level weighting.
+            
+            Args:
+                invariants: Original input
+                reconstructed: Reconstruction
+                all_z: Latents at each level
+                all_z_q: Quantized latents at each level
+                beta: Commitment loss weight
+                level_weights: Optional weights per level (default: exponential decay)
+                
+            Returns:
+                Total loss
+            """
+            # Reconstruction loss
+            rec_loss = F.mse_loss(reconstructed, invariants)
+            
+            # Per-level VQ losses
+            if level_weights is None:
+                level_weights = [0.5 ** i for i in range(len(all_z))]
+            
+            commitment = 0
+            codebook = 0
+            for i, (z, z_q) in enumerate(zip(all_z, all_z_q)):
+                commitment += level_weights[i] * F.mse_loss(z, z_q.detach())
+                codebook += level_weights[i] * F.mse_loss(z_q, z.detach())
+            
+            return rec_loss + beta * commitment + codebook
+        
+        def get_compression_stats(self, T: int, max_level: int = None) -> dict:
+            """
+            Compute compression statistics.
+            
+            Args:
+                T: Original sequence length
+                max_level: Number of levels to use
+                
+            Returns:
+                Compression statistics
+            """
+            max_level = max_level or self.num_levels
+            
+            total_tokens = 0
+            for level in range(max_level):
+                level_T = T // (self.temporal_downsample ** level)
+                total_tokens += level_T
+            
+            original_values = T * self.invariant_dim
+            token_values = total_tokens  # Each token is one index
+            
+            return {
+                "original_values": original_values,
+                "total_tokens": total_tokens,
+                "tokens_per_level": [T // (self.temporal_downsample ** l) for l in range(max_level)],
+                "compression_ratio": original_values / token_values if token_values > 0 else 1,
+                "bits_per_value": (total_tokens * np.log2(self.codebook_size)) / original_values,
+            }
+
+
+    class TemporalDownsample(nn.Module):
+        """Temporal downsampling via strided convolution."""
+        
+        def __init__(self, factor: int = 2):
+            super().__init__()
+            self.factor = factor
+        
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            # x: (B, T, D)
+            return x[:, ::self.factor, :]
+
+
+    class TemporalUpsample(nn.Module):
+        """Temporal upsampling via interpolation."""
+        
+        def __init__(self, factor: int = 2):
+            super().__init__()
+            self.factor = factor
+        
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            # x: (B, T, D)
+            B, T, D = x.shape
+            x = x.transpose(1, 2)  # (B, D, T)
+            x = F.interpolate(x, size=T * self.factor, mode='linear', align_corners=True)
+            return x.transpose(1, 2)  # (B, T*factor, D)
+
+
+    class ProgressiveTokenizer(nn.Module):
+        """
+        Progressive refinement tokenizer.
+        
+        Outputs can be truncated at any level for variable-rate decoding:
+        - 1 level: ~4x compression, coarse motion
+        - 2 levels: ~2x compression, medium detail
+        - 4 levels: ~1x compression, full fidelity
+        
+        Ideal for streaming/bandwidth-adaptive applications.
+        """
+        
+        def __init__(
+            self,
+            invariant_dim: int,
+            latent_dim: int,
+            codebook_size: int,
+            num_refinements: int = 3,
+        ):
+            super().__init__()
+            
+            self.base_tokenizer = nn.Sequential(
+                CausalConv1dEncoder(invariant_dim, latent_dim, latent_dim, 2),
+            )
+            self.base_vq = VectorQuantizer(codebook_size, latent_dim)
+            self.base_decoder = CausalConv1dEncoder(latent_dim, latent_dim, invariant_dim, 2)
+            
+            # Refinement stages
+            self.refinements = nn.ModuleList()
+            self.refine_vqs = nn.ModuleList()
+            self.refine_decoders = nn.ModuleList()
+            
+            for _ in range(num_refinements):
+                self.refinements.append(
+                    CausalConv1dEncoder(invariant_dim, latent_dim, latent_dim, 1)
+                )
+                self.refine_vqs.append(VectorQuantizer(codebook_size, latent_dim))
+                self.refine_decoders.append(
+                    CausalConv1dEncoder(latent_dim, latent_dim, invariant_dim, 1)
+                )
+        
+        def forward(self, invariants: torch.Tensor, num_refine: int = None) -> tuple:
+            """
+            Progressive tokenization.
+            
+            Args:
+                invariants: (B, T, D) input
+                num_refine: Number of refinement levels (0 = base only)
+                
+            Returns:
+                all_tokens: List of token tensors
+                reconstructed: Final reconstruction
+            """
+            if num_refine is None:
+                num_refine = len(self.refinements)
+            
+            # Base encoding
+            z_base = self.base_tokenizer(invariants)
+            tokens_base, z_q_st, z_q = self.base_vq(z_base)
+            recon = self.base_decoder(z_q_st)
+            
+            all_tokens = [tokens_base]
+            
+            # Progressive refinements
+            residual = invariants - recon
+            for i in range(min(num_refine, len(self.refinements))):
+                z_ref = self.refinements[i](residual)
+                tokens_ref, z_ref_st, z_ref_q = self.refine_vqs[i](z_ref)
+                
+                all_tokens.append(tokens_ref)
+                
+                recon = recon + self.refine_decoders[i](z_ref_st)
+                residual = invariants - recon
+            
+            return all_tokens, recon
+
+else:
+    HierarchicalTokenizer = None
+    ProgressiveTokenizer = None
+    TemporalDownsample = None
+    TemporalUpsample = None
+
+
+# Import numpy for compression stats (needed even without torch)
+import numpy as np

dhb_xr/tokenization/rvq.py

@@ -0,0 +1,178 @@
+"""Residual VQ (RVQ) tokenizer for higher capacity."""
+
+try:
+    import torch
+    import torch.nn as nn
+    from dhb_xr.tokenization.vqvae import VectorQuantizer
+    from dhb_xr.tokenization.causal_encoder import CausalConv1dEncoder
+    HAS_TORCH = True
+except ImportError:
+    HAS_TORCH = False
+
+if HAS_TORCH:
+
+    class ResidualVQTokenizer(nn.Module):
+        """RVQ: multiple codebooks on residual. invariants (B, T, C) -> list of (B, T) tokens, (B, T, C) reconstructed."""
+
+        def __init__(
+            self,
+            invariant_dim: int,
+            latent_dim: int,
+            codebook_size: int,
+            num_codebooks: int = 2,
+            num_layers: int = 2,
+        ):
+            super().__init__()
+            self.encoder = CausalConv1dEncoder(
+                invariant_dim, latent_dim, latent_dim, num_layers
+            )
+            self.vqs = nn.ModuleList([
+                VectorQuantizer(codebook_size, latent_dim) for _ in range(num_codebooks)
+            ])
+            self.decoder = CausalConv1dEncoder(
+                latent_dim, latent_dim, invariant_dim, num_layers
+            )
+            self.num_codebooks = num_codebooks
+
+        def forward(self, invariants: torch.Tensor) -> tuple:
+            z = self.encoder(invariants)
+            residuals = z
+            all_indices = []
+            z_sum = torch.zeros_like(z)
+            for vq in self.vqs:
+                indices, z_q_st, z_q = vq(residuals)
+                all_indices.append(indices)
+                z_sum = z_sum + z_q_st
+                residuals = residuals - z_q.detach()
+            reconstructed = self.decoder(z_sum)
+            return all_indices, reconstructed, z, z_sum
+
+        # ---- Flow matching integration API ----
+
+        def encode_continuous(self, invariants: torch.Tensor) -> torch.Tensor:
+            """
+            Encode invariants to continuous latent space (before quantization).
+
+            Args:
+                invariants: Input invariant sequences (B, T, C).
+
+            Returns:
+                Continuous latent z (B, T, latent_dim).
+            """
+            return self.encoder(invariants)
+
+        def decode_from_latent(self, z: torch.Tensor) -> torch.Tensor:
+            """
+            Decode from continuous latent to invariants.
+
+            Args:
+                z: Continuous latent (B, T, latent_dim).
+
+            Returns:
+                Reconstructed invariants (B, T, invariant_dim).
+            """
+            return self.decoder(z)
+
+        def quantize(self, z: torch.Tensor, num_codebooks: int = None) -> tuple:
+            """
+            Quantize continuous latent using RVQ.
+
+            Args:
+                z: Continuous latent (B, T, latent_dim).
+                num_codebooks: Number of codebooks to use (default: all).
+
+            Returns:
+                Tuple of (all_indices, z_sum) where all_indices is list of (B, T).
+            """
+            if num_codebooks is None:
+                num_codebooks = self.num_codebooks
+
+            residuals = z
+            all_indices = []
+            z_sum = torch.zeros_like(z)
+
+            for i, vq in enumerate(self.vqs[:num_codebooks]):
+                indices, z_q_st, z_q = vq(residuals)
+                all_indices.append(indices)
+                z_sum = z_sum + z_q_st
+                residuals = residuals - z_q.detach()
+
+            return all_indices, z_sum
+
+        def encode_partial(
+            self,
+            invariants: torch.Tensor,
+            num_codebooks: int,
+        ) -> tuple:
+            """
+            Encode with partial RVQ (for hierarchical VFM).
+
+            Uses only the first num_codebooks codebooks.
+
+            Args:
+                invariants: Input invariant sequences (B, T, C).
+                num_codebooks: Number of codebooks to use.
+
+            Returns:
+                Tuple of (all_indices, z_sum, reconstructed).
+            """
+            z = self.encoder(invariants)
+            all_indices, z_sum = self.quantize(z, num_codebooks)
+            reconstructed = self.decoder(z_sum)
+            return all_indices, z_sum, reconstructed
+
+        def get_codebook_embeddings(self, codebook_idx: int = 0) -> torch.Tensor:
+            """
+            Get codebook embeddings for a specific codebook.
+
+            Args:
+                codebook_idx: Index of the codebook (0 to num_codebooks-1).
+
+            Returns:
+                Codebook embeddings (codebook_size, latent_dim).
+            """
+            return self.vqs[codebook_idx].embedding.weight.data
+
+        def get_all_codebook_embeddings(self) -> list:
+            """
+            Get all codebook embeddings.
+
+            Returns:
+                List of codebook embeddings, each (codebook_size, latent_dim).
+            """
+            return [vq.embedding.weight.data for vq in self.vqs]
+
+        def embed_tokens(self, all_indices: list) -> torch.Tensor:
+            """
+            Convert RVQ token indices to summed embeddings.
+
+            Args:
+                all_indices: List of token indices, each (B, T).
+
+            Returns:
+                Summed embeddings (B, T, latent_dim).
+            """
+            z_sum = None
+            for i, (indices, vq) in enumerate(zip(all_indices, self.vqs)):
+                z_q = vq.embedding(indices)
+                if z_sum is None:
+                    z_sum = z_q
+                else:
+                    z_sum = z_sum + z_q
+            return z_sum
+
+        def decode_tokens(self, all_indices: list) -> torch.Tensor:
+            """
+            Decode RVQ token indices to invariants.
+
+            Args:
+                all_indices: List of token indices, each (B, T).
+
+            Returns:
+                Reconstructed invariants (B, T, invariant_dim).
+            """
+            z_sum = self.embed_tokens(all_indices)
+            return self.decoder(z_sum)
+
+else:
+    ResidualVQTokenizer = None

dhb_xr/tokenization/vqvae.py

@@ -0,0 +1,155 @@
+"""VQ-VAE tokenizer for DHB invariant sequences (DHB-Token)."""
+
+try:
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    from dhb_xr.tokenization.causal_encoder import CausalConv1dEncoder
+    HAS_TORCH = True
+except ImportError:
+    HAS_TORCH = False
+
+if HAS_TORCH:
+
+    class VectorQuantizer(nn.Module):
+        def __init__(self, num_embeddings: int, embedding_dim: int):
+            super().__init__()
+            self.embedding = nn.Embedding(num_embeddings, embedding_dim)
+            self.embedding.weight.data.uniform_(-1 / num_embeddings, 1 / num_embeddings)
+
+        def forward(self, z: torch.Tensor) -> tuple:
+            # z: (B, T, D)
+            B, T, D = z.shape
+            z_flat = z.reshape(-1, D)
+            d = torch.cdist(z_flat, self.embedding.weight)
+            indices = d.argmin(dim=1)
+            z_q = self.embedding(indices).reshape(B, T, D)
+            z_q_st = z + (z_q - z).detach()
+            return indices.reshape(B, T), z_q_st, z_q
+
+    class DHBTokenizer(nn.Module):
+        """
+        Causal VQ-VAE for invariant sequences.
+        invariants (B, T, C) -> tokens (B, T), reconstructed (B, T, C).
+        """
+
+        def __init__(
+            self,
+            invariant_dim: int,
+            latent_dim: int,
+            codebook_size: int,
+            num_layers: int = 2,
+            kernel_size: int = 3,
+        ):
+            super().__init__()
+            self.encoder = CausalConv1dEncoder(
+                invariant_dim, latent_dim, latent_dim, num_layers, kernel_size
+            )
+            self.vq = VectorQuantizer(codebook_size, latent_dim)
+            self.decoder = CausalConv1dEncoder(
+                latent_dim, latent_dim, invariant_dim, num_layers, kernel_size
+            )
+            self.invariant_dim = invariant_dim
+            self.latent_dim = latent_dim
+            self.codebook_size = codebook_size
+
+        def forward(self, invariants: torch.Tensor) -> tuple:
+            z = self.encoder(invariants)
+            indices, z_q_st, z_q = self.vq(z)
+            reconstructed = self.decoder(z_q_st)
+            return indices, reconstructed, z, z_q
+
+        def loss(
+            self,
+            invariants: torch.Tensor,
+            reconstructed: torch.Tensor,
+            z: torch.Tensor,
+            z_q: torch.Tensor,
+            beta: float = 0.25,
+        ) -> torch.Tensor:
+            rec_loss = F.mse_loss(reconstructed, invariants)
+            commitment = F.mse_loss(z, z_q)
+            codebook = F.mse_loss(z_q, z.detach())
+            return rec_loss + beta * commitment + codebook
+
+        # ---- Flow matching integration API ----
+
+        def encode_continuous(self, invariants: torch.Tensor) -> torch.Tensor:
+            """
+            Encode invariants to continuous latent space (before quantization).
+
+            This is useful for flow matching which operates in continuous space.
+
+            Args:
+                invariants: Input invariant sequences (B, T, C).
+
+            Returns:
+                Continuous latent z (B, T, latent_dim).
+            """
+            return self.encoder(invariants)
+
+        def decode_from_latent(self, z: torch.Tensor) -> torch.Tensor:
+            """
+            Decode from continuous latent to invariants.
+
+            Bypasses the VQ step, useful for flow matching generation.
+
+            Args:
+                z: Continuous latent (B, T, latent_dim).
+
+            Returns:
+                Reconstructed invariants (B, T, invariant_dim).
+            """
+            return self.decoder(z)
+
+        def quantize(self, z: torch.Tensor) -> tuple:
+            """
+            Quantize continuous latent to discrete tokens.
+
+            Args:
+                z: Continuous latent (B, T, latent_dim).
+
+            Returns:
+                Tuple of (indices, z_q_st, z_q).
+            """
+            return self.vq(z)
+
+        def get_codebook_embeddings(self) -> torch.Tensor:
+            """
+            Get the VQ codebook embeddings.
+
+            Useful for flow matching in embedding space or visualization.
+
+            Returns:
+                Codebook embeddings (codebook_size, latent_dim).
+            """
+            return self.vq.embedding.weight.data
+
+        def embed_tokens(self, indices: torch.Tensor) -> torch.Tensor:
+            """
+            Convert token indices to embeddings.
+
+            Args:
+                indices: Token indices (B, T).
+
+            Returns:
+                Token embeddings (B, T, latent_dim).
+            """
+            return self.vq.embedding(indices)
+
+        def decode_tokens(self, indices: torch.Tensor) -> torch.Tensor:
+            """
+            Decode token indices to invariants.
+
+            Args:
+                indices: Token indices (B, T).
+
+            Returns:
+                Reconstructed invariants (B, T, invariant_dim).
+            """
+            z_q = self.embed_tokens(indices)
+            return self.decoder(z_q)
+
+else:
+    DHBTokenizer = None
+    VectorQuantizer = None