miniworld/conditioning/poses.py

"""Camera pose -> ray-encoding utilities for pose-conditioned world model.

Adapted from `GeometryForcing/utils/geometry_utils.py` with the following
simplifications:
  * Only the bits needed for ``ray_encoding`` conditioning are kept (the
    variant the user picked as the best-performing one).
  * ``rays`` accepts independent ``(h_res, w_res)`` so non-square latents
    (e.g. 15x20) are supported without distorting the intrinsics.

All functions follow this convention:
  * Raw camera pose layout: ``(B, T, 16)`` = ``[K(4), R(9 + T(3))]`` where the
    first 4 columns are normalised intrinsics ``(fx, fy, px, py)`` (pixel-coords
    divided by image size) and the last 12 columns are a flattened ``3x4``
    world-to-camera extrinsics matrix in row-major.
  * Ray encoding output: ``(B, T, 180, H_lat, W_lat)`` (6 ray dims * 2 trig fns
    * 15 NeRF frequencies = 180). This matches what `DiT3DPose` consumes when
    ``conditioning_type=ray_encoding``.
"""

from __future__ import annotations

import math
from typing import Tuple

import torch
from einops import einsum, rearrange, repeat


def _split_pose16(raw_poses: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """``(B, T, 16)`` -> ``R (B, T, 3, 3)``, ``T (B, T, 3)``, ``K (B, T, 4)``."""
    assert raw_poses.shape[-1] == 16, f"expected 16-dim pose, got {raw_poses.shape[-1]}"
    K, RT = raw_poses.split([4, 12], dim=-1)
    RT = rearrange(RT, "b t (i j) -> b t i j", i=3, j=4)
    R = RT[..., :3, :3]
    T = RT[..., :3, 3]
    return R, T, K


def _normalize_by_first(R: torch.Tensor, T: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """Re-express all poses so the first frame is the world origin."""
    R_ref = R[:, 0]  # (B, 3, 3)
    T_ref = T[:, 0]  # (B, 3)
    R_inv = rearrange(R_ref, "b i j -> b j i")
    R_new = einsum(R, R_inv, "b t i j1, b j1 j2 -> b t i j2")
    T_new = T - einsum(R_new, T_ref, "b t i j, b j -> b t i")
    return R_new, T_new


def _compute_rays(
    R: torch.Tensor,
    T: torch.Tensor,
    K: torch.Tensor,
    h_res: int,
    w_res: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Per-pixel ray origin / direction in world coords.

    Args:
        R: ``(B, T, 3, 3)`` world->cam rotation.
        T: ``(B, T, 3)`` world->cam translation.
        K: ``(B, T, 4)`` normalised intrinsics ``(fx, fy, px, py)``.
        h_res, w_res: target ray grid resolution (independent so non-square
            latents are handled correctly).

    Returns:
        origin: ``(B, T, H, W, 3)``
        direction: ``(B, T, H, W, 3)`` (unnormalised; norm encodes depth scale)
    """
    device, dtype = K.device, K.dtype

    coord_w, coord_h = torch.meshgrid(
        torch.linspace(0, w_res - 1, w_res, device=device, dtype=dtype),
        torch.linspace(0, h_res - 1, h_res, device=device, dtype=dtype),
        indexing="xy",
    )  # (H, W) each
    coord_w = rearrange(coord_w, "h w -> 1 1 h w") + 0.5
    coord_h = rearrange(coord_h, "h w -> 1 1 h w") + 0.5

    # Normalised K -> pixel-space K (separate W / H scaling for non-square grids).
    fx = (K[..., 0] * w_res).view(*K.shape[:-1], 1, 1)  # (B, T, 1, 1)
    fy = (K[..., 1] * h_res).view(*K.shape[:-1], 1, 1)
    px = (K[..., 2] * w_res).view(*K.shape[:-1], 1, 1)
    py = (K[..., 3] * h_res).view(*K.shape[:-1], 1, 1)

    x = (coord_w - px) / fx
    y = (coord_h - py) / fy
    z = torch.ones_like(x)
    direction = torch.stack([x, y, z], dim=-1)  # (B, T, H, W, 3)

    R_inv = rearrange(R, "b t i j -> b t j i")
    direction = einsum(R_inv, direction, "b t i j, b t h w j -> b t h w i")

    origin = -einsum(R_inv, T, "b t i j, b t j -> b t i")
    origin = repeat(origin, "b t i -> b t h w i", h=h_res, w=w_res).clone()
    return origin, direction


def _normalize_translation_scale(T: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    """Per-clip translation-scale normalisation (CameraCtrl / lingbot style).

    Monocular SfM poses (e.g. RealEstate10K) have an arbitrary, per-clip metric
    scale, so raw camera translations vary wildly in magnitude across clips.
    Since the ray origin feeds NeRF frequency encoding ``sin(2^k pi x)`` -- which
    is very sensitive to the absolute magnitude of ``x`` -- this inconsistency
    hurts learning.  We rescale each clip so its largest camera displacement is
    ~1, making camera motion scale-invariant across clips.

    Args:
        T: ``(B, S, 3)`` camera translations, already expressed relative to the
            first frame (so the first frame sits at the origin).
        eps: guard so static / near-static clips (max norm ~ 0) are left
            unchanged ("only normalize when moving").

    Returns:
        ``(B, S, 3)`` translations divided by the per-clip max translation norm.
    """
    max_norm = torch.norm(T, dim=-1).amax(dim=1, keepdim=True)  # (B, 1)
    scale = torch.where(max_norm > eps, max_norm, torch.ones_like(max_norm))
    return T / scale.unsqueeze(-1)


def _nerf_pos_encoding(x: torch.Tensor, freq: int) -> torch.Tensor:
    """NeRF-style sin/cos positional encoding along the last dim."""
    scale = (
        2 ** torch.linspace(0, freq - 1, freq, device=x.device, dtype=x.dtype)
        * math.pi
    )
    encoding = rearrange(x[..., None] * scale, "b t h w i s -> b t h w (i s)")
    return torch.sin(torch.cat([encoding, encoding + 0.5 * math.pi], dim=-1))


@torch.no_grad()
@torch.autocast(device_type="cuda", enabled=False)  # always fp32 for geometry
def compute_ray_encoding(
    raw_poses: torch.Tensor,
    h_lat: int,
    w_lat: int,
    freq: int = 15,
    normalize_trans: bool = False,
) -> torch.Tensor:
    """End-to-end raw poses -> ray-encoding feature volume.

    Args:
        raw_poses: either ``(B, T, 16)`` or ``(B, T, K, 16)``. MiniWorld's RE10K
            pipeline uses ``K=4`` poses inside each WAN-VAE latent chunk.
        h_lat, w_lat: latent spatial size (= model input H, W after VAE).
        freq: NeRF frequency count. ``freq=15`` gives ``6 * 2 * 15 = 180``
            channels per pose.
        normalize_trans: if True, rescale each clip's camera
            translations so the largest displacement is ~1 (see
            ``_normalize_translation_scale``). Disabled by default to match
            DFoT's RealEstate10K preprocessing; static clips are untouched.

    Returns:
        ``(B, T, K * 6 * 2 * freq, H_lat, W_lat)`` float32. The ``K`` poses are
        ray-encoded independently then concatenated along the channel axis
        (so the spatial ``y_embedder`` sees ``K * 180`` channels). For the
        common ``(B, T, 16)`` input ``K=1`` and the output channel count is
        ``180``.
    """
    assert raw_poses.dim() in (3, 4), (
        f"raw_poses must be (B, T, 16) or (B, T, K, 16); got {raw_poses.shape}"
    )
    raw_poses = raw_poses.float()
    if raw_poses.dim() == 3:
        b, t_lat, _ = raw_poses.shape
        k_per_lat = 1
        flat = raw_poses  # (B, T, 16)
    else:
        b, t_lat, k_per_lat, _ = raw_poses.shape
        # Flatten K into the time axis so we can reuse the single-pose pipeline
        # (one shared normalisation anchor = first pose in the sequence).
        flat = raw_poses.reshape(b, t_lat * k_per_lat, 16)

    R, T, K = _split_pose16(flat)
    R, T = _normalize_by_first(R, T)
    if normalize_trans:
        T = _normalize_translation_scale(T)
    origin, direction = _compute_rays(R, T, K, h_res=h_lat, w_res=w_lat)
    enc = torch.cat(
        [
            _nerf_pos_encoding(origin, freq),
            _nerf_pos_encoding(direction, freq),
        ],
        dim=-1,
    )  # (B, T*K, H, W, 6 * 2 * freq)

    if k_per_lat == 1:
        return rearrange(enc, "b t h w c -> b t c h w").contiguous()
    return rearrange(
        enc, "b (t k) h w c -> b t (k c) h w", t=t_lat, k=k_per_lat,
    ).contiguous()


def downsample_poses_to_latent(
    raw_poses: torch.Tensor,
    t_latent: int,
) -> torch.Tensor:
    """Map per-raw-frame poses to four poses per WAN-style latent frame.

    The causal WAN VAE encodes ``T_raw = 4*(T_lat-1)+1`` raw frames into
    ``T_lat`` latents with the temporal grouping:
        * latent 0      -> raw [0]
        * latent j (>0) -> raw [4j-3, 4j-2, 4j-1, 4j]

    Latent 0 has only raw[0], so it is duplicated four times to keep the output
    shape consistent with action conditioning: ``(B, T_lat, 4, 16)``.
    """
    idx_per_latent = [[0, 0, 0, 0]]
    for j in range(1, t_latent):
        idx_per_latent.append([4 * j - 3, 4 * j - 2, 4 * j - 1, 4 * j])
    idx_flat = [i for chunk in idx_per_latent for i in chunk]
    assert raw_poses.shape[1] > max(idx_flat), (
        f"raw_poses has only {raw_poses.shape[1]} frames; need at least "
        f"{max(idx_flat) + 1} to build {t_latent} latent poses."
    )
    b = raw_poses.shape[0]
    return raw_poses[:, idx_flat].view(b, t_latent, 4, 16).contiguous()