base_fk

Class Introduction

The RobotKinematics class provides a comprehensive PyTorch-based forward and inverse kinematics solution for robotic systems. It parses URDF (Unified Robot Description Format) files to build kinematic trees and performs efficient batch computations for robot pose calculations. This class is designed to work seamlessly with PyTorch tensors and supports both CPU and GPU computation.

Key features:

URDF-based initialization: Automatically parses robot structure from URDF files
Batch processing: Efficient computation for multiple robot configurations simultaneously
PyTorch integration: Full gradient support for optimization and machine learning applications
Flexible device support: Works on both CPU and CUDA devices
Kinematic tree visualization: Built-in tools for visualizing robot structure and joint positions

⚠️Warning: We use urdf_parser_py to load URDF files, but this library is old and may not support all URDF features. If anything went wrong with this package, please consider using pybullet or pyrobot for more robust URDF parsing (and replace our RobotKinematics class with your own implementation).

Class Dependencies

import os
import numpy as np
import sys, contextlib
from urdf_parser_py.urdf import URDF
from scipy.spatial.transform import Rotation as R
import networkx as nx
import matplotlib.pyplot as plt
import torch
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import torch.nn as nn
import torch.nn.functional as F
from torch import einsum, matmul
from typing import Optional, List
from GBC.utils.base.math_utils import euler_xyz_to_rot_mat

Class Functions

🏗️ init

Module Name: GBC.utils.base.base_fk.RobotKinematics.__init__

Definition:

def __init__(self, urdf_path: str, device: str = 'cpu')

📥 Input:

urdf_path (str): Path to the URDF file
device (str): Computation device ('cpu' or 'cuda'). Default is 'cpu'

📤 Output:

Initializes RobotKinematics instance

🔧 Functionality: Initializes the RobotKinematics class by loading and parsing a URDF file. Sets up the kinematic tree, joint mappings, and prepares the robot for kinematic computations.

⚠️ Exceptions:

FileNotFoundError: If the URDF file is not found at the specified path
ValueError: If URDF doesn't have exactly one base link

💡 Example Usage:

fk = RobotKinematics(
    urdf_path="/path/to/robot.urdf", 
    device='cuda'
)

🔄 to

Module Name: GBC.utils.base.base_fk.RobotKinematics.to

Definition:

def to(self, device)

📥 Input:

device: Target device for tensor operations

📤 Output:

Returns self for method chaining

🔧 Functionality: Moves all internal tensors to the specified device and updates the device attribute.

💡 Example Usage:

fk = fk.to('cuda')

🎯 find_base_link

Module Name: GBC.utils.base.base_fk.RobotKinematics.find_base_link

Definition:

def find_base_link(self)

📥 Input:

None (self)

📤 Output:

base_link_name (str): Name of the base link

🔧 Functionality: Identifies the base link of the robot by finding the link that has no parent joint in the kinematic tree.

⚠️ Exceptions:

ValueError: If URDF doesn't have exactly one base link

💡 Example Usage:

base_link = fk.find_base_link()
print(f"Base link: {base_link}")

🔗 get_joint_names

Module Name: GBC.utils.base.base_fk.RobotKinematics.get_joint_names

Definition:

def get_joint_names(self)

📥 Input:

None (self)

📤 Output:

joint_names (list): List of all joint names (excluding fixed joints)

🔧 Functionality: Returns a list of all joint names in the robot, excluding fixed joints.

💡 Example Usage:

joints = fk.get_joint_names()
print(f"Available joints: {joints}")

🧩 get_link_names

Module Name: GBC.utils.base.base_fk.RobotKinematics.get_link_names

Definition:

def get_link_names(self)

📥 Input:

None (self)

📤 Output:

link_names (list): List of all link names

🔧 Functionality: Returns a list of all link names in the robot.

💡 Example Usage:

links = fk.get_link_names()
print(f"Available links: {links}")

🌳 build_kinematic_tree

Module Name: GBC.utils.base.base_fk.RobotKinematics.build_kinematic_tree

Definition:

def build_kinematic_tree(self)

📥 Input:

None (self)

📤 Output:

G (networkx.DiGraph): Directed graph representing the kinematic chain

🔧 Functionality: Constructs a kinematic tree using NetworkX to represent parent-child relationships between links and joints.

💡 Example Usage:

tree = fk.build_kinematic_tree()

🔄 rotation_matrix_batch

Module Name: GBC.utils.base.base_fk.RobotKinematics.rotation_matrix_batch

Definition:

def rotation_matrix_batch(self, rot_vecs)

📥 Input:

rot_vecs (torch.Tensor): Rotation vector (angle-axis), shape (B, 3)

📤 Output:

rot_mat (torch.Tensor): Rotation matrix, shape (B, 3, 3)

🔧 Functionality: Computes rotation matrices from rotation vectors using the angle-axis representation with batch processing support. Handles zero rotation vectors correctly.

💡 Example Usage:

rot_vecs = torch.tensor([[0.1, 0.2, 0.3], [0.0, 0.0, 0.5]])
rot_mats = fk.rotation_matrix_batch(rot_vecs)

⚙️ get_joint_transform_batch

Module Name: GBC.utils.base.base_fk.RobotKinematics.get_joint_transform_batch

Definition:

def get_joint_transform_batch(self, joint, joint_values)

📥 Input:

joint (urdf_parser_py.urdf.Joint): The joint object
joint_values (torch.Tensor): Joint position values, shape (B,)

📤 Output:

T (torch.Tensor): 4x4 transformation matrices, shape (B, 4, 4)

🔧 Functionality: Computes the transformation matrix for a given joint and joint values, supporting revolute, continuous, prismatic, and fixed joint types.

⚠️ Exceptions:

NotImplementedError: If joint type is not supported

💡 Example Usage:

joint = fk.joint_map['shoulder_joint']
joint_values = torch.tensor([0.5, -0.3, 0.1])
transforms = fk.get_joint_transform_batch(joint, joint_values)

📋 get_dof_names

Module Name: GBC.utils.base.base_fk.RobotKinematics.get_dof_names

Definition:

def get_dof_names(self)

📥 Input:

None (self)

📤 Output:

dof_names (list): List of joint names in DoF order

🔧 Functionality: Returns the names of degrees of freedom (DoF) for the robot in the order expected by action vectors.

💡 Example Usage:

dof_names = fk.get_dof_names()
print(f"DoF order: {dof_names}")

📏 get_dof_limits

Module Name: GBC.utils.base.base_fk.RobotKinematics.get_dof_limits

Definition:

def get_dof_limits(self)

📥 Input:

None (self)

📤 Output:

min_vals (torch.Tensor): Minimum values for each joint, shape (num_dofs,)
max_vals (torch.Tensor): Maximum values for each joint, shape (num_dofs,)

🔧 Functionality: Returns the joint limits for all degrees of freedom as defined in the URDF file. Sets limits to [-inf, inf] for joints without defined limits.

⚠️ Exceptions:

ValueError: If joint map not found (URDF not loaded)

💡 Example Usage:

min_limits, max_limits = fk.get_dof_limits()
print(f"Joint limits: {min_limits} to {max_limits}")

🔗 get_dof_related_link_names

Module Name: GBC.utils.base.base_fk.RobotKinematics.get_dof_related_link_names

Definition:

def get_dof_related_link_names(self)

📥 Input:

None (self)

📤 Output:

dof_related_link_names (list): List of link names connected to free joints

🔧 Functionality: Returns the names of links that are directly connected to movable (non-fixed) joints, sorted by DoF order.

⚠️ Exceptions:

AssertionError: If number of links connected to free joints is not equal to the number of DoFs

💡 Example Usage:

dof_links = fk.get_dof_related_link_names()
print(f"Links with DoF: {dof_links}")

🎯 get_end_effectors

Module Name: GBC.utils.base.base_fk.RobotKinematics.get_end_effectors

Definition:

def get_end_effectors(self)

📥 Input:

None (self)

📤 Output:

end_effectors (dict): Dictionary mapping end effector link names to joint names

🔧 Functionality: Identifies and returns end effectors of the robot (links with no children or only fixed children). Excludes links containing 'waist' from being considered end effectors.

💡 Example Usage:

end_effectors = fk.get_end_effectors()
print(f"End effectors: {end_effectors}")

🧩 get_certain_part_body_link_and_joint_names

Module Name: GBC.utils.base.base_fk.RobotKinematics.get_certain_part_body_link_and_joint_names

Definition:

def get_certain_part_body_link_and_joint_names(self, root_joints: List[str])

📥 Input:

root_joints (List[str]): List of root joint names

📤 Output:

lj_map (dict): Dictionary mapping link names to joint names for the specified body part

🔧 Functionality: Returns the names of links and joints that are connected to the specified root joints, traversing the kinematic tree from the given root joints.

⚠️ Exceptions:

AssertionError: If base link has no successors or if root joints are not found in base successors

💡 Example Usage:

root_joints = ['left_hip_yaw', 'right_hip_yaw']
leg_links_joints = fk.get_certain_part_body_link_and_joint_names(root_joints)

⚡ forward_kinematics

Module Name: GBC.utils.base.base_fk.RobotKinematics.forward_kinematics

Definition:

def forward_kinematics(self, dof_pos: torch.Tensor, root_trans_offset: Optional[torch.Tensor] = None, root_rot: Optional[torch.Tensor] = None, link_trans_offset: Optional[dict[str, torch.Tensor]] = None)

📥 Input:

dof_pos (torch.Tensor): Joint values, shape (B, N) where N is number of DoFs
root_trans_offset (torch.Tensor, optional): Root translation offsets, shape (B, 3)
root_rot (torch.Tensor, optional): Root rotation angles in XYZ order, shape (B, 3)
link_trans_offset (dict, optional): Dictionary of link names to translation offsets, shape (3,)

📤 Output:

global_transforms (dict): Mapping from link names to global 4x4 transformation matrices, shape (B, 4, 4)

🔧 Functionality: Computes forward kinematics for the entire robot, calculating global transformations for all links based on joint configurations. Traverses the kinematic tree in breadth-first order.

⚠️ Exceptions:

ValueError: If dof_pos is not 2D or doesn't have correct number of columns, or if joint not found in DoF order

💡 Example Usage:

dof_pos = torch.rand(10, fk.num_dofs)  # 10 robot configurations
root_offset = torch.zeros(10, 3)
global_transforms = fk.forward_kinematics(dof_pos, root_trans_offset=root_offset)

🎯 set_target_links

Module Name: GBC.utils.base.base_fk.RobotKinematics.set_target_links

Definition:

def set_target_links(self, target_links: List[str])

📥 Input:

target_links (List[str]): List of link names to track

📤 Output:

None (sets internal state)

🔧 Functionality: Sets the target links for which positions will be computed during forward kinematics.

💡 Example Usage:

fk.set_target_links(["left_hand", "right_hand", "head"])

🚀 forward

Module Name: GBC.utils.base.base_fk.RobotKinematics.forward

Definition:

def forward(self, dof_pos: torch.Tensor, root_trans_offset: Optional[torch.Tensor] = None, root_rot: Optional[torch.Tensor] = None, link_trans_offset: Optional[dict[str, torch.Tensor]] = None)

📥 Input:

dof_pos (torch.Tensor): Joint values, shape (B, num_dof)
root_trans_offset (torch.Tensor, optional): Root translation offsets, shape (B, 3)
root_rot (torch.Tensor, optional): Root rotation angles in XYZ order, shape (B, 3)
link_trans_offset (dict, optional): Link-specific translation offsets, shape (3,)

📤 Output:

target_link_positions (torch.Tensor): Positions of target links, shape (B, num_target_links, 3)

🔧 Functionality: Performs forward kinematics and returns only the positions of the previously set target links.

⚠️ Exceptions:

ValueError: If target links not set or if link has not been transformed yet

💡 Example Usage:

fk.set_target_links(["left_foot", "right_foot"])
dof_pos = torch.rand(5, fk.num_dofs)
foot_positions = fk.forward(dof_pos)  # Shape: (5, 2, 3)

🔄 inverse_kinematics

Module Name: GBC.utils.base.base_fk.RobotKinematics.inverse_kinematics

Definition:

def inverse_kinematics(self, target_link_pose: torch.Tensor, root_trans_offset: Optional[torch.Tensor] = None, root_rot: Optional[torch.Tensor] = None, max_iter: int = 1000, tol: float = 1e-6, verbose: bool = False, lr: float = 0.01, log_interval: int = 10)

📥 Input:

target_link_pose (torch.Tensor): Target positions, shape (B, num_target_links, 3)
root_trans_offset (torch.Tensor, optional): Root translation offsets, shape (B, 3)
root_rot (torch.Tensor, optional): Root rotation angles in XYZ order, shape (B, 3)
max_iter (int): Maximum optimization iterations. Default is 1000
tol (float): Convergence tolerance. Default is 1e-6
verbose (bool): Whether to print optimization progress. Default is False
lr (float): Learning rate for optimization. Default is 0.01
log_interval (int): Logging frequency. Default is 10

📤 Output:

actions (torch.Tensor): Computed joint values, shape (B, num_dofs)

🔧 Functionality: Solves inverse kinematics using gradient-based optimization with Adam optimizer. Includes joint limit constraints and action violation penalties.

⚠️ Exceptions:

ValueError: If number of target links does not match target link pose dimensions

💡 Example Usage:

fk.set_target_links(["left_hand", "right_hand"])
target_poses = torch.tensor([[[0.3, 0.2, 0.8], [-0.3, 0.2, 0.8]]])  # Both hands at specific positions
joint_solution = fk.inverse_kinematics(target_poses, verbose=True)

📊 visualize

Module Name: GBC.utils.base.base_fk.RobotKinematics.visualize

Definition:

def visualize(self)

📥 Input:

None (self)

📤 Output:

None (displays plot)

🔧 Functionality: Visualizes the kinematic tree structure using NetworkX and Matplotlib.

💡 Example Usage:

fk.visualize()  # Shows kinematic tree graph

🎨 visualize_joints

Module Name: GBC.utils.base.base_fk.RobotKinematics.visualize_joints

Definition:

def visualize_joints(self, show: bool = True, axis_length: float = 0.2, default_batch: int = 0, bm: Optional[torch.Tensor] = None, title_prefix: str = "", save_path: str = None, title_override: Optional[str] = None, fixed_axis: bool = False)

📥 Input:

show (bool): Whether to display the plot immediately. Default is True
axis_length (float): Length of orientation axes arrows. Default is 0.2
default_batch (int): Batch index to visualize. Default is 0
bm (torch.Tensor, optional): Body model vertices for comparison, shape (B, N, 3)
title_prefix (str): Prefix for plot title. Default is ""
save_path (str, optional): Path to save the plot. Default is None
title_override (str, optional): Override plot title. Default is None
fixed_axis (bool): Whether to fix plot axis ranges. Default is False

📤 Output:

None (displays/saves 3D plot)

🔧 Functionality: Creates a 3D visualization of joint positions and orientations after forward kinematics computation. Shows target links in blue, other joints in black, and optionally overlays body model vertices.

💡 Example Usage:

dof_pos = torch.rand(1, fk.num_dofs)
fk.forward_kinematics(dof_pos)
fk.visualize_joints(
    title_prefix="Robot Pose - ",
    save_path="./plots",
    show=True
)

🎭 Complete Usage Example

Here's a comprehensive example showcasing the full workflow of the RobotKinematics class:

# 🏗️ Initialize robot kinematics
fk = RobotKinematics(
    urdf_path="/path/to/robot.urdf", 
    device='cuda'
)

# 🎯 Set target links for end effector tracking
fk.set_target_links(["left_hand", "right_hand", "head"])

# 🎲 Generate random joint configurations
batch_size = 10
dof_pos = torch.rand(batch_size, fk.num_dofs, device='cuda') * 0.5

# ⚡ Perform forward kinematics
end_effector_positions = fk.forward(dof_pos)
print(f"End effector positions shape: {end_effector_positions.shape}")

# 🎨 Visualize robot configuration
fk.visualize_joints(default_batch=0, show=True)

# 🔄 Solve inverse kinematics
target_positions = torch.tensor([[[0.3, 0.2, 0.8], [-0.3, 0.2, 0.8], [0.0, 0.0, 1.2]]], device='cuda')
joint_solution = fk.inverse_kinematics(target_positions, verbose=True)

💡 Pro Tip: For optimal performance with large batch sizes, ensure your robot URDF is well-structured and consider using CUDA when available!

⚠️ Note: Always check joint limits using get_dof_limits() before setting target poses to avoid unreachable configurations.

Class Introduction​

Class Dependencies​

Class Functions​

🏗️ __init__​

🔄 to​

🎯 find_base_link​

🔗 get_joint_names​

🧩 get_link_names​

🌳 build_kinematic_tree​

🔄 rotation_matrix_batch​

⚙️ get_joint_transform_batch​

📋 get_dof_names​

📏 get_dof_limits​

🔗 get_dof_related_link_names​

🎯 get_end_effectors​

🧩 get_certain_part_body_link_and_joint_names​

⚡ forward_kinematics​

🎯 set_target_links​

🚀 forward​

🔄 inverse_kinematics​

📊 visualize​

🎨 visualize_joints​

🎭 Complete Usage Example​

Class Introduction

Class Dependencies

Class Functions

🏗️ init

🔄 to

🎯 find_base_link

🔗 get_joint_names

🧩 get_link_names

🌳 build_kinematic_tree

🔄 rotation_matrix_batch

⚙️ get_joint_transform_batch

📋 get_dof_names

📏 get_dof_limits

🔗 get_dof_related_link_names

🎯 get_end_effectors

🧩 get_certain_part_body_link_and_joint_names

⚡ forward_kinematics

🎯 set_target_links

🚀 forward

🔄 inverse_kinematics

📊 visualize

🎨 visualize_joints

🎭 Complete Usage Example