dhb-xr 0.2.1__py3-none-any.whl

dhb_xr/integration/vla/pipeline.py

@@ -0,0 +1,85 @@
+"""Minimal VLA pipeline: dataset → DHB encode → tokenization → export."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, Iterable, List, Optional
+
+import numpy as np
+
+from dhb_xr.core.types import DHBMethod, EncodingMethod
+from dhb_xr.encoder.dhb_dr import encode_dhb_dr
+from dhb_xr.tokenization.vqvae import DHBTokenizer
+
+try:
+    import torch
+    HAS_TORCH = True
+except ImportError:  # pragma: no cover - optional dependency
+    HAS_TORCH = False
+
+
+@dataclass
+class DHBVLAPipelineConfig:
+    dhb_method: DHBMethod = DHBMethod.DOUBLE_REFLECTION
+    method: EncodingMethod = EncodingMethod.POSITION
+    use_default_initial_frames: bool = True
+    codebook_size: int = 256
+    latent_dim: int = 32
+
+
+class DHBVLAPipeline:
+    """Minimal end-to-end pipeline that emits DHB tokens."""
+
+    def __init__(self, config: Optional[DHBVLAPipelineConfig] = None):
+        self.config = config or DHBVLAPipelineConfig()
+        self._tokenizer: Optional[DHBTokenizer] = None
+
+    def _get_tokenizer(self, invariant_dim: int) -> DHBTokenizer:
+        if not HAS_TORCH:
+            raise ImportError("torch is required for tokenization (pip install dhb_xr[gpu]).")
+        if self._tokenizer is None:
+            self._tokenizer = DHBTokenizer(
+                invariant_dim=invariant_dim,
+                latent_dim=self.config.latent_dim,
+                codebook_size=self.config.codebook_size,
+            )
+        return self._tokenizer
+
+    def encode_trajectory(self, positions: np.ndarray, quaternions: np.ndarray) -> Dict:
+        return encode_dhb_dr(
+            positions,
+            quaternions,
+            method=self.config.method,
+            use_default_initial_frames=self.config.use_default_initial_frames,
+            dhb_method=self.config.dhb_method,
+        )
+
+    def tokenize_invariants(self, invariants: np.ndarray) -> np.ndarray:
+        tokenizer = self._get_tokenizer(invariants.shape[1])
+        inv_batch = torch.from_numpy(invariants.astype(np.float32)).unsqueeze(0)
+        indices, _, _, _ = tokenizer(inv_batch)
+        return indices.squeeze(0).cpu().numpy()
+
+    def process_dataset(self, episodes: Iterable[Dict]) -> List[Dict]:
+        outputs: List[Dict] = []
+        for ep in episodes:
+            positions = ep["positions"]
+            quaternions = ep["quaternions"]
+            meta = ep.get("metadata", {})
+
+            enc = self.encode_trajectory(positions, quaternions)
+            invariants = np.concatenate(
+                [enc["linear_motion_invariants"], enc["angular_motion_invariants"]],
+                axis=1,
+            )
+            tokens = self.tokenize_invariants(invariants)
+
+            outputs.append(
+                {
+                    "tokens": tokens,
+                    "invariants": invariants,
+                    "initial_pose": enc["initial_pose"],
+                    "metadata": meta,
+                }
+            )
+        return outputs

dhb_xr/integration/vla/robocasa.py

@@ -0,0 +1,85 @@
+"""RoboCASA dataset adapter (HDF5/robomimic-style)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, Iterator, Optional, Tuple
+
+import numpy as np
+
+try:
+    import h5py
+    HAS_H5PY = True
+except ImportError:  # pragma: no cover - optional dependency
+    HAS_H5PY = False
+
+
+DEFAULT_POS_KEYS = (
+    "robot0_eef_pos",
+    "eef_pos",
+    "ee_pos",
+)
+DEFAULT_QUAT_KEYS = (
+    "robot0_eef_quat",
+    "eef_quat",
+    "ee_quat",
+)
+
+
+@dataclass
+class RoboCASAAdapter:
+    """
+    Minimal RoboCASA adapter that yields (positions, quaternions, metadata).
+
+    Assumes a robomimic-style HDF5 with episodes under /data/<demo_id>/obs.
+    """
+
+    pos_keys: Tuple[str, ...] = DEFAULT_POS_KEYS
+    quat_keys: Tuple[str, ...] = DEFAULT_QUAT_KEYS
+    obs_group: str = "obs"
+
+    def _find_key(self, obs_group: "h5py.Group", candidates: Tuple[str, ...]) -> Optional[str]:
+        for key in candidates:
+            if key in obs_group:
+                return key
+        return None
+
+    def load_dataset(self, dataset_path: str) -> Iterator[Dict]:
+        if not HAS_H5PY:
+            raise ImportError("h5py is required for RoboCASA adapter (pip install h5py).")
+
+        with h5py.File(dataset_path, "r") as h5:
+            data_group = h5.get("data")
+            if data_group is None:
+                raise ValueError("RoboCASA HDF5 missing /data group.")
+
+            for demo_id in data_group.keys():
+                demo = data_group[demo_id]
+                obs = demo.get(self.obs_group)
+                if obs is None:
+                    continue
+
+                pos_key = self._find_key(obs, self.pos_keys)
+                quat_key = self._find_key(obs, self.quat_keys)
+                if pos_key is None or quat_key is None:
+                    continue
+
+                positions = np.asarray(obs[pos_key], dtype=np.float64)
+                quaternions = np.asarray(obs[quat_key], dtype=np.float64)
+
+                metadata = {
+                    "demo_id": demo_id,
+                    "pos_key": pos_key,
+                    "quat_key": quat_key,
+                    "source": "robocasa",
+                }
+                if "task" in demo.attrs:
+                    metadata["task"] = demo.attrs["task"]
+                if "language" in demo.attrs:
+                    metadata["language_instruction"] = demo.attrs["language"]
+
+                yield {
+                    "positions": positions,
+                    "quaternions": quaternions,
+                    "metadata": metadata,
+                }

dhb_xr/losses/__init__.py

@@ -0,0 +1,16 @@
+"""Imitation learning losses: invariant, geodesic, hybrid."""
+
+from dhb_xr.losses.invariant_loss import invariant_matching_loss
+from dhb_xr.losses.geodesic_loss import so3_geodesic_loss, se3_geodesic_loss
+
+__all__ = [
+    "invariant_matching_loss",
+    "so3_geodesic_loss",
+    "se3_geodesic_loss",
+    "hybrid_invariant_pose_loss",
+]
+
+try:
+    from dhb_xr.losses.hybrid_loss import hybrid_invariant_pose_loss
+except ImportError:
+    hybrid_invariant_pose_loss = None

dhb_xr/losses/geodesic_loss.py

@@ -0,0 +1,91 @@
+"""SO(3) and SE(3) geodesic losses."""
+
+import numpy as np
+from typing import Optional
+
+from dhb_xr.core import geometry as geom
+
+try:
+    import torch
+    HAS_TORCH = True
+except ImportError:
+    HAS_TORCH = False
+
+
+def so3_geodesic_loss_np(R_pred: np.ndarray, R_demo: np.ndarray) -> float:
+    """||Log(R_pred^T R_demo)||^2. R_pred, R_demo: (3, 3) or (N, 3, 3)."""
+    R_pred = np.asarray(R_pred)
+    R_demo = np.asarray(R_demo)
+    if R_pred.ndim == 2:
+        R_pred = R_pred.reshape(1, 3, 3)
+        R_demo = R_demo.reshape(1, 3, 3)
+    R_diff = np.einsum("...ji,...jk->...ik", R_pred, R_demo)
+    rvec = np.array([geom.rot_to_axis_angle(R_diff[i]) for i in range(len(R_diff))])
+    return float(np.sum(rvec ** 2))
+
+
+def se3_geodesic_loss_np(
+    pos_pred: np.ndarray,
+    quat_pred: np.ndarray,
+    pos_demo: np.ndarray,
+    quat_demo: np.ndarray,
+    beta: float = 1.0,
+) -> float:
+    """Position L2 + beta * SO3 geodesic. Quaternions wxyz."""
+    loss_pos = np.sum((pos_pred - pos_demo) ** 2)
+    R_pred = geom.quat_to_rot(quat_pred)
+    R_demo = geom.quat_to_rot(quat_demo)
+    if R_pred.ndim == 2:
+        R_pred = R_pred.reshape(1, 3, 3)
+        R_demo = R_demo.reshape(1, 3, 3)
+    R_diff = np.einsum("...ji,...jk->...ik", R_pred, R_demo)
+    rvec = np.array([geom.rot_to_axis_angle(R_diff[i]) for i in range(len(R_diff))])
+    loss_rot = np.sum(rvec ** 2)
+    return float(loss_pos + beta * loss_rot)
+
+
+def so3_geodesic_loss(R_pred, R_demo):
+    """Dispatch to numpy or torch."""
+    if hasattr(R_pred, "numpy"):
+        return so3_geodesic_loss_torch(R_pred, R_demo)
+    return so3_geodesic_loss_np(R_pred, R_demo)
+
+
+def se3_geodesic_loss(
+    pos_pred, quat_pred, pos_demo, quat_demo, beta: float = 1.0
+):
+    if hasattr(pos_pred, "numpy"):
+        return se3_geodesic_loss_torch(pos_pred, quat_pred, pos_demo, quat_demo, beta)
+    return se3_geodesic_loss_np(pos_pred, quat_pred, pos_demo, quat_demo, beta)
+
+
+if HAS_TORCH:
+    try:
+        from dhb_xr.core.geometry_torch import rot_to_axis_angle_torch, quat_to_rot_torch
+    except ImportError:
+        rot_to_axis_angle_torch = None
+        quat_to_rot_torch = None
+
+    def so3_geodesic_loss_torch(R_pred: torch.Tensor, R_demo: torch.Tensor) -> torch.Tensor:
+        R_diff = R_pred.transpose(-2, -1) @ R_demo
+        rvec = rot_to_axis_angle_torch(R_diff)
+        return (rvec ** 2).sum()
+
+    def se3_geodesic_loss_torch(
+        pos_pred: torch.Tensor,
+        quat_pred: torch.Tensor,
+        pos_demo: torch.Tensor,
+        quat_demo: torch.Tensor,
+        beta: float = 1.0,
+    ) -> torch.Tensor:
+        loss_pos = ((pos_pred - pos_demo) ** 2).sum()
+        if quat_to_rot_torch is not None:
+            R_pred = quat_to_rot_torch(quat_pred)
+            R_demo = quat_to_rot_torch(quat_demo)
+            loss_rot = so3_geodesic_loss_torch(R_pred, R_demo)
+        else:
+            loss_rot = torch.tensor(0.0, device=pos_pred.device)
+        return loss_pos + beta * loss_rot
+else:
+    so3_geodesic_loss_torch = None
+    se3_geodesic_loss_torch = None

dhb_xr/losses/hybrid_loss.py

@@ -0,0 +1,36 @@
+"""Hybrid invariant + pose-space loss for imitation learning."""
+
+import numpy as np
+from typing import Optional
+
+from dhb_xr.losses.invariant_loss import invariant_matching_loss
+from dhb_xr.losses.geodesic_loss import se3_geodesic_loss_np
+
+
+def hybrid_invariant_pose_loss(
+    pred_positions: np.ndarray,
+    pred_quaternions: np.ndarray,
+    demo_positions: np.ndarray,
+    demo_quaternions: np.ndarray,
+    pred_invariants: Optional[np.ndarray] = None,
+    demo_invariants: Optional[np.ndarray] = None,
+    alpha: float = 0.5,
+    beta: float = 1.0,
+) -> float:
+    """
+    alpha * invariant_loss + (1-alpha) * pose_loss.
+    If pred_invariants/demo_invariants are None, only pose loss is used (alpha ignored for invariant part).
+    """
+    loss_pose = 0.0
+    n = len(pred_positions)
+    assert n == len(demo_positions) and n == len(pred_quaternions) and n == len(demo_quaternions)
+    for i in range(n):
+        loss_pose += se3_geodesic_loss_np(
+            pred_positions[i], pred_quaternions[i],
+            demo_positions[i], demo_quaternions[i],
+            beta=beta,
+        )
+    if pred_invariants is not None and demo_invariants is not None:
+        loss_inv = invariant_matching_loss(pred_invariants, demo_invariants)
+        return float(alpha * loss_inv + (1 - alpha) * loss_pose)
+    return float(loss_pose)

dhb_xr/losses/invariant_loss.py

@@ -0,0 +1,73 @@
+"""Imitation learning losses in invariant space."""
+
+import numpy as np
+from typing import Optional
+
+from dhb_xr.core import geometry as geom
+
+try:
+    import torch
+    import torch.nn.functional as F
+    HAS_TORCH = True
+except ImportError:
+    HAS_TORCH = False
+
+
+def quaternion_geodesic_loss_np(q1: np.ndarray, q2: np.ndarray) -> float:
+    """Sum of squared quaternion geodesic distances. q1, q2: (N, 4) wxyz."""
+    dot = np.abs(np.sum(q1 * q2, axis=-1))
+    dot = np.clip(dot, 0, 1)
+    return np.sum((2 * np.arccos(dot)) ** 2)
+
+
+def invariant_matching_loss(
+    pred_inv: np.ndarray,
+    demo_inv: np.ndarray,
+    method: str = "dhb_dr",
+    weights: Optional[np.ndarray] = None,
+) -> float:
+    """
+    pred_inv, demo_inv: (N, 2*k). method 'dhb_dr' (Euler) or 'dhb_qr' (quaternion).
+    For dhb_qr, angular quaternion part uses geodesic; else L2 with optional angle wrap.
+    """
+    pred_inv = np.asarray(pred_inv)
+    demo_inv = np.asarray(demo_inv)
+    assert pred_inv.shape == demo_inv.shape
+    k = pred_inv.shape[1] // 2
+    if weights is None:
+        weights = np.ones(pred_inv.shape[1])
+    if method == "dhb_qr":
+        m_lin = np.sum(weights[0] * (pred_inv[:, 0] - demo_inv[:, 0]) ** 2)
+        q_lin = pred_inv[:, 1:5]
+        q_lin_d = demo_inv[:, 1:5]
+        m_ang = np.sum(weights[k] * (pred_inv[:, k] - demo_inv[:, k]) ** 2)
+        q_ang = pred_inv[:, k + 1 : k + 5]
+        q_ang_d = demo_inv[:, k + 1 : k + 5]
+        loss_lin_q = quaternion_geodesic_loss_np(q_lin, q_lin_d)
+        loss_ang_q = quaternion_geodesic_loss_np(q_ang, q_ang_d)
+        return float(m_lin + m_ang + loss_lin_q + loss_ang_q)
+    diff = pred_inv - demo_inv
+    return float(np.sum(weights * (diff ** 2)))
+
+
+if HAS_TORCH:
+
+    def invariant_matching_loss_torch(
+        pred_inv: torch.Tensor,
+        demo_inv: torch.Tensor,
+        method: str = "dhb_dr",
+    ) -> torch.Tensor:
+        if method == "dhb_qr":
+            m_lin = F.mse_loss(pred_inv[..., 0], demo_inv[..., 0])
+            q_lin = pred_inv[..., 1:5]
+            q_lin_d = demo_inv[..., 1:5]
+            dot = (q_lin * q_lin_d).sum(dim=-1).abs().clamp(0, 1)
+            loss_lin_q = (2 * torch.acos(dot)).pow(2).sum()
+            k = pred_inv.shape[-1] // 2
+            m_ang = F.mse_loss(pred_inv[..., k], demo_inv[..., k])
+            q_ang = pred_inv[..., k + 1 : k + 5]
+            q_ang_d = demo_inv[..., k + 1 : k + 5]
+            dot = (q_ang * q_ang_d).sum(dim=-1).abs().clamp(0, 1)
+            loss_ang_q = (2 * torch.acos(dot)).pow(2).sum()
+            return m_lin + m_ang + loss_lin_q + loss_ang_q
+        return F.mse_loss(pred_inv, demo_inv)

dhb_xr/optimization/__init__.py

@@ -0,0 +1,72 @@
+"""Trajectory optimization: CasADi, Cusadi, Fatrop, PyTorch."""
+
+from dhb_xr.optimization.casadi_solver import generate_trajectory
+
+__all__ = ["generate_trajectory", "get_optimizer"]
+
+try:
+    from dhb_xr.optimization.cusadi_solver import batched_decode_dhb_dr, CusadiTrajectoryOptimizer
+    __all__.extend(["batched_decode_dhb_dr", "CusadiTrajectoryOptimizer"])
+except ImportError:
+    batched_decode_dhb_dr = None
+    CusadiTrajectoryOptimizer = None
+
+try:
+    from dhb_xr.optimization.torch_solver import BatchedTrajectoryOptimizer
+    __all__.append("BatchedTrajectoryOptimizer")
+except ImportError:
+    BatchedTrajectoryOptimizer = None
+
+try:
+    from dhb_xr.optimization.fatrop_solver import (
+        FatropTrajectoryGenerator,
+        ConstrainedTrajectoryGenerator,
+        generate_trajectory_fatrop,
+    )
+    __all__.extend([
+        "FatropTrajectoryGenerator",
+        "ConstrainedTrajectoryGenerator",
+        "generate_trajectory_fatrop",
+    ])
+except ImportError:
+    FatropTrajectoryGenerator = None
+    ConstrainedTrajectoryGenerator = None
+    generate_trajectory_fatrop = None
+
+
+def get_optimizer(backend="auto", batch_size=1, device="cpu", **kwargs):
+    """Factory for optimal backend selection.
+
+    Args:
+        backend: One of "auto", "torch", "cusadi", "fatrop", "ipopt"
+        batch_size: Batch size for batched optimizers
+        device: Device for torch-based optimizers
+        **kwargs: Additional arguments passed to the optimizer
+
+    Returns:
+        Optimizer instance or None
+    """
+    if backend == "auto":
+        if batch_size == 1:
+            return None  # Use CasADi generate_trajectory
+        try:
+            from dhb_xr.optimization.torch_solver import BatchedTrajectoryOptimizer
+            return BatchedTrajectoryOptimizer(device=device)
+        except ImportError:
+            return None
+    if backend == "torch":
+        from dhb_xr.optimization.torch_solver import BatchedTrajectoryOptimizer
+        return BatchedTrajectoryOptimizer(device=device)
+    if backend == "cusadi":
+        if CusadiTrajectoryOptimizer is not None:
+            return CusadiTrajectoryOptimizer(batch_size=batch_size)
+        return None
+    if backend == "fatrop":
+        if FatropTrajectoryGenerator is not None:
+            return FatropTrajectoryGenerator(use_fatrop=True, **kwargs)
+        return None
+    if backend == "ipopt":
+        if FatropTrajectoryGenerator is not None:
+            return FatropTrajectoryGenerator(use_fatrop=False, **kwargs)
+        return None
+    return None