Trajectory Optimization

Overview

DHB-XR provides optimization-based trajectory adaptation using CasADi and Fatrop solvers.

Main Functions

batched_decode_dhb_dr

batched_decode_dhb_dr

batched_decode_dhb_dr(
    linear_invariants_batch,
    angular_invariants_batch,
    initial_poses,
    method=EncodingMethod.POSITION,
    dhb_method=DHBMethod.DOUBLE_REFLECTION,
    drop_padded=True,
    use_gpu=False,
    casadi_path=DEFAULT_FN_PATH,
)

Decode multiple trajectories in batch.

Parameters:

Name	Type	Description	Default
linear_invariants_batch	ndarray	(B, T, 4) or list of (T, 4)	required
angular_invariants_batch	ndarray	(B, T, 4) or list of (T, 4)	required
initial_poses	List[Dict[str, ndarray]]	list of B dicts with 'position' (3,) and 'quaternion' (4,) wxyz.	required
method	EncodingMethod	'pos' or 'vel' for invariant interpretation	POSITION
dhb_method	DHBMethod	DHBMethod enum (DOUBLE_REFLECTION or ORIGINAL)	DOUBLE_REFLECTION
drop_padded	bool	Whether to drop padded frames	True
use_gpu	bool	If True and CUDA available, use CusADi GPU acceleration	False
casadi_path	str	Path to compiled CasADi function for GPU decode	DEFAULT_FN_PATH

Returns:

Type	Description
(positions_batch, quaternions_batch)	(B, N, 3), (B, N, 4).

Source code in src/dhb_xr/optimization/cusadi_solver.py

def batched_decode_dhb_dr(
    linear_invariants_batch: np.ndarray,
    angular_invariants_batch: np.ndarray,
    initial_poses: List[Dict[str, np.ndarray]],
    method: EncodingMethod = EncodingMethod.POSITION,
    dhb_method: DHBMethod = DHBMethod.DOUBLE_REFLECTION,
    drop_padded: bool = True,
    use_gpu: bool = False,
    casadi_path: str = DEFAULT_FN_PATH,
) -> tuple:
    """
    Decode multiple trajectories in batch.

    Args:
        linear_invariants_batch: (B, T, 4) or list of (T, 4)
        angular_invariants_batch: (B, T, 4) or list of (T, 4)
        initial_poses: list of B dicts with 'position' (3,) and 'quaternion' (4,) wxyz.
        method: 'pos' or 'vel' for invariant interpretation
        dhb_method: DHBMethod enum (DOUBLE_REFLECTION or ORIGINAL)
        drop_padded: Whether to drop padded frames
        use_gpu: If True and CUDA available, use CusADi GPU acceleration
        casadi_path: Path to compiled CasADi function for GPU decode

    Returns:
        (positions_batch, quaternions_batch): (B, N, 3), (B, N, 4).
    """
    # Try GPU decode if requested
    if use_gpu:
        if HAS_TORCH_CUDA and HAS_CUSADI:
            try:
                return batched_decode_dhb_dr_gpu(
                    linear_invariants_batch, angular_invariants_batch,
                    initial_poses, casadi_path
                )
            except Exception as e:
                print(f"GPU decode failed, falling back to CPU: {e}")
        else:
            missing = []
            if not HAS_TORCH_CUDA:
                missing.append("PyTorch CUDA")
            if not HAS_CUSADI:
                missing.append("CusADi")
            print(f"GPU decode unavailable (missing: {', '.join(missing)}), using CPU")

    # CPU decode (NumPy loop)
    if isinstance(linear_invariants_batch, np.ndarray) and linear_invariants_batch.ndim == 3:
        B = linear_invariants_batch.shape[0]
        lin_list = [linear_invariants_batch[b] for b in range(B)]
        ang_list = [angular_invariants_batch[b] for b in range(B)]
    else:
        lin_list = list(linear_invariants_batch)
        ang_list = list(angular_invariants_batch)
        B = len(lin_list)
    assert len(initial_poses) == B and len(ang_list) == B
    pos_list = []
    quat_list = []
    for b in range(B):
        decoded = decode_dhb_dr(
            lin_list[b], ang_list[b],
            initial_poses[b],
            method=method,
            dhb_method=dhb_method,
            drop_padded=drop_padded,
        )
        pos_list.append(decoded["positions"])
        quat_list.append(decoded["quaternions"])
    return np.stack(pos_list), np.stack(quat_list)

Usage Example

from dhb_xr.optimization.casadi_solver import CasadiTrajectoryOptimizer
import numpy as np

# Create optimizer
optimizer = CasadiTrajectoryOptimizer(batch_size=1)

# Sample data
lin_inv = np.random.randn(1, 20, 4)
ang_inv = np.random.randn(1, 20, 4)
init_poses = [{"position": np.zeros(3), "quaternion": np.array([1, 0, 0, 0])}]

# Optimize trajectory
positions, quaternions = optimizer.forward(lin_inv, ang_inv, init_poses)