PyRoki

Before using PyRoki, be aware of the following technical limitations:

Performance & Execution

• JIT Overhead: Because it uses JAX, JIT compilation is triggered on the first run or whenever input shapes change (e.g., changing the number of targets or obstacles). To avoid frequent recompilation, you can pre-pad arrays to vectorize over inputs with different shapes.
• Collision Speed: Collision performance may be slower than specialized toolkits like CuRobo in collision-heavy scenarios.

Modeling Constraints

• Joint Types: Supports only revolute, continuous, prismatic, and fixed joints. Other URDF joint types are treated as fixed.
• Collision Geometry: Limited to sphere, capsule, halfspace, and heightmap. Mesh collisions are approximated using capsules.
• Kinematic Structures: Supports only kinematic trees; closed-loop mechanisms and parallel manipulators are not supported.

Planning

• No Sampling-based Planners: PyRoki does not include sampling-based planners like graphs or trees.

This example demonstrates how to retarget motion from a MANO hand model to a Shadow Hand robot, incorporating costs to maintain contact with an object.

Prerequisites

• Unzip the Shadow Hand URDF at assets/hand_retargeting/shadowhand_urdf.zip.
• The example relies on pyroki_snippets implementation details, so ensure you have cloned the full PyRoki repository.

Retargeting Workflow

1. Load Robot: Load the Shadow Hand URDF using yourdfpy and initialize a pk.Robot instance.
2. Map Joints: Use get_mapping_from_mano_to_shadow(robot) to find the relationship between MANO and Shadow Hand joints, and create_conn_tree(robot, shadow_link_idx) to create a mask for connected joints.
3. Prepare Data: Load source motion (keypoints), contact information (points and joint indices), and object meshes.
4. Configure Weights: Define RetargetingWeights to control the optimization behavior:
   • local_alignment: Matches relative joint/keypoint positions and angles.
   • global_alignment: Matches keypoint positions to the robot in the world frame.
   • contact: Weight for maintaining contact between the robot and the object.
   • contact_margin: Threshold to stop penalizing contact when the robot is sufficiently close.
   • joint_smoothness: Weight for joint velocity smoothness.
   • root_smoothness: Weight for robot root translation smoothness.
5. Solve: Call solve_retargeting with the robot, target keypoints, joint mappings, masks, and contact data to obtain the world root transforms and joint configurations.

Implementation Note

The retargeting is solved as a least-squares problem using jaxls, optimizing for local alignment, scale regularization, smoothness, and contact constraints.

# Example weight configuration
def main():
    # ... (loading robot and data) ...
    
    default_weights = RetargetingWeights(
        local_alignment=10.0,
        global_alignment=1.0,
        contact=5.0,
        contact_margin=0.01,
        joint_smoothness=2.0,
        root_smoothness=2.0,
    )

    # Solve the retargeting problem
    Ts_world_root, joints = solve_retargeting(
        robot=robot,
        target_keypoints=keypoints,
        shadow_hand_link_retarget_indices=shadow_link_idx,
        mano_joint_retarget_indices=mano_joint_idx,
        mano_mask=mano_mask,
        contact_points_per_frame=jnp.array(padded_contact_points_per_frame),
        contact_indices_per_frame=jnp.array(padded_contact_indices_per_frame),
        contact_mask=jnp.array(padded_contact_mask),
        weights=default_weights,
    )

The most common way to define costs in pyroki is using automatic differentiation (autodiff). You can define a cost function by using the @Cost.create_factory decorator. This allows the library to automatically compute gradients for optimization tasks.

When defining a cost function for pose matching, you typically compute the forward kinematics to find the actual link pose, then calculate the difference between the actual pose and the target pose using position and orientation residuals.

@Cost.create_factory
def pose_cost(
    vals: VarValues,
    robot: Robot,
    joint_var: Var[Array],
    target_pose: jaxlie.SE3,
    target_link_index: Array,
    pos_weight: Array | float,
    ori_weight: Array | float,
) -> Array:
    """Computes the residual for matching link poses to target poses."""
    assert target_link_index.dtype == jnp.int32
    joint_cfg = vals[joint_var]
    Ts_link_world = robot.forward_kinematics(joint_cfg)
    pose_actual = jaxlie.SE3(Ts_link_world[..., target_link_index, :])

    # Position residual = position error * weight
    pos_residual = (pose_actual.translation() - target_pose.translation()) * pos_weight
    # Orientation residual = log(actual_inv * target) * weight
    ori_residual = (pose_actual.rotation().inverse() @ target_pose.rotation()).log() * ori_weight

    return jnp.concatenate([pos_residual, ori_residual]).flatten()

This example demonstrates how to retarget motion from a source (like SMPL keypoints) to a humanoid robot (like the G1). The process involves:

1. Loading Robot and Motion Data: Initialize the robot using pk.Robot.from_urdf(urdf) and load source keypoints, foot contact information, and a heightmap.
2. Handling Terrain: Use pk.collision.Heightmap to represent the environment and project foot keypoints onto it using heightmap.project_points(keypoints).
3. Defining Retargeting Logic: The core logic uses a least-squares optimization problem (jaxls.LeastSquaresProblem) that minimizes several costs:
   • Local Alignment: Matches relative joint/keypoint positions and angles (vectors) between the source and the robot.
   • Global Alignment: Aligns the overall keypoint positions to the robot in the world frame.
   • Smoothness: Uses pk.costs.smoothness_cost to ensure temporal continuity in joint movements.
   • Rest Pose: Uses pk.costs.rest_cost to encourage the robot to stay near its default configuration, with higher weights for specific joints (e.g., hip/waist yaw) to ensure natural motion.
   • Limits: Uses pk.costs.limit_constraint to respect robot joint limits.
4. Interactive Tuning: The example uses pk.viewer.WeightTuner integrated with viser to allow real-time adjustment of alignment weights during visualization.

# Core retargeting solver call pattern
Ts_world_root, joints = solve_retargeting(
    robot=robot,
    target_keypoints=smpl_keypoints,
    smpl_joint_retarget_indices=smpl_joint_retarget_indices,
    g1_joint_retarget_indices=g1_joint_retarget_indices,
    smpl_mask=smpl_mask,
    weights=weights.get_weights(),
)

You can perform online trajectory planning that accounts for both robot and environmental collisions. This involves defining a robot model, collision geometries (like HalfSpace or Sphere), and an IK target. The planning loop typically uses a solver to find a trajectory that moves the robot's target link to a desired position and orientation while avoiding the specified collision objects.

To run these examples, ensure you have cloned the PyRoki repository to access pyroki_snippets implementation details used in the demonstration.

import time
import numpy as np
import pyroki as pk
import viser
from pyroki.collision import HalfSpace, RobotCollision, Sphere
from robot_descriptions.loaders.yourdfpy import load_robot_description
from viser.extras import ViserUrdf

import pyroki_snippets as pks

def main():
    # 1. Setup Robot and Collisions
    urdf = load_robot_description("panda_description")
    target_link_name = "panda_hand"
    robot = pk.Robot.from_urdf(urdf)
    robot_coll = RobotCollision.from_urdf(urdf)
    
    plane_coll = HalfSpace.from_point_and_normal(
        np.array([0.0, 0.0, 0.0]), np.array([0.0, 0.0, 1.0])
    )
    sphere_coll = Sphere.from_center_and_radius(
        np.array([0.0, 0.0, 0.0]), np.array([0.05])
    )

    # 2. Planning Parameters
    len_traj, dt = 5, 0.1

    # 3. Visualization Setup (Viser)
    server = viser.ViserServer()
    urdf_vis = ViserUrdf(server, urdf, root_node_name="/robot")
    ik_target_handle = server.scene.add_transform_controls(
        "/ik_target", scale=0.2, position=(0.3, 0.0, 0.5), wxyz=(0, 0, 1, 0)
    )
    sphere_handle = server.scene.add_transform_controls(
        "/obstacle", scale=0.2, position=(0.4, 0.3, 0.4)
    )

    # 4. Planning Loop
    sol_traj = np.array(
        robot.joint_var_cls.default_factory()[None].repeat(len_traj, axis=0)
    )
    
    while True:
        # Transform obstacle to world coordinates
        sphere_coll_world_current = sphere_coll.transform_from_wxyz_position(
            wxyz=np.array(sphere_handle.wxyz),
            position=np.array(sphere_handle.position),
        )

        world_coll_list = [plane_coll, sphere_coll_world_current]

        # Solve online planning
        sol_traj, sol_pos, sol_wxyz = pks.solve_online_planning(
            robot=robot,
            robot_coll=robot_coll,
            world_coll=world_coll_list,
            target_link_name=target_link_name,
            target_position=np.array(ik_target_handle.position),
            target_wxyz=np.array(ik_target_handle.wxyz),
            timesteps=len_traj,
            dt=dt,
            start_cfg=sol_traj[0],
            prev_sols=sol_traj,
        )
        
        # Update visualization
        urdf_vis.update_cfg(sol_traj[0])

if __name__ == "__main__":
    main()