Implementation:Isaac sim IsaacGymEnvs Torch Jit Utils

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

Overview

The torch_jit_utils module provides a collection of @torch.jit.script-decorated utility functions for GPU-accelerated 3D mathematics, serving as the primary math utility layer used throughout all IsaacGymEnvs tasks.

Description

The module at isaacgymenvs/utils/torch_jit_utils.py implements the foundational GPU math functions that nearly every task in IsaacGymEnvs depends on. It provides optimized, JIT-compiled functions for quaternion operations, vector mathematics, random number generation, and coordinate transformations. Unlike the poselib rotation3d module (which serves the AMP/poselib subsystem), this module is the global utility layer imported by VecTask subclasses, Factory tasks, Allegro-Kuka tasks, and all other environments.

The quaternion functions include quat_mul() for quaternion multiplication, quat_apply() for rotating vectors by quaternions, quat_conjugate() for computing inverses, and quat_from_angle_axis() for constructing quaternions from angle-axis representation. The module uses the (x, y, z, w) quaternion convention consistent with Isaac Gym. Additional quaternion utilities include slerp() for interpolation, quat_to_exp_map() for exponential map conversion, and quat_to_angle_axis().

General math utilities include normalize() for vector normalization with epsilon clamping, tensor_clamp() for per-element clamping with tensor bounds, torch_rand_float() for generating uniform random floats in a range, and get_euler_xyz() for quaternion-to-Euler decomposition. Navigation-specific helpers like compute_heading_and_up(), calc_heading_quat(), and calc_heading_quat_inv() support locomotion tasks. The helper to_torch() (the only non-JIT function) provides convenient numpy-to-torch tensor conversion with device placement.

Usage

Use these utility functions in any IsaacGymEnvs task for GPU-accelerated 3D math operations. They are typically imported via wildcard (from isaacgymenvs.utils.torch_jit_utils import *) in task files, providing the standard math toolkit for observation computation, reward calculation, and action processing.

Code Reference

Source Location

• Repository: IsaacGymEnvs
• File: isaacgymenvs/utils/torch_jit_utils.py
• Lines: 1-669

Signature

def to_torch(x, dtype=torch.float, device='cuda:0', requires_grad=False):
    """Convert numpy array or list to torch tensor on specified device."""

@torch.jit.script
def quat_mul(a, b):
    """Quaternion multiplication with (x, y, z, w) convention."""

@torch.jit.script
def normalize(x, eps: float = 1e-9):
    """L2-normalize vectors along last dimension with epsilon clamping."""

@torch.jit.script
def quat_apply(a, b):
    """Rotate 3D vector b by quaternion a using cross-product formulation."""

@torch.jit.script
def quat_conjugate(a):
    """Compute quaternion conjugate (negate xyz, keep w)."""

@torch.jit.script
def quat_from_angle_axis(angle, axis):
    """Create quaternion from rotation angle and axis."""

@torch.jit.script
def tensor_clamp(t, min_t, max_t):
    """Per-element clamp with tensor min/max bounds."""

@torch.jit.script
def torch_rand_float(lower, upper, shape, device):
    """Generate uniform random floats in [lower, upper]."""

@torch.jit.script
def get_euler_xyz(q):
    """Extract Euler angles (roll, pitch, yaw) from quaternion."""

@torch.jit.script
def compute_heading_and_up(torso_rotation):
    """Compute heading direction and up vector from torso quaternion."""

@torch.jit.script
def calc_heading_quat(torso_rotation):
    """Compute yaw-only heading quaternion from full rotation."""

@torch.jit.script
def calc_heading_quat_inv(torso_rotation):
    """Compute inverse heading quaternion for local frame computation."""

@torch.jit.script
def slerp(q0, q1, t):
    """Spherical linear interpolation between quaternions."""

@torch.jit.script
def quat_to_exp_map(q):
    """Convert quaternion to exponential map representation."""

@torch.jit.script
def quat_to_angle_axis(q):
    """Convert quaternion to angle-axis representation."""

@torch.jit.script
def my_quat_rotate(q, v):
    """Alternative quaternion-vector rotation implementation."""

@torch.jit.script
def quat_to_tan_norm(q):
    """Convert quaternion to tangent-normal representation."""

@torch.jit.script
def normalize_angle(x):
    """Normalize angle to [-pi, pi] range."""

Import

from isaacgymenvs.utils.torch_jit_utils import quat_mul, quat_apply, to_torch

I/O Contract

Inputs

Name	Type	Required	Description
x	array_like	Yes	Input data for to_torch() conversion (numpy array, list, or scalar)
a, b	Tensor(..., 4)	Yes	Quaternion tensors in (x, y, z, w) format for quat_mul/quat_apply
angle	Tensor(...)	Yes	Rotation angles in radians for angle-axis construction
axis	Tensor(..., 3)	Yes	Unit rotation axes for angle-axis construction
t, min_t, max_t	Tensor	Yes	Tensor values for clamping or interpolation
lower, upper	float	Yes	Range bounds for random float generation
shape	Tuple[int]	Yes	Output shape for random tensor generation
device	str	Yes	Target device string (e.g., "cuda:0")
torso_rotation	Tensor(..., 4)	Yes	Torso orientation quaternion for heading computation

Outputs

Name	Type	Description
tensor	Tensor	Converted torch tensor on specified device (from to_torch)
quaternion	Tensor(..., 4)	Result quaternion from multiplication, construction, or conversion
vector	Tensor(..., 3)	Rotated or normalized 3D vector
euler	Tuple[Tensor, Tensor, Tensor]	Roll, pitch, yaw angles from get_euler_xyz
heading_quat	Tensor(..., 4)	Yaw-only quaternion from calc_heading_quat
random_tensor	Tensor	Uniform random values in specified range and shape

Usage Examples

import torch
import numpy as np
from isaacgymenvs.utils.torch_jit_utils import (
    to_torch, quat_mul, quat_apply, quat_conjugate,
    quat_from_angle_axis, normalize, tensor_clamp,
    torch_rand_float, get_euler_xyz, calc_heading_quat_inv
)

device = "cuda:0"

# Convert numpy data to GPU tensor
positions_np = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
positions = to_torch(positions_np, device=device)

# Create rotation from angle-axis
angle = torch.tensor([1.57], device=device)
axis = torch.tensor([[0.0, 0.0, 1.0]], device=device)
quat = quat_from_angle_axis(angle, axis)

# Rotate a vector
v = torch.tensor([[1.0, 0.0, 0.0]], device=device)
rotated_v = quat_apply(quat, v)

# Compose rotations
q_combined = quat_mul(quat, quat)

# Compute heading for locomotion observation
torso_quat = torch.randn(1024, 4, device=device)
torso_quat = normalize(torso_quat)
heading_inv = calc_heading_quat_inv(torso_quat)

# Random action noise
noise = torch_rand_float(-0.1, 0.1, (1024, 28), device)

# Clamp DOF targets
dof_lower = to_torch(np.array([-2.0] * 28), device=device)
dof_upper = to_torch(np.array([2.0] * 28), device=device)
targets = tensor_clamp(targets_raw, dof_lower, dof_upper)