Implementation:Isaac sim IsaacGymEnvs AMP Torch Utils

Knowledge Sources	IsaacGymEnvs
Domains	3D_Mathematics, Motion_Imitation
Last Updated	2026-02-15 11:00 GMT

Overview

AMP_Torch_Utils provides a collection of JIT-compiled PyTorch utility functions for quaternion rotation, orientation representation conversion, spherical interpolation, and heading calculation used throughout the AMP (Adversarial Motion Priors) system.

Description

This module contains a suite of performance-critical mathematical functions decorated with @torch.jit.script for GPU-accelerated execution. The functions handle quaternion-based 3D rotation operations that are fundamental to the AMP motion imitation framework.

The core rotation functions include my_quat_rotate() which rotates vectors by quaternions using an optimized manual expansion (avoiding explicit quaternion-to-matrix conversion), quat_to_angle_axis() which decomposes a unit quaternion into an angle and axis pair with numerical stability handling for small rotations, and quat_to_exp_map() which converts quaternions to exponential map (axis-angle vector) representation. The quat_to_tan_norm() function provides an alternative rotation representation using tangent and normal vectors, which can be more suitable for neural network inputs.

Interpolation is handled by the slerp() function, which performs spherical linear interpolation between two quaternions with proper handling of antipodal quaternions (flipping when dot product is negative) and degenerate cases (near-zero sine or near-unity cosine). Heading-related functions include calc_heading() which extracts the yaw angle on the XY plane from a quaternion, calc_heading_quat() which constructs a heading-only quaternion, and calc_heading_quat_inv() which provides the inverse heading quaternion. Additional conversion utilities cover exponential map to angle-axis and back to quaternion, as well as Euler XYZ to exponential map conversion.

Usage

Use these functions when working with the AMP task environments for computing observation representations, reward calculations, or motion state transformations. The heading functions are particularly important for computing heading-invariant observations, while slerp is used for motion interpolation between keyframes. The exponential map and tangent-normal conversions are used to prepare rotation data as neural network inputs.

Code Reference

Source Location

• Repository: IsaacGymEnvs
• File: isaacgymenvs/tasks/amp/utils_amp/amp_torch_utils.py
• Lines: 37-207

Signature

@torch.jit.script
def my_quat_rotate(q: Tensor, v: Tensor) -> Tensor: ...

@torch.jit.script
def quat_to_angle_axis(q: Tensor) -> Tuple[Tensor, Tensor]: ...

@torch.jit.script
def angle_axis_to_exp_map(angle: Tensor, axis: Tensor) -> Tensor: ...

@torch.jit.script
def quat_to_exp_map(q: Tensor) -> Tensor: ...

@torch.jit.script
def quat_to_tan_norm(q: Tensor) -> Tensor: ...

@torch.jit.script
def euler_xyz_to_exp_map(roll: Tensor, pitch: Tensor, yaw: Tensor) -> Tensor: ...

@torch.jit.script
def exp_map_to_angle_axis(exp_map: Tensor) -> Tuple[Tensor, Tensor]: ...

@torch.jit.script
def exp_map_to_quat(exp_map: Tensor) -> Tensor: ...

@torch.jit.script
def slerp(q0: Tensor, q1: Tensor, t: Tensor) -> Tensor: ...

@torch.jit.script
def calc_heading(q: Tensor) -> Tensor: ...

@torch.jit.script
def calc_heading_quat(q: Tensor) -> Tensor: ...

@torch.jit.script
def calc_heading_quat_inv(q: Tensor) -> Tensor: ...

Import

from isaacgymenvs.tasks.amp.utils_amp.amp_torch_utils import slerp, calc_heading, quat_to_exp_map

I/O Contract

Inputs

Name	Type	Required	Description
q	Tensor	Yes	Unit quaternion(s) in XYZW format, shape (..., 4)
v	Tensor	Yes (my_quat_rotate)	3D vector(s) to rotate, shape (N, 3)
q0	Tensor	Yes (slerp)	Start quaternion(s) for interpolation
q1	Tensor	Yes (slerp)	End quaternion(s) for interpolation
t	Tensor	Yes (slerp)	Interpolation parameter(s) in [0, 1]
angle	Tensor	Yes (angle_axis_to_exp_map)	Rotation angle(s) in radians
axis	Tensor	Yes (angle_axis_to_exp_map)	Unit rotation axis/axes, shape (..., 3)
exp_map	Tensor	Yes (exp_map_to_*)	Exponential map rotation(s), shape (..., 3)
roll, pitch, yaw	Tensor	Yes (euler_xyz_to_exp_map)	Euler angles in radians

Outputs

Name	Type	Description
rotated_v	Tensor	Rotated 3D vector(s), shape (N, 3) (from my_quat_rotate)
angle	Tensor	Rotation angle in radians (from quat_to_angle_axis)
axis	Tensor	Unit rotation axis, shape (..., 3) (from quat_to_angle_axis)
exp_map	Tensor	Exponential map representation, shape (..., 3) (from quat_to_exp_map)
norm_tan	Tensor	Concatenated tangent and normal vectors, shape (..., 6) (from quat_to_tan_norm)
new_q	Tensor	Interpolated quaternion(s), shape (..., 4) (from slerp)
heading	Tensor	Heading angle on XY plane in radians (from calc_heading)
heading_q	Tensor	Heading-only quaternion, shape (..., 4) (from calc_heading_quat)

Usage Examples

import torch
from isaacgymenvs.tasks.amp.utils_amp.amp_torch_utils import (
    my_quat_rotate, quat_to_exp_map, slerp,
    calc_heading, calc_heading_quat_inv, quat_to_tan_norm
)

# Rotate a batch of vectors by quaternions
quats = torch.randn(100, 4, device="cuda")
quats = quats / quats.norm(dim=-1, keepdim=True)  # normalize
vectors = torch.randn(100, 3, device="cuda")
rotated = my_quat_rotate(quats, vectors)

# Convert quaternions to exponential map for neural network input
exp_maps = quat_to_exp_map(quats)  # shape: (100, 3)

# Convert quaternions to tangent-normal representation
tan_norm = quat_to_tan_norm(quats)  # shape: (100, 6)

# Interpolate between two sets of quaternions
q_start = torch.randn(100, 4, device="cuda")
q_start = q_start / q_start.norm(dim=-1, keepdim=True)
q_end = torch.randn(100, 4, device="cuda")
q_end = q_end / q_end.norm(dim=-1, keepdim=True)
t = torch.full((100, 1), 0.5, device="cuda")
q_mid = slerp(q_start, q_end, t)

# Extract heading direction and compute heading-invariant rotation
heading_angle = calc_heading(quats)
heading_inv = calc_heading_quat_inv(quats)