RoMa: 3D Rotations in PyTorch

Common errors

• ModuleNotFoundError: No module named 'roma'

  cause The 'roma' package has not been installed in the current Python environment.
  fix Run `pip install roma` to install the library.

• RuntimeError: `special_procrustes` on CUDA GPUs with PyTorch < 1.8 requires `torch-batch-svd` for speed-up. Please install it using `pip install torch-batch-svd`.

  cause Using `roma.special_procrustes` on a CUDA-enabled GPU with an older PyTorch version (<1.8) without the optional `torch-batch-svd` dependency, leading to a performance warning/error (though it might still run slowly).
  fix Install the recommended dependency: `pip install torch-batch-svd`. Alternatively, upgrade your PyTorch installation to version 1.8 or newer, which includes similar optimizations natively.

Warnings

• gotcha Euler angles are prone to gimbal lock and numerical instability. While supported for convenience (e.g., user input/output), it's highly recommended to use quaternions or rotation matrices for actual computations to avoid these issues.
  Fix: Prefer `roma.rotvec_to_unitquat`, `roma.unitquat_to_rotmat`, or `roma.rotvec_to_rotmat` for internal computations and transformations.
• gotcha Some utility functions, such as `roma.utils.rotmat_inverse()`, may return a transposed 'view' of the input tensor. Performing in-place operations on these views can lead to unexpected side effects or alter the original tensor. Always clone the output if subsequent in-place modifications are intended.
  Fix: If modifying the output of `rotmat_inverse` or similar functions, ensure to call `.clone()` on the result, e.g., `R_inv = roma.utils.rotmat_inverse(R).clone()`.
• gotcha For PyTorch versions below 1.8, `roma.special_procrustes()` on CUDA GPUs can be significantly slower if `torch-batch-svd` is not installed. This is a performance regression rather than a functional breaking change, but can severely impact efficiency.
  Fix: Install `torch-batch-svd` via `pip install torch-batch-svd` if using PyTorch < 1.8 and needing fast `special_procrustes` on GPU, or upgrade PyTorch to version 1.8 or newer.

Install

• pip install roma PyPI

Imports

• roma

  import roma

• roma.utils

  import roma.utils

Quickstart

This quickstart demonstrates how to initialize `roma` and perform fundamental conversions between different 3D rotation representations: rotation vectors, unit quaternions, and rotation matrices. It also shows the use of `special_procrustes` to orthonormalize an arbitrary matrix into a valid rotation matrix.

import torch
import roma

# Example: Convert a batch of rotation vectors to unit quaternions and then to rotation matrices
batch_shape = (2, 3) # Example: 2 batches of 3 rotations each

# Generate random rotation vectors (3D tensor)
rotvec = torch.randn(batch_shape + (3,))
print(f"Rotation vector shape: {rotvec.shape}")

# Convert rotation vectors to unit quaternions (XYZW convention)
q = roma.rotvec_to_unitquat(rotvec)
print(f"Unit quaternion shape: {q.shape}")

# Convert unit quaternions to rotation matrices (3x3 tensor)
R = roma.unitquat_to_rotmat(q)
print(f"Rotation matrix shape: {R.shape}")

# Direct conversion from rotation vector to rotation matrix
R_direct = roma.rotvec_to_rotmat(rotvec)
print(f"Direct Rotation matrix shape: {R_direct.shape}")
assert torch.allclose(R, R_direct, atol=1e-6)

# Example: Special Procrustes orthonormalization
# Projects an arbitrary 3x3 matrix onto the closest rotation matrix
random_matrix = torch.randn(batch_shape + (3, 3))
R_procrustes = roma.special_procrustes(random_matrix)
print(f"Procrustes orthonormalized matrix shape: {R_procrustes.shape}")