A Hybrid Field–Graph Representation for Probabilistic World Models in Robotics

class EncodingModel(nn.Module):
    def __init__(self, latent_dim=128, max_objects=8):
        super().__init__()
        sd = 128
        self.max_objects = max_objects
        self.backbone = nn.Sequential(
            nn.Conv2d(3, 32, 7, 2, 3),
            nn.ReLU(),
            nn.Conv2d(32, 64, 5, 2, 2),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, 2, 1),
            nn.ReLU(),
            nn.Conv2d(128, 256, 3, 2, 1),
            nn.ReLU(),
        )
        self.slots_init = nn.Parameter(torch.randn(1, max_objects, sd) * 0.02)
        self.sq = nn.Linear(sd, 64)
        self.sk = nn.Linear(256, 64)
        self.sv = nn.Linear(256, sd)
        self.sgru = nn.GRUCell(sd, sd)
        self.smlp = nn.Sequential(nn.Linear(sd, sd), nn.ReLU(), nn.Linear(sd, sd))
        self.to_z = nn.Linear(sd, latent_dim)
        self.to_pose = nn.Linear(sd, 7)
        self.to_exist = nn.Linear(sd, 1)
        self.edge_net = nn.Sequential(nn.Linear(sd * 2, 64), nn.ReLU(), nn.Linear(64, 33))

    def forward(self, image):
        B, N = image.shape[0], self.max_objects
        f = self.backbone(image).flatten(2).permute(0, 2, 1)
        slots = self.slots_init.expand(B, -1, -1)

        for _ in range(3):
            q, k, v = self.sq(slots), self.sk(f), self.sv(f)
            a = F.softmax(torch.einsum("bnd,bmd->bnm", q, k) / 8, dim=1)
            u = torch.einsum("bnm,bmd->bnd", a, v)
            slots = self.sgru(u.reshape(B * N, -1), slots.reshape(B * N, -1)).reshape(B, N, -1)
            slots = slots + self.smlp(slots)

        z = self.to_z(slots)
        poses = self.to_pose(slots)
        exist = torch.sigmoid(self.to_exist(slots).squeeze(-1))
        si = slots.unsqueeze(2).expand(-1, -1, N, -1)
        sj = slots.unsqueeze(1).expand(-1, N, -1, -1)
        eo = self.edge_net(torch.cat([si, sj], -1))

        return {
            "phi": {"z": z, "poses": poses, "existence": exist},
            "graph": {
                "node_features": slots,
                "edge_features": eo[..., :32],
                "contact_probs": torch.sigmoid(eo[..., 32]),
            },
        }

class GNNLayer(nn.Module):
    def __init__(self, hd, ed):
        super().__init__()
        self.msg = nn.Sequential(nn.Linear(hd * 2 + ed, hd), nn.ReLU(), nn.Linear(hd, hd))
        self.upd = nn.Sequential(nn.Linear(hd * 2, hd), nn.ReLU(), nn.Linear(hd, hd))
        self.norm = nn.LayerNorm(hd)

    def forward(self, h, ef, cp, mask):
        hi = h.unsqueeze(2).expand(-1, -1, h.shape[1], -1)
        hj = h.unsqueeze(1).expand(-1, h.shape[1], -1, -1)
        m = self.msg(torch.cat([hi, hj, ef], -1)) * cp.unsqueeze(-1) * mask.unsqueeze(1)
        return self.norm(h + self.upd(torch.cat([h, m.sum(2)], -1))) * mask


class StateVectorizer(nn.Module):
    def __init__(self, latent_dim=128, output_dim=512):
        super().__init__()
        hd = 256
        self.node_emb = nn.Sequential(nn.Linear(latent_dim + 8, hd), nn.ReLU(), nn.Linear(hd, hd))
        self.gnns = nn.ModuleList([GNNLayer(hd, 32) for _ in range(3)])
        self.aq = nn.Linear(hd, 64)
        self.ak = nn.Linear(hd, 64)
        self.proj = nn.Sequential(nn.Linear(hd, hd), nn.ReLU(), nn.Linear(hd, output_dim))
        self.norm = nn.LayerNorm(output_dim)

    def forward(self, phi, graph):
        h = self.node_emb(
            torch.cat([phi["z"], phi["poses"], phi["existence"].unsqueeze(-1)], -1)
        )
        mask = phi["existence"].unsqueeze(-1)
        h = h * mask

        for g in self.gnns:
            h = g(h, graph["edge_features"], graph["contact_probs"], mask)

        q = self.aq(h.mean(1, keepdim=True))
        k = self.ak(h)
        a = F.softmax(torch.einsum("bid,bjd->bij", q, k) / 8 + (1 - mask.transpose(1, 2)) * -1e9, -1)
        pooled = torch.einsum("bij,bjd->bid", a, h).squeeze(1)
        return self.norm(self.proj(pooled))

class CrossEncoder(nn.Module):
    def __init__(self, state_dim=512, embed_dim=512, action_dim=6):
        super().__init__()
        self.scene_proj = nn.Sequential(nn.Linear(state_dim, embed_dim), nn.ReLU(), nn.Linear(embed_dim, embed_dim))
        self.text_emb = nn.Embedding(10000, 256)
        self.text_tf = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(256, 4, 512, batch_first=True), 2
        )
        self.img_enc = nn.Sequential(
            nn.Conv2d(3, 32, 7, 4, 3),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, 2, 1),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, 2, 1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
        )
        self.cross_attn = nn.MultiheadAttention(embed_dim, 8, batch_first=True)
        self.action_head = nn.Sequential(
            nn.Linear(embed_dim, 256), nn.ReLU(),
            nn.Linear(256, 64), nn.ReLU(),
            nn.Linear(64, action_dim), nn.Tanh(),
        )

    def forward(self, s, text_tokens=None, text_mask=None, image=None):
        se = self.scene_proj(s)
        te = self.encode_text(text_tokens, text_mask) if text_tokens is not None else None
        ie = self.encode_image(image) if image is not None else None

        toks = [se.unsqueeze(1)]
        if te is not None:
            toks.append(te.unsqueeze(1))
        if ie is not None:
            toks.append(ie.unsqueeze(1))

        seq = torch.cat(toks, 1)
        fused, _ = self.cross_attn(seq, seq, seq)
        return {
            "action": self.action_head(fused[:, 0]),
            "scene_emb": se,
            "text_emb": te,
            "image_emb": ie,
        }

class FullPipeline(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = EncodingModel(128, 8)
        self.vectorizer = StateVectorizer(128, 512)
        self.cross_encoder = CrossEncoder(512, 512, 6)

    def forward(self, image, text_tokens=None, text_mask=None):
        pg = self.encoder(image)
        s = self.vectorizer(pg["phi"], pg["graph"])
        out = self.cross_encoder(s, text_tokens, text_mask, image)
        out["phi"] = pg["phi"]
        out["graph"] = pg["graph"]
        out["state_vector"] = s
        return out

Complete Field Specification: Φ(x, d̂, t | z_i)
================================================

The belief field is a neural function that maps a query (position, direction, time)
to DISTRIBUTIONS over every physical quantity relevant for manipulation.

Design Principles:
  1. Position x ∈ R³ is always required
  2. Direction d̂ ∈ S² is required for anisotropic/directional properties
  3. Time t is required for dynamic properties
  4. Object latent z_i conditions the field per-object
  5. Every output is DISTRIBUTIONAL (mean + uncertainty)
  6. Outputs span scalars, vectors, and tensors — matching physics

This is the NeRF paradigm extended from appearance to full multiphysics:
  NeRF:  Φ(x, d̂) → (color, density)         [appearance only]
  Ours:  Φ(x, d̂, t | z) → (geometry, material, dynamics, semantics)

Architecture:
  The network has a SHARED TRUNK that processes (x, t, z_i) into a feature vector,
  then BRANCHES for different output groups. The direction d̂ enters ONLY for
  branches that need it (anisotropic material, appearance) — following the NeRF
  insight that geometry shouldn't depend on view direction.

  ┌─────────────┐
  │ x ∈ R³      │──→ Fourier Encoding ──→┐
  │ t ∈ R       │──→ Time Encoding   ──→ │
  │ z_i ∈ R^D   │──────────────────────→ ├──→ Shared Trunk (MLP) ──→ h ∈ R^H
  └─────────────┘                        │
                                         │
  ┌──────────── h ──────────────────────┐│
  │                                     ││
  │  ┌─── Geometry Branch ←── h        ││
  │  │    (position-only)               ││
  │  │    → SDF:     μ_d, σ_d          ││
  │  │    → Normal:  n̂ = ∇μ_d / |∇μ_d| ││  (computed via autograd, not predicted)
  │  │    → Curvature: κ = ∇²d         ││  (second derivative of SDF)
  │  │                                  ││
  │  ┌─── Isotropic Material Branch ← h││
  │  │    (position-only)               ││
  │  │    → Density:      μ_ρ, σ_ρ     ││  LogNormal
  │  │    → Temperature:  μ_T, σ_T     ││  Normal (or LogNormal for Kelvin)
  │  │    → Restitution:  α_e, β_e     ││  Beta[0,1]
  │  │    → Poisson ratio: α_ν, β_ν   ││  Beta[0,0.5]
  │  │                                  ││
  │  ┌─── Directional Material Branch ←─┤│← d̂ ∈ S² (direction enters here)
  │  │    (position + direction)        ││
  │  │    → Friction:    μ_f(d̂), σ_f   ││  Friction along sliding direction d̂
  │  │    → Stiffness:   E(d̂), σ_E     ││  Young's modulus along d̂
  │  │    → Thermal cond: κ(d̂), σ_κ    ││  Conductivity along d̂
  │  │                                  ││
  │  ┌─── Dynamics Branch ← h          ││
  │  │    (position-only, vector output)││
  │  │    → Velocity:  v ∈ R³           ││
  │  │    → Stress:    σ_ij (6 indep)   ││  Symmetric 3x3 → Voigt notation
  │  │    → Strain:    ε_ij (6 indep)   ││  
  │  │                                  ││
  │  ┌─── Appearance Branch ←───────────┤│← d̂ (view direction, as in NeRF)
  │  │    → Color:     c ∈ R³           ││
  │  │    → Opacity:   α ∈ [0,1]        ││
  │  │                                  ││
  │  ┌─── Semantic Branch ← h          ││
  │  │    → Class logits: R^C           ││
  │  │    → Affordance: R^A             ││
  │  └─────────────────────────────────┘│
  └─────────────────────────────────────┘
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Dict, Optional, Tuple


# ============================================================
# Encodings
# ============================================================

class FourierEncoding(nn.Module):
    """Sinusoidal positional encoding. Maps R^n → R^(n + 2nL)."""
    def __init__(self, input_dim: int = 3, n_freqs: int = 6, include_input: bool = True):
        super().__init__()
        self.input_dim = input_dim
        self.n_freqs = n_freqs
        self.include_input = include_input
        freqs = 2.0 ** torch.linspace(0, n_freqs - 1, n_freqs)
        self.register_buffer('freqs', freqs)

    @property
    def out_dim(self):
        return self.input_dim * (1 + 2 * self.n_freqs) if self.include_input else self.input_dim * 2 * self.n_freqs

    def forward(self, x):
        # x: [..., input_dim]
        x_proj = x.unsqueeze(-1) * self.freqs  # [..., input_dim, L]
        enc = torch.cat([x_proj.sin(), x_proj.cos()], dim=-1)  # [..., input_dim, 2L]
        enc = enc.reshape(*x.shape[:-1], -1)  # [..., input_dim * 2L]
        return torch.cat([x, enc], dim=-1) if self.include_input else enc


class DirectionEncoding(nn.Module):
    """Encode unit direction vector d̂ ∈ S².
    
    Can be input as:
      - Unit vector (dx, dy, dz) with |d|=1
      - Spherical angles (θ, φ)
    
    Uses spherical harmonics basis (lower frequency than position encoding
    since directional variation is typically smoother).
    """
    def __init__(self, n_freqs: int = 4):
        super().__init__()
        self.pos_enc = FourierEncoding(input_dim=3, n_freqs=n_freqs)

    @property
    def out_dim(self):
        return self.pos_enc.out_dim

    def forward(self, d: torch.Tensor):
        """d: [..., 3] unit vectors OR [..., 2] (theta, phi)"""
        if d.shape[-1] == 2:
            # Convert spherical to cartesian
            theta, phi = d[..., 0], d[..., 1]
            dx = torch.sin(theta) * torch.cos(phi)
            dy = torch.sin(theta) * torch.sin(phi)
            dz = torch.cos(theta)
            d = torch.stack([dx, dy, dz], dim=-1)
        # Normalize
        d = F.normalize(d, dim=-1)
        return self.pos_enc(d)


# ============================================================
# Distribution Heads (proper probabilistic outputs)
# ============================================================

class GaussianHead(nn.Module):
    """Predicts Gaussian(μ, σ) for unbounded scalars."""
    def __init__(self, in_dim: int, out_dim: int = 1):
        super().__init__()
        self.mu = nn.Linear(in_dim, out_dim)
        self.log_sigma = nn.Linear(in_dim, out_dim)

    def forward(self, h):
        return self.mu(h), F.softplus(self.log_sigma(h)) + 1e-4


class LogNormalHead(nn.Module):
    """Predicts LogNormal for positive scalars (density, stiffness, etc.)."""
    def __init__(self, in_dim: int, out_dim: int = 1, scale: float = 1.0):
        super().__init__()
        self.log_mu = nn.Linear(in_dim, out_dim)
        self.log_sigma = nn.Linear(in_dim, out_dim)
        self.scale = scale

    def forward(self, h):
        log_mu = self.log_mu(h)
        log_sigma = F.softplus(self.log_sigma(h)) + 1e-4
        # Mean of LogNormal = exp(log_mu + log_sigma²/2)
        mu = torch.exp(log_mu + 0.5 * log_sigma ** 2) * self.scale
        # Approximate sigma in real space
        sigma = mu * torch.sqrt(torch.exp(log_sigma ** 2) - 1)
        return mu, sigma


class BetaHead(nn.Module):
    """Predicts Beta(α, β) for [0, upper_bound] quantities."""
    def __init__(self, in_dim: int, out_dim: int = 1, upper_bound: float = 1.0):
        super().__init__()
        self.alpha = nn.Linear(in_dim, out_dim)
        self.beta = nn.Linear(in_dim, out_dim)
        self.upper_bound = upper_bound

    def forward(self, h):
        alpha = F.softplus(self.alpha(h)) + 1.01  # > 1 for unimodal
        beta = F.softplus(self.beta(h)) + 1.01
        mu = alpha / (alpha + beta) * self.upper_bound
        var = alpha * beta / ((alpha + beta) ** 2 * (alpha + beta + 1))
        sigma = torch.sqrt(var) * self.upper_bound
        return mu, sigma


class VectorHead(nn.Module):
    """Predicts a 3D vector with per-component uncertainty."""
    def __init__(self, in_dim: int):
        super().__init__()
        self.mu = nn.Linear(in_dim, 3)
        self.log_sigma = nn.Linear(in_dim, 3)

    def forward(self, h):
        return self.mu(h), F.softplus(self.log_sigma(h)) + 1e-4


class SymmetricTensorHead(nn.Module):
    """Predicts a 3×3 symmetric tensor via 6 Voigt components.
    
    Voigt ordering: [σ_xx, σ_yy, σ_zz, σ_yz, σ_xz, σ_xy]
    Ensures symmetry by construction.
    """
    def __init__(self, in_dim: int):
        super().__init__()
        self.voigt = nn.Linear(in_dim, 6)  # 6 independent components
        self.log_sigma = nn.Linear(in_dim, 6)

    def forward(self, h):
        v = self.voigt(h)
        s = F.softplus(self.log_sigma(h)) + 1e-4
        return v, s

    @staticmethod
    def voigt_to_matrix(voigt):
        """Convert [B, 6] Voigt to [B, 3, 3] symmetric matrix."""
        B = voigt.shape[0]
        mat = torch.zeros(B, 3, 3, device=voigt.device)
        mat[:, 0, 0] = voigt[:, 0]  # xx
        mat[:, 1, 1] = voigt[:, 1]  # yy
        mat[:, 2, 2] = voigt[:, 2]  # zz
        mat[:, 1, 2] = mat[:, 2, 1] = voigt[:, 3]  # yz
        mat[:, 0, 2] = mat[:, 2, 0] = voigt[:, 4]  # xz
        mat[:, 0, 1] = mat[:, 1, 0] = voigt[:, 5]  # xy
        return mat


# ============================================================
# The Complete Object Field
# ============================================================

class ObjectFieldComplete(nn.Module):
    """Complete per-object neural field.
    
    Φ(x, d̂, t | z_i) → full multiphysics state
    
    Architecture follows the NeRF insight:
      - Shared trunk processes (x, t, z_i) → features h
      - Geometry branch uses h only (no direction dependence for SDF)
      - Directional branches take h + encoded d̂
      - This ensures SDF is view/direction invariant
    
    Args:
        latent_dim:  dimension of object latent z_i
        hidden_dim:  width of trunk MLP
        n_layers:    depth of trunk
        pos_freqs:   Fourier frequencies for position encoding
        dir_freqs:   Fourier frequencies for direction encoding
        n_semantic_classes: number of semantic categories
        skip_at:     layer index for skip connection
    """

    def __init__(
        self,
        latent_dim: int = 128,
        hidden_dim: int = 256,
        n_layers: int = 6,
        pos_freqs: int = 6,
        dir_freqs: int = 4,
        n_semantic_classes: int = 32,
        skip_at: int = 3,
    ):
        super().__init__()
        self.latent_dim = latent_dim
        self.hidden_dim = hidden_dim
        self.skip_at = skip_at

        # --- Encodings ---
        self.pos_enc = FourierEncoding(input_dim=3, n_freqs=pos_freqs)
        self.dir_enc = DirectionEncoding(n_freqs=dir_freqs)
        self.time_enc = FourierEncoding(input_dim=1, n_freqs=4)

        trunk_input_dim = self.pos_enc.out_dim + self.time_enc.out_dim + latent_dim

        # --- Shared Trunk ---
        trunk_layers = []
        dims = [trunk_input_dim] + [hidden_dim] * n_layers
        for i in range(n_layers):
            in_d = dims[i] + (trunk_input_dim if i == skip_at else 0)
            trunk_layers.append(nn.Linear(in_d, hidden_dim))
        self.trunk = nn.ModuleList(trunk_layers)

        # --- Geometry Branch (position-only) ---
        self.sdf_head = GaussianHead(hidden_dim)
        # Note: surface normal is computed as ∇SDF via autograd, not predicted
        # Curvature can be computed as ∇²SDF (Laplacian)

        # --- Isotropic Material Branch (position-only) ---
        self.density_head = LogNormalHead(hidden_dim, scale=1000.0)    # kg/m³
        self.temperature_head = GaussianHead(hidden_dim)                # Kelvin
        self.restitution_head = BetaHead(hidden_dim, upper_bound=1.0)   # [0, 1]
        self.poisson_head = BetaHead(hidden_dim, upper_bound=0.5)       # [0, 0.5]
        self.acoustic_damping_head = LogNormalHead(hidden_dim, scale=1.0)

        # --- Directional Material Branch (position + direction) ---
        dir_input_dim = hidden_dim + self.dir_enc.out_dim
        self.dir_material_trunk = nn.Sequential(
            nn.Linear(dir_input_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, hidden_dim // 2),
            nn.ReLU(),
        )
        dir_h_dim = hidden_dim // 2
        self.friction_head = LogNormalHead(dir_h_dim, scale=1.0)        # μ ∈ R+ (typ 0.1-1.0)
        self.stiffness_head = LogNormalHead(dir_h_dim, scale=1e9)       # Pa
        self.thermal_cond_head = LogNormalHead(dir_h_dim, scale=1.0)    # W/(m·K)

        # --- Dynamics Branch (position-only, vector/tensor outputs) ---
        self.velocity_head = VectorHead(hidden_dim)                      # m/s
        self.stress_head = SymmetricTensorHead(hidden_dim)               # Pa (6 Voigt components)
        self.strain_head = SymmetricTensorHead(hidden_dim)               # dimensionless (6 Voigt)

        # --- Appearance Branch (position + direction, as in NeRF) ---
        app_input_dim = hidden_dim + self.dir_enc.out_dim
        self.appearance_trunk = nn.Sequential(
            nn.Linear(app_input_dim, hidden_dim // 2),
            nn.ReLU(),
        )
        self.color_head = nn.Linear(hidden_dim // 2, 3)     # RGB ∈ [0,1]
        self.opacity_head = nn.Linear(hidden_dim // 2, 1)    # α ∈ [0,1]

        # --- Semantic Branch ---
        self.semantic_head = nn.Linear(hidden_dim, n_semantic_classes)
        self.affordance_head = nn.Linear(hidden_dim, 8)  # graspable, pushable, pourable, etc.

    def trunk_forward(self, x: torch.Tensor, t: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        """Shared trunk: (x, t, z) → feature vector h."""
        x_enc = self.pos_enc(x)
        t_enc = self.time_enc(t.unsqueeze(-1) if t.dim() == 1 else t)
        if z.dim() == 1:
            z = z.unsqueeze(0).expand(x.shape[0], -1)

        h = torch.cat([x_enc, t_enc, z], dim=-1)
        h_input = h

        for i, layer in enumerate(self.trunk):
            if i == self.skip_at:
                h = torch.cat([h, h_input], dim=-1)
            h = F.relu(layer(h))

        return h

    def forward(
        self,
        x: torch.Tensor,                      # [B, 3] position in object-local frame
        z: torch.Tensor,                       # [D] or [B, D] object latent
        t: Optional[torch.Tensor] = None,      # [B] or [B, 1] time
        d: Optional[torch.Tensor] = None,      # [B, 3] or [B, 2] query direction
        compute_normals: bool = False,          # whether to compute ∇SDF
        compute_curvature: bool = False,        # whether to compute ∇²SDF
    ) -> Dict[str, torch.Tensor]:
        """
        Full field query.
        
        Returns dict with all physical quantities as (mean, uncertainty) pairs.
        """
        B = x.shape[0]

        if t is None:
            t = torch.zeros(B, 1, device=x.device)
        elif t.dim() == 1:
            t = t.unsqueeze(-1)

        # Enable gradient tracking for normal/curvature computation
        if compute_normals or compute_curvature:
            x = x.detach().requires_grad_(True)

        # --- Shared trunk ---
        h = self.trunk_forward(x, t, z)

        result = {}

        # --- Geometry (position-only) ---
        sdf_mu, sdf_sigma = self.sdf_head(h)
        result['sdf_mu'] = sdf_mu
        result['sdf_sigma'] = sdf_sigma

        # Surface normal via autograd
        if compute_normals:
            grad_sdf = torch.autograd.grad(
                sdf_mu.sum(), x, create_graph=compute_curvature, retain_graph=True
            )[0]  # [B, 3]
            normal = F.normalize(grad_sdf, dim=-1)
            result['normal'] = normal
            result['sdf_gradient'] = grad_sdf

            if compute_curvature:
                # Laplacian of SDF ≈ mean curvature
                curvature = torch.zeros(B, device=x.device)
                for dim in range(3):
                    grad2 = torch.autograd.grad(
                        grad_sdf[:, dim].sum(), x, retain_graph=True
                    )[0][:, dim]
                    curvature += grad2
                result['mean_curvature'] = curvature

        # --- Isotropic materials (position-only) ---
        result['density_mu'], result['density_sigma'] = self.density_head(h)
        result['temperature_mu'], result['temperature_sigma'] = self.temperature_head(h)
        result['restitution_mu'], result['restitution_sigma'] = self.restitution_head(h)
        result['poisson_mu'], result['poisson_sigma'] = self.poisson_head(h)
        result['damping_mu'], result['damping_sigma'] = self.acoustic_damping_head(h)

        # --- Directional materials (position + direction) ---
        if d is not None:
            d_enc = self.dir_enc(d)
            h_dir = self.dir_material_trunk(torch.cat([h, d_enc], dim=-1))
            result['friction_mu'], result['friction_sigma'] = self.friction_head(h_dir)
            result['stiffness_mu'], result['stiffness_sigma'] = self.stiffness_head(h_dir)
            result['thermal_cond_mu'], result['thermal_cond_sigma'] = self.thermal_cond_head(h_dir)
        else:
            # If no direction given, query with canonical directions and average
            # (gives isotropic estimate)
            canonical = torch.tensor([[1,0,0],[0,1,0],[0,0,1]], dtype=torch.float, device=x.device)
            fric_sum = torch.zeros(B, 1, device=x.device)
            stiff_sum = torch.zeros(B, 1, device=x.device)
            for cd in canonical:
                cd_batch = cd.unsqueeze(0).expand(B, -1)
                d_enc = self.dir_enc(cd_batch)
                h_dir = self.dir_material_trunk(torch.cat([h, d_enc], dim=-1))
                f_mu, _ = self.friction_head(h_dir)
                s_mu, _ = self.stiffness_head(h_dir)
                fric_sum += f_mu
                stiff_sum += s_mu
            result['friction_mu'] = fric_sum / 3
            result['stiffness_mu'] = stiff_sum / 3
            # Higher uncertainty for averaged isotropic estimate
            result['friction_sigma'] = torch.ones_like(result['friction_mu']) * 0.1
            result['stiffness_sigma'] = result['stiffness_mu'] * 0.2

        # --- Dynamics ---
        result['velocity_mu'], result['velocity_sigma'] = self.velocity_head(h)
        stress_v, stress_s = self.stress_head(h)
        result['stress_voigt'] = stress_v        # [B, 6]
        result['stress_sigma'] = stress_s
        strain_v, strain_s = self.strain_head(h)
        result['strain_voigt'] = strain_v
        result['strain_sigma'] = strain_s

        # --- Appearance (needs direction) ---
        if d is not None:
            d_enc = self.dir_enc(d)
            h_app = self.appearance_trunk(torch.cat([h, d_enc], dim=-1))
            result['color'] = torch.sigmoid(self.color_head(h_app))
            result['opacity'] = torch.sigmoid(self.opacity_head(h_app))

        # --- Semantics ---
        result['semantic_logits'] = self.semantic_head(h)
        result['affordance_logits'] = self.affordance_head(h)

        # --- Store features for gating ---
        result['features'] = h

        return result


# ============================================================
# Convenience: Derived Quantities
# ============================================================

def contact_response(field_output: Dict, contact_force: torch.Tensor, contact_direction: torch.Tensor):
    """Given field output at a contact point, predict contact response.
    
    Args:
        field_output: from ObjectFieldComplete.forward(x, z, d=contact_normal)
        contact_force: [B, 1] applied normal force
        contact_direction: [B, 3] contact normal
    
    Returns:
        dict with:
            - max_static_friction: F_max = μ_s * F_normal
            - deformation: δ ≈ F / (E * A) for small deformations
            - will_slip: bool, whether tangential force exceeds friction
    """
    mu = field_output['friction_mu']
    E = field_output['stiffness_mu']

    max_friction = mu * contact_force.abs()
    # Hertzian contact approximation: deformation ~ (F / E)^(2/3)
    deformation = (contact_force.abs() / (E + 1e-6)) ** 0.667

    return {
        'max_static_friction': max_friction,
        'deformation': deformation,
        'friction': mu,
        'stiffness': E,
    }


# ============================================================
# Test
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Complete Object Field — Architecture Test")
    print("=" * 60)

    field = ObjectFieldComplete(
        latent_dim=128,
        hidden_dim=256,
        n_layers=6,
        pos_freqs=6,
        dir_freqs=4,
    )

    params = sum(p.numel() for p in field.parameters())
    print(f"\nTotal parameters: {params:,}")

    # Itemize by component
    components = {
        'trunk': sum(p.numel() for l in field.trunk for p in l.parameters()),
        'geometry': sum(p.numel() for p in field.sdf_head.parameters()),
        'isotropic_material': sum(
            sum(p.numel() for p in head.parameters())
            for head in [field.density_head, field.temperature_head,
                         field.restitution_head, field.poisson_head, field.acoustic_damping_head]
        ),
        'directional_material': sum(p.numel() for p in field.dir_material_trunk.parameters()) +
                                sum(p.numel() for p in field.friction_head.parameters()) +
                                sum(p.numel() for p in field.stiffness_head.parameters()) +
                                sum(p.numel() for p in field.thermal_cond_head.parameters()),
        'dynamics': sum(p.numel() for p in field.velocity_head.parameters()) +
                    sum(p.numel() for p in field.stress_head.parameters()) +
                    sum(p.numel() for p in field.strain_head.parameters()),
        'appearance': sum(p.numel() for p in field.appearance_trunk.parameters()) +
                      sum(p.numel() for p in field.color_head.parameters()) +
                      sum(p.numel() for p in field.opacity_head.parameters()),
        'semantic': sum(p.numel() for p in field.semantic_head.parameters()) +
                    sum(p.numel() for p in field.affordance_head.parameters()),
    }
    for name, count in components.items():
        pct = count / params * 100
        print(f"  {name:25s} {count:>8,} ({pct:5.1f}%)")

    # Forward pass
    B = 64
    x = torch.randn(B, 3) * 0.05
    z = torch.randn(128)
    t = torch.zeros(B)
    d = F.normalize(torch.randn(B, 3), dim=-1)

    print(f"\nQuery: {B} points, with direction, with normals+curvature")
    out = field(x, z, t=t, d=d, compute_normals=True, compute_curvature=True)

    print(f"\nOutputs:")
    for k, v in sorted(out.items()):
        if isinstance(v, torch.Tensor):
            print(f"  {k:25s} shape={str(list(v.shape)):12s} range=[{v.min().item():.4f}, {v.max().item():.4f}]")

    # Query without direction (isotropic average)
    print(f"\nQuery: {B} points, NO direction (isotropic estimate)")
    out_iso = field(x, z, t=t, d=None, compute_normals=True)
    print(f"  friction_mu (isotropic avg): {out_iso['friction_mu'].mean().item():.4f}")
    print(f"  normal (from ∇SDF):          shape={list(out_iso['normal'].shape)}")

    # Contact response
    print(f"\nContact response at query points:")
    F_contact = torch.ones(B, 1) * 5.0  # 5N
    response = contact_response(out, F_contact, d)
    for k, v in response.items():
        if isinstance(v, torch.Tensor):
            print(f"  {k:25s} mean={v.mean().item():.4f}")

    print("\nDone!")