├── .gitignore ├── LICENSE.txt ├── README.md ├── examples └── example.py ├── ikflow ├── __init__.py ├── config.py ├── evaluation_utils.py ├── ikflow_solver.py ├── math_utils.py ├── model.py ├── model_descriptions.yaml ├── model_loading.py ├── supporting_types.py ├── thirdparty │ └── ranger │ │ ├── __init__.py │ │ ├── ranger.py │ │ ├── ranger913A.py │ │ └── rangerqh.py ├── training │ ├── lt_data.py │ ├── lt_model.py │ └── training_utils.py ├── utils.py ├── visualization_resources │ └── klampt_resources │ │ ├── objects │ │ ├── README.md │ │ ├── block.obj │ │ ├── block_small.obj │ │ ├── centeredcube.off │ │ ├── claraio_downloads │ │ │ ├── mug.mlp │ │ │ ├── mug.mtl │ │ │ ├── mug.obj │ │ │ └── mug.off │ │ ├── convert.sh │ │ ├── cube.geom │ │ ├── cube.off │ │ ├── cylinder.obj │ │ ├── cylinder.off │ │ ├── cylinder_y.off │ │ ├── flat_foot.off │ │ ├── flatblock.obj │ │ ├── flatblock2.obj │ │ ├── flatblock3.obj │ │ ├── flatblock4.obj │ │ ├── flatblock5.obj │ │ ├── mug.obj │ │ ├── mug.obj.mtl │ │ ├── mug.off │ │ ├── mug_klampt.obj │ │ ├── paddle_cube.off │ │ ├── sphere.geom │ │ ├── sphere.obj │ │ ├── sphere.off │ │ ├── sphere_10cm.obj │ │ ├── sphere_5cm.obj │ │ ├── srimugsmooth.obj │ │ ├── srimugsmooth.off │ │ ├── thincube.off │ │ └── urdf_primitives │ │ │ ├── box_ori_center.off │ │ │ ├── box_ori_center.tri │ │ │ ├── cylinder_ori_center.off │ │ │ ├── cylinder_ori_center.tri │ │ │ ├── sphere_ori_center.off │ │ │ └── sphere_ori_center.tri │ │ ├── simulation_test_worlds │ │ ├── bounce_test.xml │ │ ├── cuppile.xml │ │ ├── displaytest.xml │ │ ├── fastdroptest.xml │ │ ├── heavyblocktest.xml │ │ ├── jointtest.xml │ │ ├── plane_10.off │ │ ├── sensortest.xml │ │ ├── slideblock.xml │ │ ├── sphere_slope.xml │ │ ├── stacktest.xml │ │ ├── stiffness_test.xml │ │ └── tx90cuppile.xml │ │ └── terrains │ │ ├── block.off │ │ ├── cube.off │ │ ├── drc_ladder.off │ │ ├── drc_rough_terrain.stl │ │ ├── fractal_terrain_1.off │ │ ├── fractal_terrain_2.off │ │ ├── fractal_terrain_3.off │ │ ├── fractal_terrain_4.off │ │ ├── plane.env │ │ ├── plane.off │ │ ├── plane_for_atlas.off │ │ ├── ridges.off │ │ ├── smallplane.off │ │ └── stair.off └── visualizations.py ├── lint ├── notebooks ├── ik_convergence_analysis.ipynb ├── inference_optimization.ipynb ├── robot_visualizations.ipynb └── solution_refinement_runtime_plotting.ipynb ├── pyproject.toml ├── scripts ├── benchmark_generate_exact_solutions.py ├── benchmark_runtime.py ├── build_dataset.py ├── download_model_from_wandb_checkpoint.py ├── evaluate.py ├── evaluate_inference_speed.py ├── train.py ├── train_from_checkpoint.py ├── visualize.py └── visualize_robot.py ├── test ├── tests ├── evaluation_utils_test.py ├── ikflow_solver_test.py ├── lt_model_test.py ├── model_loading_test.py └── model_test.py └── uv.lock /.gitignore: -------------------------------------------------------------------------------- 1 | # Pycache 2 | __pycache__/ 3 | thirdparty/mdn/mdn/__pycache__/ 4 | thirdparty/mdn/mdn/__pycache__/__init__.cpython-38.pyc 5 | thirdparty/mdn/mdn/__pycache__/models.cpython-38.pyc 6 | 7 | # Virtual Environments 8 | venv/ 9 | venvTemp/ 10 | venv-test/ 11 | 12 | # IDE Config 13 | .vscode/ 14 | 15 | # Third party folders 16 | wandb/ 17 | artifacts/ 18 | 19 | # Unit testing generated files 20 | _tests/models/ 21 | 22 | # Pytl training data 23 | checkpoints/ 24 | training_logs/ 25 | old/ 26 | notes_to_self 27 | notes 28 | 29 | # Datasets 30 | datasets/*/ 31 | 32 | # Slurm log files 33 | *.out 34 | SBATCH_train.sh 35 | SBATCH_resume_train.sh 36 | SBATCH_script.sh 37 | SBATCH_wandb_sweep.sh 38 | 39 | sourcewandbenvvars 40 | 41 | # Jupyter notebook checkpoints 42 | .ipynb_checkpoints/ 43 | lr_scheduler_scrathpad.py 44 | 45 | *.pkl 46 | 47 | dist/ 48 | ikflow.egg-info 49 | build/ 50 | *.pdf 51 | *.png 52 | 53 | .pytest_cache/ 54 | .ruff_cache/ 55 | .venv/ -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | #TODO -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # IKFlow 2 | Normalizing flows for Inverse Kinematics. Open source implementation to the paper ["IKFlow: Generating Diverse Inverse Kinematics Solutions"](https://ieeexplore.ieee.org/abstract/document/9793576) 3 | 4 | [![arxiv.org](https://img.shields.io/badge/cs.RO-%09arXiv%3A2111.08933-red)](https://arxiv.org/abs/2111.08933) 5 | 6 | 7 | Runtime curve for getting *exact* IK solutions for the Franka Panda (maximum positional/rotational error: 1mm, .572 deg) (generated with ` 8 | python scripts/benchmark_generate_exact_solutions.py --model_name=panda__full__lp191_5.25m`): 9 | 10 | ![alt text](../media/exact_ik_runtime__model:panda__full__lp191_5.25m.png?raw=true) 11 | 12 | 13 | ## Setup 14 | 15 | The following section outlines the setup procedures required to run the visualizer that this project uses. The only supported OS is Ubuntu. Visualization may work on Mac and Windows, I haven't tried it though. For Ubuntu, there are different system wide dependencies for `Ubuntu > 21` and `Ubuntu < 21`. For example, `qt5-default` is not in the apt repository for Ubuntu 21.0+ so can't be installed. See https://askubuntu.com/questions/1335184/qt5-default-not-in-ubuntu-21-04. 16 | 17 | Ubuntu >= 21.04 18 | ``` 19 | sudo apt-get install -y qtbase5-dev qtchooser qt5-qmake qtbase5-dev-tools libosmesa6 build-essential qtcreator 20 | export PYOPENGL_PLATFORM=osmesa # this needs to be run every time you run a visualization script in a new terminal - annoying, I know 21 | ``` 22 | Ubuntu <= 20.x.y 23 | 24 | (This includes 20.04 LTS, 18.04 LTS, ...) 25 | ``` 26 | sudo apt-get install -y qt5-default build-essential qtcreator 27 | ``` 28 | 29 | Lastly, install with uv: 30 | ``` bash 31 | git clone https://github.com/jstmn/ikflow.git && cd ikflow 32 | uv sync 33 | uv pip install -e . 34 | ``` 35 | 36 | 37 | ## Getting started 38 | 39 | **> Example 1: Use IKFlow to generate approximate IK solutions for the Franka Panda** 40 | 41 | Evaluate a pretrained IKFlow model for the Franka Panda arm. Note that the value for `model_name` - in this case `panda__full__lp191_5.25m` should match an entry in `model_descriptions.yaml` 42 | ``` 43 | uv run python scripts/evaluate.py --testset_size=500 --model_name=panda__full__lp191_5.25m 44 | ``` 45 | 46 | **> Example 2: Use IKFlow to generate exact IK solutions for the Franka Panda** 47 | 48 | Additional examples are provided in examples/example.py. This file includes examples of collision checking and pose error calculation, among other utilities. 49 | 50 | ``` 51 | ik_solver, _ = get_ik_solver("panda__full__lp191_5.25m") 52 | target_poses = torch.tensor( 53 | [ 54 | [0.25, 0, 0.5, 1, 0, 0, 0], 55 | [0.35, 0, 0.5, 1, 0, 0, 0], 56 | [0.45, 0, 0.5, 1, 0, 0, 0], 57 | ], 58 | device=device, 59 | ) 60 | solutions, _ = ik_solver.generate_exact_ik_solutions(target_poses) 61 | ``` 62 | 63 | 64 | **> Example 3: Visualize the solutions returned by the `fetch_arm__large__mh186_9.25m` model** 65 | 66 | Run the following: 67 | ``` 68 | uv run python scripts/visualize.py --model_name=fetch_arm__large__mh186_9.25m --demo_name=oscillate_target 69 | ``` 70 | ![ikflow solutions for oscillating target pose](../media/ikflow__fetcharm__oscillating-target.gif?raw=true) 71 | 72 | Run an interactive notebook: `jupyter notebook notebooks/robot_visualizations.ipynb` 73 | 74 | 75 | ## Notes 76 | This project uses the `w,x,y,z` format for quaternions. That is all. 77 | 78 | 79 | ## Training new models 80 | 81 | The training code uses [Pytorch Lightning](https://www.pytorchlightning.ai/) to setup and perform the training and [Weights and Biases](https://wandb.ai/) ('wandb') to track training runs and experiments. WandB isn't required for training but it's what this project is designed around. Changing the code to use Tensorboard should be straightforward (so feel free to put in a pull request for this if you want it :)). 82 | 83 | First, create a dataset for the robot: 84 | ``` 85 | uv run python scripts/build_dataset.py --robot_name=panda --training_set_size=25000000 --only_non_self_colliding 86 | ``` 87 | 88 | Then start a training run: 89 | ``` 90 | # Login to wandb account - Only needs to be run once 91 | uv run wandb login 92 | 93 | # Set wandb project name and entity 94 | export WANDB_PROJECT=ikflow 95 | export WANDB_ENTITY= 96 | 97 | uv run python scripts/train.py \ 98 | --robot_name=panda \ 99 | --nb_nodes=12 \ 100 | --batch_size=128 \ 101 | --learning_rate=0.0005 102 | ``` 103 | 104 | ## Common errors 105 | 106 | 1. GLUT font retrieval function when running a visualizer. Run `export PYOPENGL_PLATFORM=osmesa` and then try again. See https://bytemeta.vip/repo/MPI-IS/mesh/issues/66 107 | 108 | ``` 109 | Traceback (most recent call last): 110 | File "visualize.py", line 4, in 111 | from ikflow.visualizations import _3dDemo 112 | File "/home/jstm/Projects/ikflow/utils/visualizations.py", line 10, in 113 | from klampt import vis 114 | File "/home/jstm/Projects/ikflow/venv/lib/python3.8/site-packages/klampt/vis/__init__.py", line 3, in 115 | from .glprogram import * 116 | File "/home/jstm/Projects/ikflow/venv/lib/python3.8/site-packages/klampt/vis/glprogram.py", line 11, in 117 | from .glviewport import GLViewport 118 | File "/home/jstm/Projects/ikflow/venv/lib/python3.8/site-packages/klampt/vis/glviewport.py", line 8, in 119 | from . import gldraw 120 | File "/home/jstm/Projects/ikflow/venv/lib/python3.8/site-packages/klampt/vis/gldraw.py", line 10, in 121 | from OpenGL import GLUT 122 | File "/home/jstm/Projects/ikflow/venv/lib/python3.8/site-packages/OpenGL/GLUT/__init__.py", line 5, in 123 | from OpenGL.GLUT.fonts import * 124 | File "/home/jstm/Projects/ikflow/venv/lib/python3.8/site-packages/OpenGL/GLUT/fonts.py", line 20, in 125 | p = platform.getGLUTFontPointer( name ) 126 | File "/home/jstm/Projects/ikflow/venv/lib/python3.8/site-packages/OpenGL/platform/baseplatform.py", line 350, in getGLUTFontPointer 127 | raise NotImplementedError( 128 | NotImplementedError: Platform does not define a GLUT font retrieval function 129 | ``` 130 | 131 | 2. If you get this error: `tkinter.TclError: no display name and no $DISPLAY environment variable`, add the lines below to the top of `ik_solvers.py` (anywhere before `import matplotlib.pyplot as plt` should work). 132 | ``` python 133 | import matplotlib 134 | matplotlib.use("Agg") 135 | ``` 136 | 137 | 138 | ## Citation 139 | ``` 140 | @ARTICLE{9793576, 141 | author={Ames, Barrett and Morgan, Jeremy and Konidaris, George}, 142 | journal={IEEE Robotics and Automation Letters}, 143 | title={IKFlow: Generating Diverse Inverse Kinematics Solutions}, 144 | year={2022}, 145 | volume={7}, 146 | number={3}, 147 | pages={7177-7184}, 148 | doi={10.1109/LRA.2022.3181374} 149 | } 150 | ``` -------------------------------------------------------------------------------- /examples/example.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from ikflow.utils import set_seed 4 | from ikflow.model_loading import get_ik_solver 5 | from ikflow.config import DEVICE 6 | 7 | import torch 8 | 9 | set_seed() 10 | 11 | 12 | """ Usage 13 | 14 | uv run python examples/example.py --model_name=panda__full__lp191_5.25m 15 | 16 | """ 17 | 18 | if __name__ == "__main__": 19 | parser = argparse.ArgumentParser(prog="minimial example of running IKFlow") 20 | parser.add_argument( 21 | "--model_name", 22 | type=str, 23 | required=True, 24 | help=( 25 | "Name of the saved model ('panda__full__lp191_5.25m' should work). Defined in" 26 | " ikflow/model_descriptions.yaml" 27 | ), 28 | ) 29 | args = parser.parse_args() 30 | 31 | # Build IKFlowSolver and set weights 32 | ik_solver, hyper_parameters = get_ik_solver(args.model_name) 33 | robot = ik_solver.robot 34 | 35 | """SINGLE TARGET-POSE 36 | 37 | The following code is for when you want to run IKFlow on a single target poses. In this example we are getting 38 | number_of_solutions=5 solutions for the target pose. 39 | """ 40 | target_pose = torch.tensor( 41 | [0.5, 0.5, 0.5, 1, 0, 0, 0], device=DEVICE 42 | ) # Note: quaternions format for ikflow is [w x y z] (I have learned this the hard way...) 43 | number_of_solutions = 5 44 | 45 | # -> Get some approximate solutions 46 | solutions, positional_errors, rotational_errors, joint_limits_exceeded, self_colliding, runtime = ( 47 | ik_solver.generate_ik_solutions(target_pose, n=number_of_solutions, return_detailed=True) 48 | ) 49 | print( 50 | "\nGot {} ikflow solutions in {} ms. Solution positional errors = {} (cm)".format( 51 | number_of_solutions, round(runtime * 1000, 3), positional_errors * 100 52 | ) 53 | ) 54 | 55 | # -> Get some exact solutions 56 | solutions, _ = ik_solver.generate_exact_ik_solutions(target_pose.expand((number_of_solutions, 7))) 57 | print( 58 | "Got {} exact ikflow solutions in {} ms. Solution positional errors = {} (cm)".format( 59 | len(solutions), round(runtime * 1000, 3), "(TODO)" 60 | ) 61 | ) 62 | 63 | """MULTIPLE TARGET-POSES 64 | 65 | The following code is for when you want to run IKFlow on multiple target poses at once. 66 | """ 67 | target_poses = torch.tensor( 68 | [ 69 | [0.25, 0, 0.5, 1, 0, 0, 0], 70 | [0.35, 0, 0.5, 1, 0, 0, 0], 71 | [0.45, 0, 0.5, 1, 0, 0, 0], 72 | [0.55, 0, 0.5, 1, 0, 0, 0], 73 | [0.65, 0, 0.5, 1, 0, 0, 0], 74 | ], 75 | device=DEVICE, 76 | ) 77 | 78 | # -> approximate solutions 79 | solutions, positional_errors, rotational_errors, joint_limits_exceeded, self_colliding, runtime = ( 80 | ik_solver.generate_ik_solutions(target_poses, n=len(target_poses), refine_solutions=False, return_detailed=True) 81 | ) 82 | print( 83 | "\nGot {} ikflow solutions in {} ms. The positional error of the solutions = {} (cm)".format( 84 | target_poses.shape[0], round(runtime * 1000, 3), positional_errors * 100 85 | ) 86 | ) 87 | 88 | # -> exact solutions 89 | solutions, _ = ik_solver.generate_exact_ik_solutions(target_poses) 90 | print( 91 | "Got {} exact ikflow solutions in {} ms. The positional error of the solutions = {} (cm)".format( 92 | len(solutions), round(runtime * 1000, 3), "(TODO)" 93 | ) 94 | ) 95 | 96 | print("\nDone") 97 | -------------------------------------------------------------------------------- /ikflow/__init__.py: -------------------------------------------------------------------------------- 1 | from . import config 2 | from . import evaluation_utils 3 | from . import ikflow_solver 4 | from . import math_utils 5 | from . import model_loading 6 | from . import model 7 | from . import supporting_types 8 | from . import utils 9 | from . import visualizations 10 | 11 | __all__ = [ 12 | "config", 13 | "evaluation_utils", 14 | "ikflow_solver", 15 | "math_utils", 16 | "model_loading", 17 | "model", 18 | "supporting_types", 19 | "utils", 20 | "visualizations", 21 | ] 22 | -------------------------------------------------------------------------------- /ikflow/config.py: -------------------------------------------------------------------------------- 1 | """Global configuration""" 2 | 3 | import os 4 | import torch 5 | 6 | from jrl.config import DEVICE 7 | 8 | DEFAULT_TORCH_DTYPE = torch.float32 9 | print(f"ikflow/config.py | Using device: '{DEVICE}'") 10 | 11 | # ~/.cache/ikflow/ 12 | DEFAULT_DATA_DIR = os.path.join(os.path.expanduser("~"), ".cache/ikflow/") 13 | 14 | # 15 | WANDB_CACHE_DIR = os.path.join(DEFAULT_DATA_DIR, "wandb/") 16 | DATASET_DIR = os.path.join(DEFAULT_DATA_DIR, "datasets/") 17 | TRAINING_LOGS_DIR = os.path.join(DEFAULT_DATA_DIR, "training_logs/") 18 | MODELS_DIR = os.path.join(DEFAULT_DATA_DIR, "models/") 19 | 20 | # Dataset tags 21 | DATASET_TAG_NON_SELF_COLLIDING = "non-self-colliding" 22 | 23 | ALL_DATASET_TAGS = [DATASET_TAG_NON_SELF_COLLIDING] 24 | 25 | 26 | # Training 27 | # This is the absolute maximum value that a value on the joint angle side of the network can take on when it is passed 28 | # through the network, when sigmoid activation is enabled. Note this only applies for values that are in columns past 29 | # the columns reserved for the joints. On other terms, this only applies for the columns with indexes greater than 30 | # n_dofs. 31 | SIGMOID_SCALING_ABS_MAX = 1.0 32 | -------------------------------------------------------------------------------- /ikflow/evaluation_utils.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, List, Optional 2 | 3 | import numpy as np 4 | import torch 5 | 6 | from jrl.robot import Robot 7 | from jrl.math_utils import geodesic_distance_between_quaternions 8 | from jrl.config import PT_NP_TYPE 9 | 10 | """ Description of 'SOLUTION_EVALUATION_RESULT_TYPE': 11 | - torch.Tensor: [n] tensor of positional errors of the IK solutions. The error is the L2 norm of the realized poses of 12 | the solutions and the target pose. 13 | - torch.Tensor: [n] tensor of angular errors of the IK solutions. The error is the angular geodesic distance between the 14 | realized orientation of the robot's end effector from the IK solutions and the targe orientation. 15 | - torch.Tensor: [n] tensor of bools indicating whether each IK solutions has exceeded the robots joint limits. 16 | - torch.Tensor: [n] tensor of bools indicating whether each IK solutions is self colliding. 17 | - float: Runtime 18 | """ 19 | SOLUTION_EVALUATION_RESULT_TYPE = Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, float] 20 | 21 | 22 | def _get_target_pose_batch(target_pose: PT_NP_TYPE, n_solutions: int) -> torch.Tensor: 23 | """Return a batch of target poses. If the target_pose is a single pose, then it will be tiled to match the number of 24 | solutions. If the target_pose is a batch of poses, then it will be returned as is. 25 | 26 | Args: 27 | target_pose (PT_NP_TYPE): [7] or [n x 7] the target pose(s) to calculate errors for 28 | n_solutions (int): The number of solutions we are calculating errors for 29 | """ 30 | if target_pose.shape[0] == 7: 31 | if isinstance(target_pose, torch.Tensor): 32 | return target_pose.repeat(n_solutions, 1) 33 | return np.tile(target_pose, (n_solutions, 1)) 34 | return target_pose 35 | 36 | 37 | def pose_errors( 38 | poses_1: PT_NP_TYPE, poses_2: PT_NP_TYPE, acos_epsilon: Optional[float] = None 39 | ) -> Tuple[PT_NP_TYPE, PT_NP_TYPE]: 40 | """Return the positional and rotational angular error between two batch of poses.""" 41 | assert poses_1.shape == poses_2.shape, f"Poses are of different shape: {poses_1.shape} != {poses_2.shape}" 42 | 43 | if isinstance(poses_1, torch.Tensor): 44 | l2_errors = torch.norm(poses_1[:, 0:3] - poses_2[:, 0:3], dim=1) 45 | else: 46 | l2_errors = np.linalg.norm(poses_1[:, 0:3] - poses_2[:, 0:3], axis=1) 47 | angular_errors = geodesic_distance_between_quaternions( 48 | poses_1[:, 3 : 3 + 4], poses_2[:, 3 : 3 + 4], acos_epsilon=acos_epsilon 49 | ) 50 | assert l2_errors.shape == angular_errors.shape 51 | return l2_errors, angular_errors 52 | 53 | 54 | def pose_errors_cm_deg( 55 | poses_1: PT_NP_TYPE, poses_2: PT_NP_TYPE, acos_epsilon: Optional[float] = None 56 | ) -> Tuple[PT_NP_TYPE, PT_NP_TYPE]: 57 | """Return the positional and rotational angular error between two batch of poses in cm and degrees""" 58 | assert poses_1.shape == poses_2.shape, f"Poses are of different shape: {poses_1.shape} != {poses_2.shape}" 59 | l2_errors, angular_errors = pose_errors(poses_1, poses_2, acos_epsilon=acos_epsilon) 60 | if isinstance(poses_1, torch.Tensor): 61 | return 100 * l2_errors, torch.rad2deg(angular_errors) 62 | return 100 * l2_errors, np.rad2deg(angular_errors) 63 | 64 | 65 | def solution_pose_errors( 66 | robot: Robot, solutions: torch.Tensor, target_poses: torch.Tensor 67 | ) -> Tuple[np.ndarray, np.ndarray]: 68 | """Return the L2 and angular errors of calculated ik solutions for a given target_pose. Note: this function expects 69 | multiple solutions but only a single target_pose. All of the solutions are assumed to be for the given target_pose 70 | 71 | Args: 72 | robot (Robot): The Robot which contains the FK function we will use 73 | solutions torch.Tensor: [n x 7] IK solutions for the given target pose 74 | target_pose (np.ndarray): [7] the target pose the IK solutions were generated for 75 | 76 | Returns: 77 | Tuple[np.ndarray, np.ndarray]: The L2, and angular (rad) errors of IK solutions for the given target_pose 78 | """ 79 | assert isinstance(target_poses, torch.Tensor), f"target_poses must be a torch.Tensor (got {type(target_poses)})" 80 | assert isinstance(solutions, torch.Tensor), f"solutions must be a torch.Tensor (got {type(solutions)})" 81 | n_solutions = solutions.shape[0] 82 | if n_solutions >= 1000: 83 | print("Heads up: It may be faster to run solution_pose_errors() with pytorch directly on the cpu/gpu") 84 | 85 | target_poses = _get_target_pose_batch(target_poses, solutions.shape[0]) 86 | ee_pose_ikflow = robot.forward_kinematics(solutions[:, 0 : robot.ndof]) 87 | rot_output = ee_pose_ikflow[:, 3:] 88 | 89 | # Positional Error 90 | l2_errors = torch.norm(ee_pose_ikflow[:, 0:3] - target_poses[:, 0:3], dim=1) 91 | rot_target = target_poses[:, 3:] 92 | assert rot_target.shape == rot_output.shape 93 | 94 | # Surprisingly, this is almost always faster to calculate on the gpu than on the cpu. I would expect the opposite 95 | # for low number of solutions (< 200). 96 | ang_errors = geodesic_distance_between_quaternions(rot_target, rot_output) 97 | return l2_errors, ang_errors 98 | 99 | 100 | def calculate_joint_limits_exceeded(configs: torch.Tensor, joint_limits: List[Tuple[float, float]]) -> torch.Tensor: 101 | """Calculate if the given configs have exceeded the specified joint limits 102 | 103 | Args: 104 | configs (torch.Tensor): [batch x ndof] tensor of robot configurations 105 | joint_limits (List[Tuple[float, float]]): The joint limits for the robot. Should be a list of tuples, where each 106 | tuple contains (lower, upper). 107 | Returns: 108 | torch.Tensor: [batch] tensor of bools indicating if the given configs have exceeded the specified joint limits 109 | """ 110 | toolarge = configs > torch.tensor([x[1] for x in joint_limits], dtype=torch.float32, device=configs.device) 111 | toosmall = configs < torch.tensor([x[0] for x in joint_limits], dtype=torch.float32, device=configs.device) 112 | return torch.logical_or(toolarge, toosmall).any(dim=1) 113 | 114 | 115 | def calculate_self_collisions(robot: Robot, configs: torch.Tensor) -> torch.Tensor: 116 | """Calculate if the given configs will cause the robot to collide with itself 117 | 118 | Args: 119 | robot (Robot): The robot 120 | configs (torch.Tensor): [batch x ndof] tensor of robot configurations 121 | 122 | Returns: 123 | torch.Tensor: [batch] tensor of bools indicating if the given configs will self-collide 124 | """ 125 | is_self_colliding = [robot.config_self_collides(config) for config in configs] 126 | return torch.tensor(is_self_colliding, dtype=torch.bool) 127 | 128 | 129 | # TODO: What type should this return? Right now its a mix of torch.Tensor and np.ndarray 130 | def evaluate_solutions( 131 | robot: Robot, target_poses: PT_NP_TYPE, solutions: torch.Tensor 132 | ) -> SOLUTION_EVALUATION_RESULT_TYPE: 133 | """Return detailed information about a batch of solutions. See the comment at the top of this file for a description 134 | of the SOLUTION_EVALUATION_RESULT_TYPE type - that's the type that this function returns 135 | 136 | Example usage: 137 | l2_errors, ang_errors, joint_limits_exceeded, self_collisions = evaluate_solutions( 138 | robot, ee_pose_target, samples 139 | ) 140 | """ 141 | assert isinstance(target_poses, torch.Tensor), f"target_poses must be a torch.Tensor (got {type(target_poses)})" 142 | assert isinstance(solutions, torch.Tensor), f"solutions must be a torch.Tensor (got {type(solutions)})" 143 | target_poses = _get_target_pose_batch(target_poses, solutions.shape[0]) 144 | l2_errors, angular_errors = solution_pose_errors(robot, solutions, target_poses) 145 | joint_limits_exceeded = calculate_joint_limits_exceeded(solutions, robot.actuated_joints_limits) 146 | self_collisions_respected = calculate_self_collisions(robot, solutions) 147 | return l2_errors, angular_errors, joint_limits_exceeded, self_collisions_respected 148 | -------------------------------------------------------------------------------- /ikflow/math_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | 5 | 6 | def MMD_multiscale(x, y, c_list, a_list, reduce=True): 7 | """Example usage: 8 | MMD_multiscale(x0, x1, rev_kernel_width, reverse_loss_a, reduce=False) 9 | 10 | 11 | Example usage in toy-inverse-kinematics: 12 | 13 | rev_kernel_width = 1.1827009364464547 14 | 15 | `backward_mmd(x0, x1, *y_args)`: 16 | mmd = MMD_multiscale(x0, x1, [c.rev_kernel_width, c.rev_kernel_width, c.rev_kernel_width], [0.2, 1.0, 2.0]) 17 | 18 | `latent_mmd(y0, y1)`: 19 | mmd = MMD_multiscale(y0, y1, [0.1, 0.2, 0.5], [0.5, 1.0, 2.0]) 20 | 21 | """ 22 | xx, yy, xy = torch.mm(x, x.t()), torch.mm(y, y.t()), torch.mm(x, y.t()) 23 | 24 | rx = xx.diag().unsqueeze(0).expand_as(xx) 25 | ry = yy.diag().unsqueeze(0).expand_as(yy) 26 | 27 | dxx = torch.clamp(rx.t() + rx - 2.0 * xx, 0, np.inf) 28 | dyy = torch.clamp(ry.t() + ry - 2.0 * yy, 0, np.inf) 29 | dxy = torch.clamp(rx.t() + ry - 2.0 * xy, 0, np.inf) 30 | 31 | XX = torch.zeros(xx.shape).to(xx.device) 32 | YY = torch.zeros(xx.shape).to(xx.device) 33 | XY = torch.zeros(xx.shape).to(xx.device) 34 | 35 | for C, a in zip(c_list, a_list): 36 | XX += C**a * ((C + dxx) / a) ** -a 37 | YY += C**a * ((C + dyy) / a) ** -a 38 | XY += C**a * ((C + dxy) / a) ** -a 39 | 40 | if reduce: 41 | return torch.mean(XX + YY - 2.0 * XY) 42 | else: 43 | return XX + YY - 2.0 * XY 44 | -------------------------------------------------------------------------------- /ikflow/model_descriptions.yaml: -------------------------------------------------------------------------------- 1 | # Notes: 2 | # 1. 'tpm' is short for 'top performing model'. 3 | # 2. 'nsc' is short for 'non self colliding'. This means the dataset used for training has only non self-colliding configs. 4 | # These models will typically return 3-6% self colliding solutions 5 | 6 | # =========================================================== 7 | # === Panda 8 | # 9 | 10 | panda__full__lp191_5.25m: 11 | nb_nodes: 12 12 | dim_latent_space: 7 13 | coeff_fn_config: 3 14 | coeff_fn_internal_size: 1024 15 | rnvp_clamp: 2.5 16 | robot_name: 'panda' 17 | model_weights_url: 'https://storage.googleapis.com/ikflow_models/panda/panda__lyric-puddle-191__global_step%3D5.25M.pkl' 18 | 19 | panda_lite_tpm: 20 | nb_nodes: 6 21 | dim_latent_space: 7 22 | coeff_fn_config: 3 23 | coeff_fn_internal_size: 1024 24 | rnvp_clamp: 2.5 25 | robot_name: 'panda' 26 | model_weights_url: 'https://storage.googleapis.com/ikflow_models/panda/panda_arm-young-night-84.pkl' 27 | 28 | 29 | # =========================================================== 30 | # === Fetch 31 | # 32 | 33 | fetch_full_temp_nsc_tpm: 34 | nb_nodes: 12 35 | dim_latent_space: 8 36 | coeff_fn_config: 3 37 | coeff_fn_internal_size: 1024 38 | rnvp_clamp: 2.5 39 | robot_name: 'fetch' 40 | model_weights_url: 'https://storage.googleapis.com/ikflow_models/fetch/fetch__sleek-microwave-65__global_step%3D9.25M.pkl' 41 | 42 | fetch__large__ns183_9.75m: 43 | nb_nodes: 16 44 | dim_latent_space: 8 45 | coeff_fn_config: 3 46 | coeff_fn_internal_size: 1024 47 | rnvp_clamp: 2.5 48 | robot_name: 'fetch' 49 | model_weights_url: 'https://storage.googleapis.com/ikflow_models/fetch/fetch__northern-sea-183__global_step%3D9.75M.pkl' 50 | 51 | 52 | # =========================================================== 53 | # === FetchArm 54 | # 55 | 56 | fetch_arm__large__mh186_9.25m: 57 | nb_nodes: 16 58 | dim_latent_space: 10 59 | coeff_fn_config: 3 60 | coeff_fn_internal_size: 1024 61 | rnvp_clamp: 2.5 62 | robot_name: 'fetch_arm' 63 | model_weights_url: 'https://storage.googleapis.com/ikflow_models/fetch_arm/fetch_arm__major-hill-186__global_step%3D9.25M.pkl' 64 | 65 | 66 | 67 | 68 | 69 | # TODO: add weights back once Iiwa7 is added to jrl (currently ommited because missing capsules) 70 | 71 | # =========================================================== 72 | # === Iiwa7 73 | # 74 | 75 | # iiwa7_full_temp_nsc_tpm: 76 | # nb_nodes: 12 77 | # dim_latent_space: 10 78 | # coeff_fn_config: 3 79 | # coeff_fn_internal_size: 1024 80 | # rnvp_clamp: 2.5 81 | # robot_name: 'iiwa7' 82 | # model_weights_url: 'https://storage.googleapis.com/ikflow_models/iiwa7/iiwa7__serene-snowball-78__global_step%3D10.75M.pkl' 83 | 84 | 85 | # =========================================================== 86 | # === Rizon4 87 | # 88 | 89 | # TODO: Upload a better model when one becomes available 90 | rizon4__snowy-brook-208__global_step=2.75M: 91 | nb_nodes: 12 92 | dim_latent_space: 7 93 | coeff_fn_config: 3 94 | coeff_fn_internal_size: 1024 95 | rnvp_clamp: 2.5 96 | robot_name: 'rizon4' 97 | model_weights_url: 'https://storage.googleapis.com/ikflow_models/rizon4/rizon4__snowy-brook-208__global_step%3D2.75M.pkl' 98 | -------------------------------------------------------------------------------- /ikflow/model_loading.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, Optional, Dict 2 | import yaml 3 | import os 4 | from urllib.request import urlretrieve 5 | 6 | from jrl.robot import Robot 7 | from jrl.robots import get_robot 8 | from ikflow.utils import safe_mkdir, get_filepath 9 | from ikflow.ikflow_solver import IKFlowSolver 10 | from ikflow.model import IkflowModelParameters 11 | from ikflow.config import MODELS_DIR 12 | 13 | with open(get_filepath("model_descriptions.yaml"), "r") as f: 14 | MODEL_DESCRIPTIONS = yaml.safe_load(f) 15 | 16 | 17 | def _assert_model_downloaded_correctly(filepath: str): 18 | filesize_bytes = os.path.getsize(filepath) 19 | filesize_mb = filesize_bytes * 0.000001 # Appriximate MB 20 | assert filesize_mb > 10, ( 21 | f"Model weights saved at '{filepath}' has only {filesize_mb} MB - was it saved correctly? Tip: Check that the" 22 | " file is publically available on GCP." 23 | ) 24 | 25 | 26 | def get_all_model_names() -> Tuple[str]: 27 | """Return a tuple of the model names""" 28 | return tuple(MODEL_DESCRIPTIONS.keys()) 29 | 30 | 31 | def download_model(url: str, download_dir: Optional[str] = None) -> str: 32 | """Download the model at the url `url` to the given directory. download_dir defaults to MODELS_DIR 33 | 34 | Args: 35 | model_url (str): _description_ 36 | download_dir (str): _description_ 37 | """ 38 | if download_dir is None: 39 | download_dir = MODELS_DIR 40 | safe_mkdir(download_dir) 41 | assert os.path.isdir(download_dir), f"Download directory {download_dir} does not exist" 42 | model_name = model_filename(url) 43 | save_filepath = os.path.join(download_dir, model_name) 44 | if os.path.isfile(save_filepath): 45 | _assert_model_downloaded_correctly(save_filepath) 46 | return save_filepath 47 | urlretrieve(url, filename=save_filepath) # returns (filename, headers) 48 | _assert_model_downloaded_correctly(save_filepath) 49 | return save_filepath 50 | 51 | 52 | def model_filename(url: str) -> str: 53 | """Return the model alias given the url of the model (stored in google cloud presumably). 54 | 55 | Example input/output: https://storage.googleapis.com/ikflow_models/atlas_desert-sweep-6.pkl -> atlas_desert-sweep-6.pkl 56 | """ 57 | return url.split("/")[-1] 58 | 59 | 60 | def get_ik_solver( 61 | model_name: str, robot: Optional[Robot] = None, compile_model: Optional[Dict] = None 62 | ) -> Tuple[IKFlowSolver, IkflowModelParameters]: 63 | """Build and return the `IKFlowSolver` for the given model. The input `model_name` should match and index in `model_descriptions.yaml` 64 | 65 | Returns: 66 | Tuple[IKFlowSolver, IkflowModelParameters]: A `IKFlowSolver` solver and the corresponding 67 | `IkflowModelParameters` parameters object 68 | """ 69 | assert model_name in MODEL_DESCRIPTIONS, f"Model name '{model_name}' not found in model descriptions" 70 | model_weights_url = MODEL_DESCRIPTIONS[model_name]["model_weights_url"] 71 | robot_name = MODEL_DESCRIPTIONS[model_name]["robot_name"] 72 | hparams = MODEL_DESCRIPTIONS[model_name] 73 | assert isinstance(robot_name, str), f"robot_name must be a string, got {type(robot_name)}" 74 | assert isinstance(hparams, dict), f"model_hyperparameters must be a Dict, got {type(hparams)}" 75 | 76 | model_weights_filepath = download_model(model_weights_url) 77 | assert os.path.isfile(model_weights_filepath), ( 78 | f"File '{model_weights_filepath}' was not found. Unable to load model weights" 79 | ) 80 | 81 | if robot is None: 82 | robot = get_robot(robot_name) 83 | assert robot.name == robot_name 84 | 85 | # Build IKFlowSolver and set weights 86 | hyper_parameters = IkflowModelParameters() 87 | hyper_parameters.__dict__.update(hparams) 88 | ik_solver = IKFlowSolver(hyper_parameters, robot, compile_model=compile_model) 89 | ik_solver.load_state_dict(model_weights_filepath) 90 | return ik_solver, hyper_parameters 91 | 92 | 93 | if __name__ == "__main__": 94 | ik_solver, hyper_parameters = get_ik_solver("panda_tpm") 95 | -------------------------------------------------------------------------------- /ikflow/supporting_types.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jstmn/ikflow/19634d436eb805f22578b4c5b6526f96d91d0323/ikflow/supporting_types.py -------------------------------------------------------------------------------- /ikflow/thirdparty/ranger/__init__.py: -------------------------------------------------------------------------------- 1 | from .ranger import Ranger 2 | from .ranger913A import RangerVA 3 | from .rangerqh import RangerQH 4 | -------------------------------------------------------------------------------- /ikflow/thirdparty/ranger/ranger.py: -------------------------------------------------------------------------------- 1 | #Ranger deep learning optimizer - RAdam + Lookahead combined. 2 | #https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer 3 | 4 | #Ranger has now been used to capture 12 records on the FastAI leaderboard. 5 | 6 | #This version = 9.3.19 7 | 8 | #Credits: 9 | #RAdam --> https://github.com/LiyuanLucasLiu/RAdam 10 | #Lookahead --> rewritten by lessw2020, but big thanks to Github @LonePatient and @RWightman for ideas from their code. 11 | #Lookahead paper --> MZhang,G Hinton https://arxiv.org/abs/1907.08610 12 | 13 | #summary of changes: 14 | #full code integration with all updates at param level instead of group, moves slow weights into state dict (from generic weights), 15 | #supports group learning rates (thanks @SHolderbach), fixes sporadic load from saved nn_model issues. 16 | #changes 8/31/19 - fix references to *self*.N_sma_threshold; 17 | #changed eps to 1e-5 as better default than 1e-8. 18 | 19 | import math 20 | import torch 21 | from torch.optim.optimizer import Optimizer, required 22 | import itertools as it 23 | 24 | 25 | 26 | class Ranger(Optimizer): 27 | 28 | def __init__(self, params, lr=1e-3, alpha=0.5, k=6, N_sma_threshhold=5, betas=(.95,0.999), eps=1e-5, weight_decay=0): 29 | #parameter checks 30 | if not 0.0 <= alpha <= 1.0: 31 | raise ValueError(f'Invalid slow update rate: {alpha}') 32 | if not 1 <= k: 33 | raise ValueError(f'Invalid lookahead steps: {k}') 34 | if not lr > 0: 35 | raise ValueError(f'Invalid Learning Rate: {lr}') 36 | if not eps > 0: 37 | raise ValueError(f'Invalid eps: {eps}') 38 | 39 | #parameter comments: 40 | # beta1 (momentum) of .95 seems to work better than .90... 41 | #N_sma_threshold of 5 seems better in testing than 4. 42 | #In both cases, worth testing on your trdata (.90 vs .95, 4 vs 5) to make sure which works best for you. 43 | 44 | #prep defaults and init torch.optim base 45 | defaults = dict(lr=lr, alpha=alpha, k=k, step_counter=0, betas=betas, N_sma_threshhold=N_sma_threshhold, eps=eps, weight_decay=weight_decay) 46 | super().__init__(params,defaults) 47 | 48 | #adjustable threshold 49 | self.N_sma_threshhold = N_sma_threshhold 50 | 51 | #now we can get to work... 52 | #removed as we now use step from RAdam...no need for duplicate step counting 53 | #for group in self.param_groups: 54 | # group["step_counter"] = 0 55 | #print("group step counter init") 56 | 57 | #look ahead params 58 | self.alpha = alpha 59 | self.k = k 60 | 61 | #radam buffer for state 62 | self.radam_buffer = [[None,None,None] for ind in range(10)] 63 | 64 | #self.first_run_check=0 65 | 66 | #lookahead weights 67 | #9/2/19 - lookahead param tensors have been moved to state storage. 68 | #This should resolve issues with load/save where weights were left in GPU memory from first load, slowing down future runs. 69 | 70 | #self.slow_weights = [[p.clone().detach() for p in group['params']] 71 | # for group in self.param_groups] 72 | 73 | #don't use grad for lookahead weights 74 | #for w in it.chain(*self.slow_weights): 75 | # w.requires_grad = False 76 | 77 | def __setstate__(self, state): 78 | print("set state called") 79 | super(Ranger, self).__setstate__(state) 80 | 81 | 82 | def step(self, closure=None): 83 | loss = None 84 | #note - below is commented out b/c I have other work that passes back the loss as a float, and thus not a callable closure. 85 | #Uncomment if you need to use the actual closure... 86 | 87 | #if closure is not None: 88 | #loss = closure() 89 | 90 | #Evaluate averages and grad, update param tensors 91 | for group in self.param_groups: 92 | 93 | for p in group['params']: 94 | if p.grad is None: 95 | continue 96 | grad = p.grad.data.float() 97 | if grad.is_sparse: 98 | raise RuntimeError('Ranger optimizer does not support sparse gradients') 99 | 100 | p_data_fp32 = p.data.float() 101 | 102 | state = self.state[p] #get state dict for this param 103 | 104 | if len(state) == 0: #if first time to run...init dictionary with our desired entries 105 | #if self.first_run_check==0: 106 | #self.first_run_check=1 107 | #print("Initializing slow buffer...should not see this at load from saved nn_model!") 108 | state['step'] = 0 109 | state['exp_avg'] = torch.zeros_like(p_data_fp32) 110 | state['exp_avg_sq'] = torch.zeros_like(p_data_fp32) 111 | 112 | #look ahead weight storage now in state dict 113 | state['slow_buffer'] = torch.empty_like(p.data) 114 | state['slow_buffer'].copy_(p.data) 115 | 116 | else: 117 | state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32) 118 | state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32) 119 | 120 | #begin computations 121 | exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] 122 | beta1, beta2 = group['betas'] 123 | 124 | #compute variance mov avg 125 | exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) 126 | #compute mean moving avg 127 | exp_avg.mul_(beta1).add_(1 - beta1, grad) 128 | 129 | state['step'] += 1 130 | 131 | 132 | buffered = self.radam_buffer[int(state['step'] % 10)] 133 | if state['step'] == buffered[0]: 134 | N_sma, step_size = buffered[1], buffered[2] 135 | else: 136 | buffered[0] = state['step'] 137 | beta2_t = beta2 ** state['step'] 138 | N_sma_max = 2 / (1 - beta2) - 1 139 | N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t) 140 | buffered[1] = N_sma 141 | if N_sma > self.N_sma_threshhold: 142 | step_size = math.sqrt((1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / (N_sma_max - 2)) / (1 - beta1 ** state['step']) 143 | else: 144 | step_size = 1.0 / (1 - beta1 ** state['step']) 145 | buffered[2] = step_size 146 | 147 | if group['weight_decay'] != 0: 148 | p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32) 149 | 150 | if N_sma > self.N_sma_threshhold: 151 | denom = exp_avg_sq.sqrt().add_(group['eps']) 152 | p_data_fp32.addcdiv_(-step_size * group['lr'], exp_avg, denom) 153 | else: 154 | p_data_fp32.add_(-step_size * group['lr'], exp_avg) 155 | 156 | p.data.copy_(p_data_fp32) 157 | 158 | #integrated look ahead... 159 | #we do it at the param level instead of group level 160 | if state['step'] % group['k'] == 0: 161 | slow_p = state['slow_buffer'] #get access to slow param tensor 162 | slow_p.add_(self.alpha, p.data - slow_p) #(fast weights - slow weights) * alpha 163 | p.data.copy_(slow_p) #copy interpolated weights to RAdam param tensor 164 | 165 | return loss 166 | -------------------------------------------------------------------------------- /ikflow/thirdparty/ranger/ranger913A.py: -------------------------------------------------------------------------------- 1 | # Ranger deep learning optimizer - RAdam + Lookahead + calibrated adaptive LR combined. 2 | # https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer 3 | 4 | # Ranger has now been used to capture 12 records on the FastAI leaderboard. 5 | 6 | #This version = 9.13.19A 7 | 8 | #Credits: 9 | #RAdam --> https://github.com/LiyuanLucasLiu/RAdam 10 | #Lookahead --> rewritten by lessw2020, but big thanks to Github @LonePatient and @RWightman for ideas from their code. 11 | #Lookahead paper --> MZhang,G Hinton https://arxiv.org/abs/1907.08610 12 | # Calibrated anisotropic adaptive learning rates - https://arxiv.org/abs/1908.00700v2 13 | 14 | #summary of changes: 15 | #full code integration with all updates at param level instead of group, moves slow weights into state dict (from generic weights), 16 | #supports group learning rates (thanks @SHolderbach), fixes sporadic load from saved nn_model issues. 17 | #changes 8/31/19 - fix references to *self*.N_sma_threshold; 18 | #changed eps to 1e-5 as better default than 1e-8. 19 | 20 | import math 21 | import torch 22 | from torch.optim.optimizer import Optimizer, required 23 | import itertools as it 24 | 25 | 26 | 27 | class RangerVA(Optimizer): 28 | 29 | def __init__(self, params, lr=1e-3, alpha=0.5, k=6, n_sma_threshhold=5, betas=(.95,0.999), 30 | eps=1e-5, weight_decay=0, amsgrad=True, transformer='softplus', smooth=50, 31 | grad_transformer='square'): 32 | #parameter checks 33 | if not 0.0 <= alpha <= 1.0: 34 | raise ValueError(f'Invalid slow update rate: {alpha}') 35 | if not 1 <= k: 36 | raise ValueError(f'Invalid lookahead steps: {k}') 37 | if not lr > 0: 38 | raise ValueError(f'Invalid Learning Rate: {lr}') 39 | if not eps > 0: 40 | raise ValueError(f'Invalid eps: {eps}') 41 | 42 | #parameter comments: 43 | # beta1 (momentum) of .95 seems to work better than .90... 44 | #N_sma_threshold of 5 seems better in testing than 4. 45 | #In both cases, worth testing on your trdata (.90 vs .95, 4 vs 5) to make sure which works best for you. 46 | 47 | #prep defaults and init torch.optim base 48 | defaults = dict(lr=lr, alpha=alpha, k=k, step_counter=0, betas=betas, 49 | n_sma_threshhold=n_sma_threshhold, eps=eps, weight_decay=weight_decay, 50 | smooth=smooth, transformer=transformer, grad_transformer=grad_transformer, 51 | amsgrad=amsgrad) 52 | super().__init__(params,defaults) 53 | 54 | #adjustable threshold 55 | self.n_sma_threshhold = n_sma_threshhold 56 | 57 | #look ahead params 58 | self.alpha = alpha 59 | self.k = k 60 | 61 | #radam buffer for state 62 | self.radam_buffer = [[None,None,None] for ind in range(10)] 63 | 64 | #self.first_run_check=0 65 | 66 | #lookahead weights 67 | #9/2/19 - lookahead param tensors have been moved to state storage. 68 | #This should resolve issues with load/save where weights were left in GPU memory from first load, slowing down future runs. 69 | 70 | #self.slow_weights = [[p.clone().detach() for p in group['params']] 71 | # for group in self.param_groups] 72 | 73 | #don't use grad for lookahead weights 74 | #for w in it.chain(*self.slow_weights): 75 | # w.requires_grad = False 76 | 77 | def __setstate__(self, state): 78 | super(RangerVA, self).__setstate__(state) 79 | 80 | def step(self, closure=None): 81 | loss = None 82 | #note - below is commented out b/c I have other work that passes back the loss as a float, and thus not a callable closure. 83 | #Uncomment if you need to use the actual closure... 84 | 85 | #if closure is not None: 86 | #loss = closure() 87 | 88 | #Evaluate averages and grad, update param tensors 89 | for group in self.param_groups: 90 | 91 | for p in group['params']: 92 | if p.grad is None: 93 | continue 94 | grad = p.grad.data.float() 95 | if grad.is_sparse: 96 | raise RuntimeError('Ranger optimizer does not support sparse gradients') 97 | 98 | amsgrad = group['amsgrad'] 99 | smooth = group['smooth'] 100 | grad_transformer = group['grad_transformer'] 101 | 102 | p_data_fp32 = p.data.float() 103 | 104 | state = self.state[p] #get state dict for this param 105 | 106 | if len(state) == 0: #if first time to run...init dictionary with our desired entries 107 | #if self.first_run_check==0: 108 | #self.first_run_check=1 109 | #print("Initializing slow buffer...should not see this at load from saved nn_model!") 110 | state['step'] = 0 111 | state['exp_avg'] = torch.zeros_like(p_data_fp32) 112 | state['exp_avg_sq'] = torch.zeros_like(p_data_fp32) 113 | if amsgrad: 114 | # Maintains max of all exp. moving avg. of sq. grad. values 115 | state['max_exp_avg_sq'] = torch.zeros_like(p.data) 116 | 117 | #look ahead weight storage now in state dict 118 | state['slow_buffer'] = torch.empty_like(p.data) 119 | state['slow_buffer'].copy_(p.data) 120 | 121 | else: 122 | state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32) 123 | state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32) 124 | 125 | 126 | #begin computations 127 | exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] 128 | beta1, beta2 = group['betas'] 129 | if amsgrad: 130 | max_exp_avg_sq = state['max_exp_avg_sq'] 131 | 132 | 133 | #compute variance mov avg 134 | exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) 135 | #compute mean moving avg 136 | exp_avg.mul_(beta1).add_(1 - beta1, grad) 137 | 138 | 139 | 140 | 141 | ##transformer 142 | if grad_transformer == 'square': 143 | grad_tmp = grad**2 144 | elif grad_transformer == 'abs': 145 | grad_tmp = grad.abs() 146 | 147 | 148 | exp_avg_sq.mul_(beta2).add_((1 - beta2)*grad_tmp) 149 | 150 | 151 | 152 | if amsgrad: 153 | # Maintains the maximum of all 2nd moment running avg. till now 154 | torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq) 155 | # Use the max. for normalizing running avg. of gradient 156 | denomc = max_exp_avg_sq.clone() 157 | else: 158 | denomc = exp_avg_sq.clone() 159 | 160 | if grad_transformer == 'square': 161 | #pdb.set_trace() 162 | denomc.sqrt_() 163 | 164 | 165 | 166 | 167 | 168 | state['step'] += 1 169 | 170 | 171 | 172 | 173 | 174 | if group['weight_decay'] != 0: 175 | p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32) 176 | 177 | 178 | bias_correction1 = 1 - beta1 ** state['step'] 179 | bias_correction2 = 1 - beta2 ** state['step'] 180 | step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1 181 | 182 | 183 | # ...let's use calibrated alr 184 | if group['transformer'] =='softplus': 185 | sp = torch.nn.Softplus( smooth) 186 | denomf = sp( denomc) 187 | p_data_fp32.addcdiv_(-step_size, exp_avg, denomf ) 188 | 189 | else: 190 | 191 | denom = exp_avg_sq.sqrt().add_(group['eps']) 192 | p_data_fp32.addcdiv_(-step_size * group['lr'], exp_avg, denom) 193 | 194 | 195 | p.data.copy_(p_data_fp32) 196 | 197 | #integrated look ahead... 198 | #we do it at the param level instead of group level 199 | if state['step'] % group['k'] == 0: 200 | slow_p = state['slow_buffer'] #get access to slow param tensor 201 | slow_p.add_(self.alpha, p.data - slow_p) #(fast weights - slow weights) * alpha 202 | p.data.copy_(slow_p) #copy interpolated weights to RAdam param tensor 203 | 204 | return loss 205 | -------------------------------------------------------------------------------- /ikflow/thirdparty/ranger/rangerqh.py: -------------------------------------------------------------------------------- 1 | # RangerQH - @lessw2020 github 2 | # Combines Quasi Hyperbolic momentum with Hinton Lookahead. 3 | 4 | # https://arxiv.org/abs/1810.06801v4 (QH paper) 5 | # #Lookahead paper --> MZhang,G Hinton https://arxiv.org/abs/1907.08610 6 | 7 | 8 | 9 | 10 | # Some portions = Copyright (c) Facebook, Inc. and its affiliates. 11 | # 12 | # This source code is licensed under the MIT license found in the 13 | # LICENSE file in the root directory of this source tree. 14 | 15 | import torch 16 | from torch.optim.optimizer import Optimizer 17 | 18 | #from ..common import param_conv 19 | 20 | 21 | class RangerQH(Optimizer): 22 | r"""Implements the QHAdam optimization algorithm `(Ma and Yarats, 2019)`_. 23 | Along with Hinton/Zhang Lookahead. 24 | Args: 25 | params (iterable): 26 | iterable of parameters to optimize or dicts defining parameter 27 | groups 28 | lr (float, optional): learning rate (:math:`\alpha` from the paper) 29 | (default: 1e-3) 30 | betas (Tuple[float, float], optional): coefficients used for computing 31 | running averages of the gradient and its square 32 | (default: (0.9, 0.999)) 33 | nus (Tuple[float, float], optional): immediate discount factors used to 34 | estimate the gradient and its square 35 | (default: (1.0, 1.0)) 36 | eps (float, optional): term added to the denominator to improve 37 | numerical stability 38 | (default: 1e-8) 39 | weight_decay (float, optional): weight decay (default: 0.0) 40 | decouple_weight_decay (bool, optional): whether to decouple the weight 41 | decay from the gradient-based optimization step 42 | (default: False) 43 | Example: 44 | >>> optimizer = qhoptim.pyt.QHAdam( 45 | ... nn_model.parameters(), 46 | ... lr=3e-4, nus=(0.8, 1.0), betas=(0.99, 0.999)) 47 | >>> optimizer.zero_grad() 48 | >>> loss_fn(nn_model(input), target).backward() 49 | >>> optimizer.step() 50 | .. _`(Ma and Yarats, 2019)`: https://arxiv.org/abs/1810.06801 51 | """ 52 | 53 | def __init__( 54 | self, 55 | params, 56 | lr=1e-3, 57 | betas=(0.9, 0.999), 58 | nus=(.7, 1.0), 59 | weight_decay=0.0, 60 | k=6, 61 | alpha=.5, 62 | decouple_weight_decay=False, 63 | eps=1e-8, 64 | ): 65 | if not 0.0 <= lr: 66 | raise ValueError("Invalid learning rate: {}".format(lr)) 67 | if not 0.0 <= eps: 68 | raise ValueError("Invalid epsilon value: {}".format(eps)) 69 | if not 0.0 <= betas[0] < 1.0: 70 | raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) 71 | if not 0.0 <= betas[1] < 1.0: 72 | raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) 73 | if weight_decay < 0.0: 74 | raise ValueError("Invalid weight_decay value: {}".format(weight_decay)) 75 | 76 | defaults = { 77 | "lr": lr, 78 | "betas": betas, 79 | "nus": nus, 80 | "weight_decay": weight_decay, 81 | "decouple_weight_decay": decouple_weight_decay, 82 | "eps": eps, 83 | } 84 | super().__init__(params, defaults) 85 | 86 | #look ahead params 87 | self.alpha = alpha 88 | self.k = k 89 | 90 | 91 | def step(self, closure=None): 92 | """Performs a single optimization step. 93 | Args: 94 | closure (callable, optional): 95 | A closure that reevaluates the nn_model and returns the loss. 96 | """ 97 | loss = None 98 | if closure is not None: 99 | loss = closure() 100 | 101 | for group in self.param_groups: 102 | lr = group["lr"] 103 | beta1, beta2 = group["betas"] 104 | nu1, nu2 = group["nus"] 105 | weight_decay = group["weight_decay"] 106 | decouple_weight_decay = group["decouple_weight_decay"] 107 | eps = group["eps"] 108 | 109 | for p in group["params"]: 110 | if p.grad is None: 111 | continue 112 | 113 | d_p = p.grad.data 114 | if d_p.is_sparse: 115 | raise RuntimeError("QHAdam does not support sparse gradients") 116 | 117 | 118 | 119 | if weight_decay != 0: 120 | if decouple_weight_decay: 121 | p.data.mul_(1 - lr * weight_decay) 122 | else: 123 | d_p.add_(weight_decay, p.data) 124 | 125 | d_p_sq = d_p.mul(d_p) 126 | 127 | #prep for saved param loading 128 | param_state = self.state[p] 129 | 130 | if len(param_state) == 0: 131 | param_state["beta1_weight"] = 0.0 132 | param_state["beta2_weight"] = 0.0 133 | param_state['step'] = 0 134 | param_state["exp_avg"] = torch.zeros_like(p.data) 135 | param_state["exp_avg_sq"] = torch.zeros_like(p.data) 136 | #look ahead weight storage now in state dict 137 | param_state['slow_buffer'] = torch.empty_like(p.data) 138 | param_state['slow_buffer'].copy_(p.data) 139 | 140 | 141 | param_state['step'] += 1 142 | 143 | param_state["beta1_weight"] = 1.0 + beta1 * param_state["beta1_weight"] 144 | param_state["beta2_weight"] = 1.0 + beta2 * param_state["beta2_weight"] 145 | 146 | beta1_weight = param_state["beta1_weight"] 147 | beta2_weight = param_state["beta2_weight"] 148 | exp_avg = param_state["exp_avg"] 149 | exp_avg_sq = param_state["exp_avg_sq"] 150 | 151 | beta1_adj = 1.0 - (1.0 / beta1_weight) 152 | beta2_adj = 1.0 - (1.0 / beta2_weight) 153 | exp_avg.mul_(beta1_adj).add_(1.0 - beta1_adj, d_p) 154 | exp_avg_sq.mul_(beta2_adj).add_(1.0 - beta2_adj, d_p_sq) 155 | 156 | avg_grad = exp_avg.mul(nu1) 157 | if nu1 != 1.0: 158 | avg_grad.add_(1.0 - nu1, d_p) 159 | 160 | avg_grad_rms = exp_avg_sq.mul(nu2) 161 | if nu2 != 1.0: 162 | avg_grad_rms.add_(1.0 - nu2, d_p_sq) 163 | avg_grad_rms.sqrt_() 164 | if eps != 0.0: 165 | avg_grad_rms.add_(eps) 166 | 167 | p.data.addcdiv_(-lr, avg_grad, avg_grad_rms) 168 | 169 | #integrated look ahead... 170 | #we do it at the param level instead of group level 171 | if param_state['step'] % self.k ==0: #group['k'] == 0: 172 | slow_p = param_state['slow_buffer'] #get access to slow param tensor 173 | slow_p.add_(self.alpha, p.data - slow_p) #(fast weights - slow weights) * alpha 174 | p.data.copy_(slow_p) #copy interpolated weights to RAdam param tensor 175 | 176 | 177 | return loss 178 | 179 | @classmethod 180 | def _params_to_dict(cls, params): 181 | return {"lr": params.alpha, "nus": (params.nu1, params.nu2), "betas": (params.beta1, params.beta2)} 182 | 183 | -------------------------------------------------------------------------------- /ikflow/training/lt_data.py: -------------------------------------------------------------------------------- 1 | import os 2 | from typing import List, Dict 3 | 4 | from torch.utils.data import DataLoader 5 | from pytorch_lightning.core.datamodule import LightningDataModule 6 | import wandb 7 | import torch 8 | 9 | from ikflow.config import ALL_DATASET_TAGS, device 10 | from ikflow.utils import get_sum_joint_limit_range, get_dataset_directory, get_dataset_filepaths 11 | 12 | 13 | class IkfLitDataset(LightningDataModule): 14 | def __init__( 15 | self, robot_name: str, batch_size: int, val_set_size: int, dataset_tags: List[str], prepare_data_per_node=True 16 | ): 17 | for tag in dataset_tags: 18 | assert tag in ALL_DATASET_TAGS 19 | 20 | self._robot_name = robot_name 21 | self._batch_size = batch_size 22 | self._val_set_size = val_set_size 23 | 24 | # If set to True will call prepare_data() on LOCAL_RANK=0 for every node. If set to False will only call from NODE_RANK=0, LOCAL_RANK=0. 25 | self.prepare_data_per_node = prepare_data_per_node 26 | self._log_hyperparams = True 27 | 28 | dataset_directory = get_dataset_directory(self._robot_name) 29 | assert os.path.isdir(dataset_directory), ( 30 | f"Directory '{dataset_directory}' doesn't exist - have you created the dataset for this robot yet? Your probably want to run" 31 | f" `uv run python scripts/build_dataset.py --robot_name={robot_name} --training_set_size=10000000 --only_non_self_colliding`)" 32 | ) 33 | samples_tr_file_path, poses_tr_file_path, samples_te_file_path, poses_te_file_path, _ = get_dataset_filepaths( 34 | dataset_directory, dataset_tags 35 | ) 36 | self._samples_tr = torch.load(samples_tr_file_path).to(device) 37 | self._endpoints_tr = torch.load(poses_tr_file_path).to(device) 38 | self._samples_te = torch.load(samples_te_file_path).to(device) 39 | self._endpoints_te = torch.load(poses_te_file_path).to(device) 40 | 41 | self._sum_joint_limit_range = get_sum_joint_limit_range(self._samples_tr) 42 | self.allow_zero_length_dataloader_with_multiple_devices = False 43 | 44 | def add_dataset_hashes_to_cfg(self, cfg: Dict): 45 | cfg.update( 46 | { 47 | "dataset_hashes": str( 48 | [ 49 | self._samples_tr.sum().item(), 50 | self._endpoints_tr.sum().item(), 51 | self._samples_te.sum().item(), 52 | self._endpoints_te.sum().item(), 53 | ] 54 | ) 55 | } 56 | ) 57 | 58 | def log_dataset_sizes(self, epoch=0, batch_nb=0): 59 | """Log the training and testset size to wandb""" 60 | assert self._samples_tr.shape[0] == self._endpoints_tr.shape[0] 61 | assert self._samples_te.shape[0] == self._endpoints_te.shape[0] 62 | wandb.log( 63 | { 64 | "epoch": epoch, 65 | "batch_nb": batch_nb, 66 | "small_dataset": self._use_small_dataset, 67 | "sum_joint_limit_range": self._sum_joint_limit_range, 68 | "dataset_size_tr": self._samples_tr.shape[0], 69 | "dataset_size_te": self._samples_te.shape[0], 70 | } 71 | ) 72 | 73 | def train_dataloader(self): 74 | return DataLoader( 75 | torch.utils.data.TensorDataset(self._samples_tr, self._endpoints_tr), 76 | batch_size=self._batch_size, 77 | shuffle=True, 78 | drop_last=True, 79 | # see https://github.com/dbolya/yolact/issues/664#issuecomment-975051339 80 | generator=torch.Generator(device=device), 81 | ) 82 | 83 | def val_dataloader(self): 84 | return DataLoader( 85 | torch.utils.data.TensorDataset( 86 | self._samples_te[0 : self._val_set_size], self._endpoints_te[0 : self._val_set_size] 87 | ), 88 | batch_size=1, 89 | shuffle=False, 90 | drop_last=True, 91 | ) 92 | -------------------------------------------------------------------------------- /ikflow/training/training_utils.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | import os 3 | 4 | import numpy as np 5 | import torch 6 | import wandb 7 | 8 | from ikflow import config 9 | 10 | 11 | def get_softflow_noise(x: torch.Tensor, softflow_noise_scale: float): 12 | """ 13 | Return noise and noise magnitude for softflow. See https://arxiv.org/abs/2006.04604 14 | Return: 15 | c (torch.Tensor): An (batch_sz x 1) array storing the magnitude of the sampled gausian `eps` along each row 16 | eps (torch.Tensor): a (batch_sz x dim_x) array, where v[i] is a sample drawn from a zero mean gaussian with variance = c[i]**2 17 | """ 18 | dim_x = x.shape[1] 19 | 20 | # (batch_size x 1) from uniform(0, 1) 21 | c = torch.rand_like(x[:, 0]).unsqueeze(1) 22 | 23 | # (batch_size x dim_y) * (batch_size x dim_y) | element wise multiplication 24 | eps = torch.randn_like(x) * torch.cat(dim_x * [c], dim=1) * softflow_noise_scale 25 | return c, eps 26 | 27 | 28 | def _np_to_str(x: np.ndarray) -> str: 29 | assert x.ndim == 1 30 | return str([round(x_i, 4) for x_i in x.tolist()]) 31 | 32 | 33 | def _datetime_str() -> str: 34 | now = datetime.now() 35 | return now.strftime("%b.%d.%Y_%I:%-M:%p") 36 | 37 | 38 | def get_checkpoint_dir(robot_name: str) -> str: 39 | if wandb.run is not None: 40 | ckpt_dir_name = f"{robot_name}--{_datetime_str()}--wandb-run-name:{wandb.run.name}/" 41 | else: 42 | ckpt_dir_name = f"{robot_name}--{_datetime_str()}/" 43 | dir_filepath = os.path.join(config.TRAINING_LOGS_DIR, ckpt_dir_name) 44 | return dir_filepath 45 | 46 | 47 | # def log_ik_solutions_to_wandb_table(self, k: int = 3): 48 | # """Log `k` solutions for each of the target poses in ROBOT_TARGET_POSE_POINTS_FOR_PLOTTING to wandb""" 49 | # if wandb.run is None: 50 | # print("wandb not initialized, skipping logging solution table") 51 | # return 52 | 53 | # # TODO: 54 | # # see https://github.com/wandb/wandb/issues/1826 55 | # # 56 | # robot = self.ik_solver.robot 57 | # if robot.name not in ROBOT_TARGET_POSE_POINTS_FOR_PLOTTING: 58 | # return 59 | 60 | # target_poses = ROBOT_TARGET_POSE_POINTS_FOR_PLOTTING[robot.name] 61 | # joint_types = robot.actuated_joint_types 62 | 63 | # for i, target_pose in enumerate(target_poses): 64 | # solutions = self.generate_ik_solutions(target_pose, k) # (k, n_dofs) 65 | # realized_poses = robot.forward_kinematics(solutions.cpu().detach().numpy()) 66 | # errors_xyz = realized_poses[:, 0:3] - target_pose[0:3] 67 | # q_target = np.tile(target_pose[3:7], (k, 1)) 68 | # q_realized = realized_poses[:, 3:7] 69 | # # see https://gamedev.stackexchange.com/a/68184 70 | # q_error = quaternion_product(q_target, quaternion_inverse(q_realized)) 71 | # errors = np.hstack([errors_xyz, q_error]) 72 | # assert errors.shape == (k, 7) 73 | # for j in range(k): 74 | # ik_solution_table.add_data( 75 | # int(self.global_step), 76 | # _np_to_str(target_pose), 77 | # _np_to_str(solutions[j]), 78 | # _np_to_str(realized_poses[j]), 79 | # _np_to_str(errors[j]), 80 | # str(joint_types), 81 | # ) 82 | 83 | # new_table = wandb.Table(data=ik_solution_table.data, columns=_IK_SOLUTION_TABLE_COLUMNS) 84 | # wandb.log({f"{self.ik_solver.robot.name}_solution_table": new_table, "global_step": float(self.global_step)}) 85 | -------------------------------------------------------------------------------- /ikflow/utils.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, Optional, Callable, List 2 | import pathlib 3 | import os 4 | import random 5 | import pkg_resources 6 | 7 | import numpy as np 8 | import torch 9 | 10 | from ikflow import config 11 | 12 | 13 | def get_wandb_project() -> Tuple[str, str]: 14 | """Get the wandb entity and project. Reads from environment variables""" 15 | 16 | wandb_project = os.getenv("WANDB_PROJECT") 17 | wandb_entity = os.getenv("WANDB_ENTITY") 18 | assert wandb_project is not None, ( 19 | "The 'WANDB_PROJECT' environment variable is not set (try `export WANDB_PROJECT=`)" 20 | ) 21 | assert wandb_entity is not None, ( 22 | "The 'WANDB_ENTITY' environment variable is not set (try `export WANDB_PROJECT=`)" 23 | ) 24 | return wandb_entity, wandb_project 25 | 26 | 27 | def get_dataset_directory(robot: str) -> str: 28 | """Return the path of the directory""" 29 | return os.path.join(config.DATASET_DIR, robot) 30 | 31 | 32 | def get_dataset_filepaths(dataset_directory: str, tags: List[str]): 33 | """Return the filepaths of the tensors in a dataset""" 34 | 35 | def filename_w_tags(filename: str): 36 | for i, tag in enumerate(tags): 37 | filename = filename + f"__tag{i}={tag}" 38 | return filename 39 | 40 | samples_tr_file_path = os.path.join(dataset_directory, filename_w_tags("samples_tr.pt")) 41 | poses_tr_file_path = os.path.join(dataset_directory, filename_w_tags("endpoints_tr.pt")) 42 | samples_te_file_path = os.path.join(dataset_directory, filename_w_tags("samples_te.pt")) 43 | poses_te_file_path = os.path.join(dataset_directory, filename_w_tags("endpoints_te.pt")) 44 | info_filepath = os.path.join(dataset_directory, filename_w_tags("info.txt")) 45 | return samples_tr_file_path, poses_tr_file_path, samples_te_file_path, poses_te_file_path, info_filepath 46 | 47 | 48 | def get_filepath(local_filepath: str): 49 | return pkg_resources.resource_filename(__name__, local_filepath) 50 | 51 | 52 | # _____________ 53 | # Pytorch utils 54 | 55 | 56 | def assert_joint_angle_tensor_in_joint_limits( 57 | joints_limits: List[Tuple[float, float]], x: torch.Tensor, description: str, eps: float 58 | ): 59 | """Validate that a tensor of joint angles is within the joint limits of the robot.""" 60 | for i, (lower, upper) in enumerate(joints_limits): 61 | max_elem = torch.max(x[:, i]).item() 62 | min_elem = torch.min(x[:, i]).item() 63 | error_lower = min_elem - (lower - eps) 64 | error_upper = max_elem - (upper + eps) 65 | assert min_elem >= lower - eps, ( 66 | f"[{description}] Joint angle {min_elem} is less than lower limit {lower} (minus eps={eps}) for joint {i} -" 67 | f" error = {error_lower}\n limits(joint_{i}) = ({lower}, {upper})" 68 | ) 69 | assert max_elem <= upper + eps, ( 70 | f"[{description}] Max element {max_elem} is greater than upper limit {upper} (plus eps={eps}) for joint" 71 | f" {i} - error = {error_upper}\n limits(joint_{i}) = ({lower}, {upper})" 72 | ) 73 | 74 | 75 | def set_seed(seed=0): 76 | torch.manual_seed(seed) 77 | torch.cuda.manual_seed_all(seed) 78 | torch.backends.cudnn.deterministic = True 79 | torch.backends.cudnn.benchmark = False 80 | np.random.seed(seed) 81 | random.seed(seed) 82 | os.environ["PYTHONHASHSEED"] = str(0) 83 | print("set_seed() - random int: ", torch.randint(0, 1000, (1, 1)).item()) 84 | 85 | 86 | def cuda_info(): 87 | """Printout the current cuda status""" 88 | cuda_available = torch.cuda.is_available() 89 | print("\n____________\ncuda_info()") 90 | print(f"cuda_available: {cuda_available}") 91 | 92 | if cuda_available: 93 | print(f" current_device: {torch.cuda.current_device()}") 94 | print(f" device(0): {torch.cuda.device(0)}") 95 | print(f" device_count: {torch.cuda.device_count()}") 96 | print(f" get_device_name(0): {torch.cuda.get_device_name(0)}") 97 | print() 98 | 99 | 100 | # __________________ 101 | # Printing functions 102 | 103 | 104 | def print_tensor_stats( 105 | arr, 106 | name="", 107 | writable: Optional[ 108 | Callable[ 109 | [ 110 | str, 111 | ], 112 | None, 113 | ] 114 | ] = None, 115 | ): 116 | if writable is None: 117 | writable = lambda _s: None 118 | 119 | round_amt = 4 120 | 121 | s = f"\n\t\tmin,\tmax,\tmean,\tstd - for '{name}'" 122 | print(s) 123 | writable.write(s + "\n") 124 | 125 | for i in range(arr.shape[1]): 126 | if "torch" in str(type(arr)): 127 | min_ = round(torch.min(arr[:, i]).item(), round_amt) 128 | max_ = round(torch.max(arr[:, i]).item(), round_amt) 129 | mean = round(torch.mean(arr[:, i]).item(), round_amt) 130 | std = round(torch.std(arr[:, i]).item(), round_amt) 131 | else: 132 | min_ = round(np.min(arr[:, i]), round_amt) 133 | max_ = round(np.max(arr[:, i]), round_amt) 134 | mean = round(np.mean(arr[:, i]), round_amt) 135 | std = round(np.std(arr[:, i]), round_amt) 136 | s = f" col_{i}:\t{min_}\t{max_}\t{mean}\t{std}" 137 | print(s) 138 | writable.write(s + "\n") 139 | 140 | 141 | def get_sum_joint_limit_range(samples): 142 | """Return the total joint limit range""" 143 | sum_joint_range = 0 144 | for joint_i in range(samples.shape[1]): 145 | min_sample = torch.min(samples[:, joint_i]) 146 | max_sample = torch.max(samples[:, joint_i]) 147 | sum_joint_range += max_sample - min_sample 148 | return sum_joint_range 149 | 150 | 151 | # ___________________ 152 | # Scripting functions 153 | 154 | 155 | def boolean_string(s): 156 | if isinstance(s, bool): 157 | return s 158 | if s.upper() not in {"FALSE", "TRUE"}: 159 | raise ValueError(f'input: "{s}" ("{type(s)}") is not a valid boolean string') 160 | return s.upper() == "TRUE" 161 | 162 | 163 | def non_private_dict(d): 164 | r = {} 165 | for k, v in d.items(): 166 | if k[0] == "_": 167 | continue 168 | r[k] = v 169 | return r 170 | 171 | 172 | # _____________________ 173 | # File system utilities 174 | 175 | 176 | def safe_mkdir(dir_name: str): 177 | """Create a directory `dir_name`. May include multiple levels of new directories""" 178 | pathlib.Path(dir_name).mkdir(exist_ok=True, parents=True) 179 | 180 | 181 | # ______________ 182 | # Training utils 183 | 184 | 185 | def grad_stats(params_trainable) -> Tuple[float, float, float]: 186 | """ 187 | Return the average and max. gradient from the parameters in params_trainable 188 | """ 189 | ave_grads = [] 190 | abs_ave_grads = [] 191 | max_grad = 0.0 192 | for p in params_trainable: 193 | if p.grad is not None: 194 | ave_grads.append(p.grad.mean().item()) 195 | abs_ave_grads.append(p.grad.abs().mean().item()) 196 | max_grad = max(max_grad, p.grad.data.max().item()) 197 | return np.average(ave_grads), np.average(abs_ave_grads), max_grad 198 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/README.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | # Working with object files and klampt 5 | 6 | 7 | Klampt can load object files either as Terrains, or as RigidObject ( see `loadRigidObject(self, fn)`, `loadTerrain(self, fn)` in robotsim.py). Calling `loadElement(self, fn)` will load a mesh as a Terrain automatically. Terrains in klampt's vis are tan/brown with checkered squares. This looks fine for the ground, but doesn't look good / realistic for objects, so make sure to load objects as RigidObject. 8 | 9 | Process for creating a rigid body: 10 | 1. Go to clara.io and download the object you want to add as a .obj file. Move the file to models/klampt_resources/objects and add a copy to models/klampt_resources/objects/claraio_downloads. This is so if you mess up the file you have a backup 11 | 2. Run `ctmconv .obj .off`. Copy `srimugsmooth.obj`, and rename it `klampt_.obj`. Update the first line to `".off"` 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | # Objects 25 | 26 | 1. Mug. From downloaded from https://clara.io/view/e71cd85a-d179-4905-827a-5febbc842987# 27 | 28 | 29 | 30 | 31 | 32 | # Not imported 33 | 34 | 1. Conveyor belt https://clara.io/view/fd7007b6-3499-4e5d-9297-dfdac27195f9 35 | 2. Table https://clara.io/view/c162c073-ec50-428c-8e1e-22cd34535ee2 -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/block.obj: -------------------------------------------------------------------------------- 1 | mesh "cube.off" 2 | geomscale 0.4 0.4 0.4 3 | geomtranslate -0.2 -0.2 -0.2 4 | T 1 0 0 0 1 0 0 0 1 1.0 0.0 0.3 5 | #T 1 0 0 0 1 0 0 0 1 0.0 0.0 2.0 6 | mass 0.1 7 | autoMass 8 | kStiffness 10000 9 | kDamping 3000 10 | kRestitution 0 11 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/block_small.obj: -------------------------------------------------------------------------------- 1 | mesh "cube.off" 2 | geomscale 0.08 0.08 0.08 3 | geomtranslate -0.04 -0.04 -0.04 4 | T 1 0 0 0 1 0 0 0 1 0.5 0.17 0.1 5 | mass 0.1 6 | autoMass 7 | kFriction 1 8 | kStiffness 10000 9 | kDamping 3000 10 | kRestitution 0 11 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/centeredcube.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 8 12 0 3 | -0.500000 -0.500000 -0.500000 4 | -0.500000 -0.500000 0.500000 5 | -0.500000 0.500000 -0.500000 6 | -0.500000 0.500000 0.500000 7 | 0.500000 -0.500000 -0.500000 8 | 0.500000 -0.500000 0.500000 9 | 0.500000 0.500000 -0.500000 10 | 0.500000 0.500000 0.500000 11 | 3 0 1 3 12 | 3 0 3 2 13 | 3 4 6 7 14 | 3 4 7 5 15 | 3 0 4 5 16 | 3 0 5 1 17 | 3 2 3 7 18 | 3 2 7 6 19 | 3 0 2 6 20 | 3 0 6 4 21 | 3 1 5 7 22 | 3 1 7 3 23 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/claraio_downloads/mug.mlp: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 1 0 0 0 7 | 0 1 0 0 8 | 0 0 1 0 9 | 0 0 0 1 10 | 11 | 101001000000010000000100000001010101010010100000110100111011100000001001 12 | 13 | 14 | 15 | 1 0 0 0 16 | 0 1 0 0 17 | 0 0 1 0 18 | 0 0 0 1 19 | 20 | 001111000000010000000100000000000000000010100000110100111011110000001101 21 | 22 | 23 | 24 | 25 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/claraio_downloads/mug.mtl: -------------------------------------------------------------------------------- 1 | # Blender MTL File: 'None' 2 | # Material Count: 1 3 | 4 | newmtl b0b0b0 5 | Ns 96.078431 6 | Ka 0.000000 0.000000 0.000000 7 | Kd 0.690196 0.690196 0.690196 8 | Ks 0.009961 0.009961 0.009961 9 | Ni 1.000000 10 | d 1.000000 11 | illum 2 12 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/convert.sh: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | ctmconv mug.obj mug.off -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/cube.geom: -------------------------------------------------------------------------------- 1 | Primitive 2 | AABB 3 | 0 0 0 1 1 1 -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/cube.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 8 12 0 3 | 0.000000 0.000000 0.000000 4 | 0.000000 0.000000 1.000000 5 | 0.000000 1.000000 0.000000 6 | 0.000000 1.000000 1.000000 7 | 1.000000 0.000000 0.000000 8 | 1.000000 0.000000 1.000000 9 | 1.000000 1.000000 0.000000 10 | 1.000000 1.000000 1.000000 11 | 3 0 1 3 12 | 3 0 3 2 13 | 3 4 6 7 14 | 3 4 7 5 15 | 3 0 4 5 16 | 3 0 5 1 17 | 3 2 3 7 18 | 3 2 7 6 19 | 3 0 2 6 20 | 3 0 6 4 21 | 3 1 5 7 22 | 3 1 7 3 23 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/cylinder.obj: -------------------------------------------------------------------------------- 1 | mesh "cylinder.off" 2 | geomscale 0.025 0.025 0.2 3 | geomtranslate 0 0 -0.1 4 | T 1 0 0 0 1 0 0 0 1 0.0 0.0 0.0 5 | mass 0.1 6 | autoMass 7 | kStiffness 10000 8 | kDamping 3000 9 | kRestitution 0 10 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/cylinder.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 66 128 0 3 | 0.000000 0.000000 0.000000 4 | 1.000000 0.000000 0.000000 5 | 0.980785 0.195090 0.000000 6 | 0.923880 0.382683 0.000000 7 | 0.831470 0.555570 0.000000 8 | 0.707107 0.707107 0.000000 9 | 0.555570 0.831470 0.000000 10 | 0.382683 0.923880 0.000000 11 | 0.195090 0.980785 0.000000 12 | 0.000000 1.000000 0.000000 13 | -0.195090 0.980785 0.000000 14 | -0.382683 0.923880 0.000000 15 | -0.555570 0.831470 0.000000 16 | -0.707107 0.707107 0.000000 17 | -0.831470 0.555570 0.000000 18 | -0.923880 0.382683 0.000000 19 | -0.980785 0.195090 0.000000 20 | -1.000000 0.000000 0.000000 21 | -0.980785 -0.195090 0.000000 22 | -0.923880 -0.382683 0.000000 23 | -0.831470 -0.555570 0.000000 24 | -0.707107 -0.707107 0.000000 25 | -0.555570 -0.831470 0.000000 26 | -0.382683 -0.923880 0.000000 27 | -0.195090 -0.980785 0.000000 28 | -0.000000 -1.000000 0.000000 29 | 0.195090 -0.980785 0.000000 30 | 0.382683 -0.923880 0.000000 31 | 0.555570 -0.831470 0.000000 32 | 0.707107 -0.707107 0.000000 33 | 0.831470 -0.555570 0.000000 34 | 0.923880 -0.382683 0.000000 35 | 0.980785 -0.195090 0.000000 36 | 1.000000 0.000000 1.000000 37 | 0.980785 0.195090 1.000000 38 | 0.923880 0.382683 1.000000 39 | 0.831470 0.555570 1.000000 40 | 0.707107 0.707107 1.000000 41 | 0.555570 0.831470 1.000000 42 | 0.382683 0.923880 1.000000 43 | 0.195090 0.980785 1.000000 44 | 0.000000 1.000000 1.000000 45 | -0.195090 0.980785 1.000000 46 | -0.382683 0.923880 1.000000 47 | -0.555570 0.831470 1.000000 48 | -0.707107 0.707107 1.000000 49 | -0.831470 0.555570 1.000000 50 | -0.923880 0.382683 1.000000 51 | -0.980785 0.195090 1.000000 52 | -1.000000 0.000000 1.000000 53 | -0.980785 -0.195090 1.000000 54 | -0.923880 -0.382683 1.000000 55 | -0.831470 -0.555570 1.000000 56 | -0.707107 -0.707107 1.000000 57 | -0.555570 -0.831470 1.000000 58 | -0.382683 -0.923880 1.000000 59 | -0.195090 -0.980785 1.000000 60 | -0.000000 -1.000000 1.000000 61 | 0.195090 -0.980785 1.000000 62 | 0.382683 -0.923880 1.000000 63 | 0.555570 -0.831470 1.000000 64 | 0.707107 -0.707107 1.000000 65 | 0.831470 -0.555570 1.000000 66 | 0.923880 -0.382683 1.000000 67 | 0.980785 -0.195090 1.000000 68 | 0.000000 0.000000 1.000000 69 | 3 0 2 1 70 | 3 0 3 2 71 | 3 0 4 3 72 | 3 0 5 4 73 | 3 0 6 5 74 | 3 0 7 6 75 | 3 0 8 7 76 | 3 0 9 8 77 | 3 0 10 9 78 | 3 0 11 10 79 | 3 0 12 11 80 | 3 0 13 12 81 | 3 0 14 13 82 | 3 0 15 14 83 | 3 0 16 15 84 | 3 0 17 16 85 | 3 0 18 17 86 | 3 0 19 18 87 | 3 0 20 19 88 | 3 0 21 20 89 | 3 0 22 21 90 | 3 0 23 22 91 | 3 0 24 23 92 | 3 0 25 24 93 | 3 0 26 25 94 | 3 0 27 26 95 | 3 0 28 27 96 | 3 0 29 28 97 | 3 0 30 29 98 | 3 0 31 30 99 | 3 0 32 31 100 | 3 0 1 32 101 | 3 65 33 34 102 | 3 65 34 35 103 | 3 65 35 36 104 | 3 65 36 37 105 | 3 65 37 38 106 | 3 65 38 39 107 | 3 65 39 40 108 | 3 65 40 41 109 | 3 65 41 42 110 | 3 65 42 43 111 | 3 65 43 44 112 | 3 65 44 45 113 | 3 65 45 46 114 | 3 65 46 47 115 | 3 65 47 48 116 | 3 65 48 49 117 | 3 65 49 50 118 | 3 65 50 51 119 | 3 65 51 52 120 | 3 65 52 53 121 | 3 65 53 54 122 | 3 65 54 55 123 | 3 65 55 56 124 | 3 65 56 57 125 | 3 65 57 58 126 | 3 65 58 59 127 | 3 65 59 60 128 | 3 65 60 61 129 | 3 65 61 62 130 | 3 65 62 63 131 | 3 65 63 64 132 | 3 65 64 33 133 | 3 1 2 33 134 | 3 2 34 33 135 | 3 2 3 34 136 | 3 3 35 34 137 | 3 3 4 35 138 | 3 4 36 35 139 | 3 4 5 36 140 | 3 5 37 36 141 | 3 5 6 37 142 | 3 6 38 37 143 | 3 6 7 38 144 | 3 7 39 38 145 | 3 7 8 39 146 | 3 8 40 39 147 | 3 8 9 40 148 | 3 9 41 40 149 | 3 9 10 41 150 | 3 10 42 41 151 | 3 10 11 42 152 | 3 11 43 42 153 | 3 11 12 43 154 | 3 12 44 43 155 | 3 12 13 44 156 | 3 13 45 44 157 | 3 13 14 45 158 | 3 14 46 45 159 | 3 14 15 46 160 | 3 15 47 46 161 | 3 15 16 47 162 | 3 16 48 47 163 | 3 16 17 48 164 | 3 17 49 48 165 | 3 17 18 49 166 | 3 18 50 49 167 | 3 18 19 50 168 | 3 19 51 50 169 | 3 19 20 51 170 | 3 20 52 51 171 | 3 20 21 52 172 | 3 21 53 52 173 | 3 21 22 53 174 | 3 22 54 53 175 | 3 22 23 54 176 | 3 23 55 54 177 | 3 23 24 55 178 | 3 24 56 55 179 | 3 24 25 56 180 | 3 25 57 56 181 | 3 25 26 57 182 | 3 26 58 57 183 | 3 26 27 58 184 | 3 27 59 58 185 | 3 27 28 59 186 | 3 28 60 59 187 | 3 28 29 60 188 | 3 29 61 60 189 | 3 29 30 61 190 | 3 30 62 61 191 | 3 30 31 62 192 | 3 31 63 62 193 | 3 31 32 63 194 | 3 32 64 63 195 | 3 32 1 64 196 | 3 1 33 64 197 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/cylinder_y.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 66 128 0 3 | 0.000000 1.000000 0.000000 4 | 1.000000 1.000000 0.000000 5 | 0.980785 1.000000 0.195090 6 | 0.923880 1.000000 0.382683 7 | 0.831470 1.000000 0.555570 8 | 0.707107 1.000000 0.707107 9 | 0.555570 1.000000 0.831470 10 | 0.382683 1.000000 0.923880 11 | 0.195090 1.000000 0.980785 12 | 0.000000 1.000000 1.000000 13 | -0.195090 1.000000 0.980785 14 | -0.382683 1.000000 0.923880 15 | -0.555570 1.000000 0.831470 16 | -0.707107 1.000000 0.707107 17 | -0.831470 1.000000 0.555570 18 | -0.923880 1.000000 0.382683 19 | -0.980785 1.000000 0.195090 20 | -1.000000 1.000000 0.000000 21 | -0.980785 1.000000 -0.195090 22 | -0.923880 1.000000 -0.382683 23 | -0.831470 1.000000 -0.555570 24 | -0.707107 1.000000 -0.707107 25 | -0.555570 1.000000 -0.831470 26 | -0.382683 1.000000 -0.923880 27 | -0.195090 1.000000 -0.980785 28 | -0.000000 1.000000 -1.000000 29 | 0.195090 1.000000 -0.980785 30 | 0.382683 1.000000 -0.923880 31 | 0.555570 1.000000 -0.831470 32 | 0.707107 1.000000 -0.707107 33 | 0.831470 1.000000 -0.555570 34 | 0.923880 1.000000 -0.382683 35 | 0.980785 1.000000 -0.195090 36 | 1.000000 -1.000000 0.000000 37 | 0.980785 -1.000000 0.195090 38 | 0.923880 -1.000000 0.382683 39 | 0.831470 -1.000000 0.555570 40 | 0.707107 -1.000000 0.707107 41 | 0.555570 -1.000000 0.831470 42 | 0.382683 -1.000000 0.923880 43 | 0.195090 -1.000000 0.980785 44 | 0.000000 -1.000000 1.000000 45 | -0.195090 -1.000000 0.980785 46 | -0.382683 -1.000000 0.923880 47 | -0.555570 -1.000000 0.831470 48 | -0.707107 -1.000000 0.707107 49 | -0.831470 -1.000000 0.555570 50 | -0.923880 -1.000000 0.382683 51 | -0.980785 -1.000000 0.195090 52 | -1.000000 -1.000000 0.000000 53 | -0.980785 -1.000000 -0.195090 54 | -0.923880 -1.000000 -0.382683 55 | -0.831470 -1.000000 -0.555570 56 | -0.707107 -1.000000 -0.707107 57 | -0.555570 -1.000000 -0.831470 58 | -0.382683 -1.000000 -0.923880 59 | -0.195090 -1.000000 -0.980785 60 | -0.000000 -1.000000 -1.000000 61 | 0.195090 -1.000000 -0.980785 62 | 0.382683 -1.000000 -0.923880 63 | 0.555570 -1.000000 -0.831470 64 | 0.707107 -1.000000 -0.707107 65 | 0.831470 -1.000000 -0.555570 66 | 0.923880 -1.000000 -0.382683 67 | 0.980785 -1.000000 -0.195090 68 | 0.000000 -1.000000 0.000000 69 | 3 0 2 1 70 | 3 0 3 2 71 | 3 0 4 3 72 | 3 0 5 4 73 | 3 0 6 5 74 | 3 0 7 6 75 | 3 0 8 7 76 | 3 0 9 8 77 | 3 0 10 9 78 | 3 0 11 10 79 | 3 0 12 11 80 | 3 0 13 12 81 | 3 0 14 13 82 | 3 0 15 14 83 | 3 0 16 15 84 | 3 0 17 16 85 | 3 0 18 17 86 | 3 0 19 18 87 | 3 0 20 19 88 | 3 0 21 20 89 | 3 0 22 21 90 | 3 0 23 22 91 | 3 0 24 23 92 | 3 0 25 24 93 | 3 0 26 25 94 | 3 0 27 26 95 | 3 0 28 27 96 | 3 0 29 28 97 | 3 0 30 29 98 | 3 0 31 30 99 | 3 0 32 31 100 | 3 0 1 32 101 | 3 65 33 34 102 | 3 65 34 35 103 | 3 65 35 36 104 | 3 65 36 37 105 | 3 65 37 38 106 | 3 65 38 39 107 | 3 65 39 40 108 | 3 65 40 41 109 | 3 65 41 42 110 | 3 65 42 43 111 | 3 65 43 44 112 | 3 65 44 45 113 | 3 65 45 46 114 | 3 65 46 47 115 | 3 65 47 48 116 | 3 65 48 49 117 | 3 65 49 50 118 | 3 65 50 51 119 | 3 65 51 52 120 | 3 65 52 53 121 | 3 65 53 54 122 | 3 65 54 55 123 | 3 65 55 56 124 | 3 65 56 57 125 | 3 65 57 58 126 | 3 65 58 59 127 | 3 65 59 60 128 | 3 65 60 61 129 | 3 65 61 62 130 | 3 65 62 63 131 | 3 65 63 64 132 | 3 65 64 33 133 | 3 1 2 33 134 | 3 2 34 33 135 | 3 2 3 34 136 | 3 3 35 34 137 | 3 3 4 35 138 | 3 4 36 35 139 | 3 4 5 36 140 | 3 5 37 36 141 | 3 5 6 37 142 | 3 6 38 37 143 | 3 6 7 38 144 | 3 7 39 38 145 | 3 7 8 39 146 | 3 8 40 39 147 | 3 8 9 40 148 | 3 9 41 40 149 | 3 9 10 41 150 | 3 10 42 41 151 | 3 10 11 42 152 | 3 11 43 42 153 | 3 11 12 43 154 | 3 12 44 43 155 | 3 12 13 44 156 | 3 13 45 44 157 | 3 13 14 45 158 | 3 14 46 45 159 | 3 14 15 46 160 | 3 15 47 46 161 | 3 15 16 47 162 | 3 16 48 47 163 | 3 16 17 48 164 | 3 17 49 48 165 | 3 17 18 49 166 | 3 18 50 49 167 | 3 18 19 50 168 | 3 19 51 50 169 | 3 19 20 51 170 | 3 20 52 51 171 | 3 20 21 52 172 | 3 21 53 52 173 | 3 21 22 53 174 | 3 22 54 53 175 | 3 22 23 54 176 | 3 23 55 54 177 | 3 23 24 55 178 | 3 24 56 55 179 | 3 24 25 56 180 | 3 25 57 56 181 | 3 25 26 57 182 | 3 26 58 57 183 | 3 26 27 58 184 | 3 27 59 58 185 | 3 27 28 59 186 | 3 28 60 59 187 | 3 28 29 60 188 | 3 29 61 60 189 | 3 29 30 61 190 | 3 30 62 61 191 | 3 30 31 62 192 | 3 31 63 62 193 | 3 31 32 63 194 | 3 32 64 63 195 | 3 32 1 64 196 | 3 1 33 64 197 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/flat_foot.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 8 12 0 3 | 0.000000 -0.250000 -0.025000 4 | 0.000000 -0.250000 0.025000 5 | 0.000000 0.250000 -0.025000 6 | 0.000000 0.250000 0.025000 7 | 0.200000 -0.250000 -0.025000 8 | 0.200000 -0.250000 0.025000 9 | 0.200000 0.250000 -0.025000 10 | 0.200000 0.250000 0.025000 11 | 3 0 1 3 12 | 3 0 3 2 13 | 3 4 6 7 14 | 3 4 7 5 15 | 3 0 4 5 16 | 3 0 5 1 17 | 3 2 3 7 18 | 3 2 7 6 19 | 3 0 2 6 20 | 3 0 6 4 21 | 3 1 5 7 22 | 3 1 7 3 23 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/flatblock.obj: -------------------------------------------------------------------------------- 1 | mesh "cube.off" 2 | geomscale 0.4 0.4 0.04 3 | geomtranslate -0.2 -0.2 -0.02 4 | T 1 0 0 0 1 0 0 0 1 1.0 0.0 0.06 5 | mass 0.05 6 | autoMass 7 | kStiffness 10000 8 | kDamping 3000 9 | kRestitution 0 10 | kFriction 0.5 11 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/flatblock2.obj: -------------------------------------------------------------------------------- 1 | mesh "cube.off" 2 | geomscale 0.4 0.4 0.06 3 | geomtranslate -0.2 -0.2 -0.03 4 | T 1 0 0 0 1 0 0 0 1 1.0 0.0 0.06 5 | mass 0.05 6 | autoMass 7 | kStiffness 10000 8 | kDamping 3000 9 | kRestitution 0 10 | kFriction 0.5 11 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/flatblock3.obj: -------------------------------------------------------------------------------- 1 | mesh "cube.off" 2 | geomscale 0.44 0.44 0.06 3 | geomtranslate -0.21 -0.21 -0.02 4 | T 1 0 0 0 1 0 0 0 1 1.0 0.0 0.06 5 | mass 0.05 6 | autoMass 7 | kStiffness 10000 8 | kDamping 3000 9 | kRestitution 0 10 | kFriction 0.5 11 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/flatblock4.obj: -------------------------------------------------------------------------------- 1 | mesh "cube.off" 2 | geomscale 0.4 0.4 0.03 3 | geomtranslate -0.2 -0.2 -0.015 4 | T 1 0 0 0 1 0 0 0 1 1.0 0.0 0.06 5 | mass 0.05 6 | autoMass 7 | kStiffness 10000 8 | kDamping 3000 9 | kRestitution 0 10 | kFriction 2 11 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/flatblock5.obj: -------------------------------------------------------------------------------- 1 | mesh "cube.off" 2 | geomscale 1.0 0.4 0.10 3 | geomtranslate -0.5 -0.2 -0.05 4 | T 1 0 0 0 1 0 0 0 1 0.0 0.0 0.0 5 | mass 0.05 6 | autoMass 7 | kStiffness 10000 8 | kDamping 3000 9 | kRestitution 0 10 | kFriction 0.5 11 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/mug.obj.mtl: -------------------------------------------------------------------------------- 1 | # 2 | # Wavefront material file 3 | # Converted by Meshlab Group 4 | # 5 | 6 | newmtl material_0 7 | Ka 0.200000 0.200000 0.200000 8 | Kd 0.752941 0.752941 0.752941 9 | Ks 1.000000 1.000000 1.000000 10 | Tr 1.000000 11 | illum 2 12 | Ns 0.000000 13 | 14 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/mug_klampt.obj: -------------------------------------------------------------------------------- 1 | mesh "mug.off" 2 | T 1 0 0 0 1 0 0 0 1 0.5 0.2 0.2 3 | mass 0.1 4 | autoMass 5 | kStiffness 10000 6 | kDamping 3000 7 | kRestitution 0 8 | kFriction 0.5 9 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/paddle_cube.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 8 12 0 3 | 0.000000 -0.200000 -0.025000 4 | 0.000000 -0.200000 0.025000 5 | 0.000000 0.200000 -0.025000 6 | 0.000000 0.200000 0.025000 7 | 1.000000 -0.200000 -0.025000 8 | 1.000000 -0.200000 0.025000 9 | 1.000000 0.200000 -0.025000 10 | 1.000000 0.200000 0.025000 11 | 3 0 1 3 12 | 3 0 3 2 13 | 3 4 6 7 14 | 3 4 7 5 15 | 3 0 4 5 16 | 3 0 5 1 17 | 3 2 3 7 18 | 3 2 7 6 19 | 3 0 2 6 20 | 3 0 6 4 21 | 3 1 5 7 22 | 3 1 7 3 23 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/sphere.geom: -------------------------------------------------------------------------------- 1 | Primitive 2 | Sphere 3 | 0 0 0 1 4 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/sphere.obj: -------------------------------------------------------------------------------- 1 | mesh "sphere.off" 2 | geomscale 0.1 0.1 0.1 3 | mass 1 4 | autoMass 5 | kStiffness 10000 6 | kDamping 3000 7 | kRestitution 0 8 | 9 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/sphere_10cm.obj: -------------------------------------------------------------------------------- 1 | mesh "sphere.off" 2 | geomscale 0.05 0.05 0.05 3 | mass 1 4 | autoMass 5 | kStiffness 10000 6 | kDamping 3000 7 | kRestitution 0 8 | 9 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/sphere_5cm.obj: -------------------------------------------------------------------------------- 1 | mesh "sphere.off" 2 | geomscale 0.025 0.025 0.025 3 | mass 1 4 | autoMass 5 | kStiffness 10000 6 | kDamping 3000 7 | kRestitution 0 8 | 9 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/srimugsmooth.obj: -------------------------------------------------------------------------------- 1 | mesh "srimugsmooth.off" 2 | T 1 0 0 0 1 0 0 0 1 0.5 0.2 0.2 3 | mass 0.1 4 | autoMass 5 | kStiffness 10000 6 | kDamping 3000 7 | kRestitution 0 8 | kFriction 0.5 9 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/thincube.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 8 12 0 3 | 0.000000 -0.050000 -0.050000 4 | 0.000000 -0.050000 0.050000 5 | 0.000000 0.050000 -0.050000 6 | 0.000000 0.050000 0.050000 7 | 1.000000 -0.050000 -0.050000 8 | 1.000000 -0.050000 0.050000 9 | 1.000000 0.050000 -0.050000 10 | 1.000000 0.050000 0.050000 11 | 3 0 1 3 12 | 3 0 3 2 13 | 3 4 6 7 14 | 3 4 7 5 15 | 3 0 4 5 16 | 3 0 5 1 17 | 3 2 3 7 18 | 3 2 7 6 19 | 3 0 2 6 20 | 3 0 6 4 21 | 3 1 5 7 22 | 3 1 7 3 23 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/urdf_primitives/box_ori_center.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 8 12 0 3 | 0.500000 0.500000 0.500000 4 | -0.500000 0.500000 0.500000 5 | -0.500000 -0.500000 0.500000 6 | 0.500000 -0.500000 0.500000 7 | 0.500000 0.500000 -0.500000 8 | -0.500000 0.500000 -0.500000 9 | -0.500000 -0.500000 -0.500000 10 | 0.500000 -0.500000 -0.500000 11 | 3 0 1 2 12 | 3 0 2 3 13 | 3 0 4 5 14 | 3 0 5 1 15 | 3 1 5 6 16 | 3 1 6 2 17 | 3 2 6 7 18 | 3 2 7 3 19 | 3 3 7 4 20 | 3 3 4 0 21 | 3 4 7 6 22 | 3 4 6 5 23 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/urdf_primitives/box_ori_center.tri: -------------------------------------------------------------------------------- 1 | 8 2 | 0.5 0.5 0.5 3 | -0.5 0.5 0.5 4 | -0.5 -0.5 0.5 5 | 0.5 -0.5 0.5 6 | 0.5 0.5 -0.5 7 | -0.5 0.5 -0.5 8 | -0.5 -0.5 -0.5 9 | 0.5 -0.5 -0.5 10 | 12 11 | 0 1 2 12 | 0 2 3 13 | 0 4 5 14 | 0 5 1 15 | 1 5 6 16 | 1 6 2 17 | 2 6 7 18 | 2 7 3 19 | 3 7 4 20 | 3 4 0 21 | 4 7 6 22 | 4 6 5 23 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/urdf_primitives/cylinder_ori_center.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 40 76 0 3 | 1.000000 0.000000 -0.500000 4 | 0.951057 0.309017 -0.500000 5 | 0.809017 0.587785 -0.500000 6 | 0.587785 0.809017 -0.500000 7 | 0.309017 0.951057 -0.500000 8 | 0.000000 1.000000 -0.500000 9 | -0.309017 0.951057 -0.500000 10 | -0.587785 0.809017 -0.500000 11 | -0.809017 0.587785 -0.500000 12 | -0.951057 0.309017 -0.500000 13 | -1.000000 0.000000 -0.500000 14 | -0.951057 -0.309017 -0.500000 15 | -0.809017 -0.587785 -0.500000 16 | -0.587785 -0.809017 -0.500000 17 | -0.309017 -0.951056 -0.500000 18 | -0.000000 -1.000000 -0.500000 19 | 0.309017 -0.951057 -0.500000 20 | 0.587785 -0.809017 -0.500000 21 | 0.809017 -0.587785 -0.500000 22 | 0.951056 -0.309017 -0.500000 23 | 1.000000 0.000000 0.500000 24 | 0.951057 0.309017 0.500000 25 | 0.809017 0.587785 0.500000 26 | 0.587785 0.809017 0.500000 27 | 0.309017 0.951057 0.500000 28 | 0.000000 1.000000 0.500000 29 | -0.309017 0.951057 0.500000 30 | -0.587785 0.809017 0.500000 31 | -0.809017 0.587785 0.500000 32 | -0.951057 0.309017 0.500000 33 | -1.000000 0.000000 0.500000 34 | -0.951057 -0.309017 0.500000 35 | -0.809017 -0.587785 0.500000 36 | -0.587785 -0.809017 0.500000 37 | -0.309017 -0.951056 0.500000 38 | -0.000000 -1.000000 0.500000 39 | 0.309017 -0.951057 0.500000 40 | 0.587785 -0.809017 0.500000 41 | 0.809017 -0.587785 0.500000 42 | 0.951056 -0.309017 0.500000 43 | 3 0 2 1 44 | 3 0 3 2 45 | 3 0 4 3 46 | 3 0 5 4 47 | 3 0 6 5 48 | 3 0 7 6 49 | 3 0 8 7 50 | 3 0 9 8 51 | 3 0 10 9 52 | 3 0 11 10 53 | 3 0 12 11 54 | 3 0 13 12 55 | 3 0 14 13 56 | 3 0 15 14 57 | 3 0 16 15 58 | 3 0 17 16 59 | 3 0 18 17 60 | 3 0 19 18 61 | 3 22 20 21 62 | 3 23 20 22 63 | 3 24 20 23 64 | 3 25 20 24 65 | 3 26 20 25 66 | 3 27 20 26 67 | 3 28 20 27 68 | 3 29 20 28 69 | 3 30 20 29 70 | 3 31 20 30 71 | 3 32 20 31 72 | 3 33 20 32 73 | 3 34 20 33 74 | 3 35 20 34 75 | 3 36 20 35 76 | 3 37 20 36 77 | 3 38 20 37 78 | 3 39 20 38 79 | 3 20 0 1 80 | 3 20 1 21 81 | 3 21 1 2 82 | 3 21 2 22 83 | 3 22 2 3 84 | 3 22 3 23 85 | 3 23 3 4 86 | 3 23 4 24 87 | 3 24 4 5 88 | 3 24 5 25 89 | 3 25 5 6 90 | 3 25 6 26 91 | 3 26 6 7 92 | 3 26 7 27 93 | 3 27 7 8 94 | 3 27 8 28 95 | 3 28 8 9 96 | 3 28 9 29 97 | 3 29 9 10 98 | 3 29 10 30 99 | 3 30 10 11 100 | 3 30 11 31 101 | 3 31 11 12 102 | 3 31 12 32 103 | 3 32 12 13 104 | 3 32 13 33 105 | 3 33 13 14 106 | 3 33 14 34 107 | 3 34 14 15 108 | 3 34 15 35 109 | 3 35 15 16 110 | 3 35 16 36 111 | 3 36 16 17 112 | 3 36 17 37 113 | 3 37 17 18 114 | 3 37 18 38 115 | 3 38 18 19 116 | 3 38 19 39 117 | 3 39 19 0 118 | 3 39 0 20 119 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/objects/urdf_primitives/cylinder_ori_center.tri: -------------------------------------------------------------------------------- 1 | 40 2 | 1 0 -0.5 3 | 0.951057 0.309017 -0.5 4 | 0.809017 0.587785 -0.5 5 | 0.587785 0.809017 -0.5 6 | 0.309017 0.951057 -0.5 7 | 2.67949e-08 1 -0.5 8 | -0.309017 0.951057 -0.5 9 | -0.587785 0.809017 -0.5 10 | -0.809017 0.587785 -0.5 11 | -0.951057 0.309017 -0.5 12 | -1 5.35898e-08 -0.5 13 | -0.951057 -0.309017 -0.5 14 | -0.809017 -0.587785 -0.5 15 | -0.587785 -0.809017 -0.5 16 | -0.309017 -0.951056 -0.5 17 | -8.03847e-08 -1 -0.5 18 | 0.309017 -0.951057 -0.5 19 | 0.587785 -0.809017 -0.5 20 | 0.809017 -0.587785 -0.5 21 | 0.951056 -0.309017 -0.5 22 | 1 0 0.5 23 | 0.951057 0.309017 0.5 24 | 0.809017 0.587785 0.5 25 | 0.587785 0.809017 0.5 26 | 0.309017 0.951057 0.5 27 | 2.67949e-08 1 0.5 28 | -0.309017 0.951057 0.5 29 | -0.587785 0.809017 0.5 30 | -0.809017 0.587785 0.5 31 | -0.951057 0.309017 0.5 32 | -1 5.35898e-08 0.5 33 | -0.951057 -0.309017 0.5 34 | -0.809017 -0.587785 0.5 35 | -0.587785 -0.809017 0.5 36 | -0.309017 -0.951056 0.5 37 | -8.03847e-08 -1 0.5 38 | 0.309017 -0.951057 0.5 39 | 0.587785 -0.809017 0.5 40 | 0.809017 -0.587785 0.5 41 | 0.951056 -0.309017 0.5 42 | 76 43 | 0 2 1 44 | 0 3 2 45 | 0 4 3 46 | 0 5 4 47 | 0 6 5 48 | 0 7 6 49 | 0 8 7 50 | 0 9 8 51 | 0 10 9 52 | 0 11 10 53 | 0 12 11 54 | 0 13 12 55 | 0 14 13 56 | 0 15 14 57 | 0 16 15 58 | 0 17 16 59 | 0 18 17 60 | 0 19 18 61 | 22 20 21 62 | 23 20 22 63 | 24 20 23 64 | 25 20 24 65 | 26 20 25 66 | 27 20 26 67 | 28 20 27 68 | 29 20 28 69 | 30 20 29 70 | 31 20 30 71 | 32 20 31 72 | 33 20 32 73 | 34 20 33 74 | 35 20 34 75 | 36 20 35 76 | 37 20 36 77 | 38 20 37 78 | 39 20 38 79 | 20 0 1 80 | 20 1 21 81 | 21 1 2 82 | 21 2 22 83 | 22 2 3 84 | 22 3 23 85 | 23 3 4 86 | 23 4 24 87 | 24 4 5 88 | 24 5 25 89 | 25 5 6 90 | 25 6 26 91 | 26 6 7 92 | 26 7 27 93 | 27 7 8 94 | 27 8 28 95 | 28 8 9 96 | 28 9 29 97 | 29 9 10 98 | 29 10 30 99 | 30 10 11 100 | 30 11 31 101 | 31 11 12 102 | 31 12 32 103 | 32 12 13 104 | 32 13 33 105 | 33 13 14 106 | 33 14 34 107 | 34 14 15 108 | 34 15 35 109 | 35 15 16 110 | 35 16 36 111 | 36 16 17 112 | 36 17 37 113 | 37 17 18 114 | 37 18 38 115 | 38 18 19 116 | 38 19 39 117 | 39 19 0 118 | 39 0 20 119 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/simulation_test_worlds/bounce_test.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/simulation_test_worlds/cuppile.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/simulation_test_worlds/displaytest.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 51 | 52 | 53 | 54 | 55 | 56 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/simulation_test_worlds/fastdroptest.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/simulation_test_worlds/heavyblocktest.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/simulation_test_worlds/jointtest.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/simulation_test_worlds/plane_10.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 42 40 0 3 | -5.000000 -5.000000 0.000000 4 | -5.000000 5.000000 0.000000 5 | -4.500000 -5.000000 0.000000 6 | -4.500000 5.000000 0.000000 7 | -4.000000 -5.000000 0.000000 8 | -4.000000 5.000000 0.000000 9 | -3.500000 -5.000000 0.000000 10 | -3.500000 5.000000 0.000000 11 | -3.000000 -5.000000 0.000000 12 | -3.000000 5.000000 0.000000 13 | -2.500000 -5.000000 0.000000 14 | -2.500000 5.000000 0.000000 15 | -2.000000 -5.000000 0.000000 16 | -2.000000 5.000000 0.000000 17 | -1.500000 -5.000000 0.000000 18 | -1.500000 5.000000 0.000000 19 | -1.000000 -5.000000 0.000000 20 | -1.000000 5.000000 0.000000 21 | -0.500000 -5.000000 0.000000 22 | -0.500000 5.000000 0.000000 23 | 0.000000 -5.000000 0.000000 24 | 0.000000 5.000000 0.000000 25 | 0.500000 -5.000000 0.000000 26 | 0.500000 5.000000 0.000000 27 | 1.000000 -5.000000 0.000000 28 | 1.000000 5.000000 0.000000 29 | 1.500000 -5.000000 0.000000 30 | 1.500000 5.000000 0.000000 31 | 2.000000 -5.000000 0.000000 32 | 2.000000 5.000000 0.000000 33 | 2.500000 -5.000000 0.000000 34 | 2.500000 5.000000 0.000000 35 | 3.000000 -5.000000 0.000000 36 | 3.000000 5.000000 0.000000 37 | 3.500000 -5.000000 0.000000 38 | 3.500000 5.000000 0.000000 39 | 4.000000 -5.000000 0.000000 40 | 4.000000 5.000000 0.000000 41 | 4.500000 -5.000000 0.000000 42 | 4.500000 5.000000 0.000000 43 | 5.000000 -5.000000 0.000000 44 | 5.000000 5.000000 0.000000 45 | 3 0 2 1 46 | 3 1 2 3 47 | 3 2 4 3 48 | 3 3 4 5 49 | 3 4 6 5 50 | 3 5 6 7 51 | 3 6 8 7 52 | 3 7 8 9 53 | 3 8 10 9 54 | 3 9 10 11 55 | 3 10 12 11 56 | 3 11 12 13 57 | 3 12 14 13 58 | 3 13 14 15 59 | 3 14 16 15 60 | 3 15 16 17 61 | 3 16 18 17 62 | 3 17 18 19 63 | 3 18 20 19 64 | 3 19 20 21 65 | 3 20 22 21 66 | 3 21 22 23 67 | 3 22 24 23 68 | 3 23 24 25 69 | 3 24 26 25 70 | 3 25 26 27 71 | 3 26 28 27 72 | 3 27 28 29 73 | 3 28 30 29 74 | 3 29 30 31 75 | 3 30 32 31 76 | 3 31 32 33 77 | 3 32 34 33 78 | 3 33 34 35 79 | 3 34 36 35 80 | 3 35 36 37 81 | 3 36 38 37 82 | 3 37 38 39 83 | 3 38 40 39 84 | 3 39 40 41 85 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/simulation_test_worlds/sensortest.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/simulation_test_worlds/slideblock.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/simulation_test_worlds/sphere_slope.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/simulation_test_worlds/stacktest.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83 | 84 | 85 | 86 | 87 | 88 | 89 | 90 | 91 | 92 | 93 | 94 | 95 | 96 | 97 | 98 | 99 | 100 | 101 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/simulation_test_worlds/stiffness_test.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/simulation_test_worlds/tx90cuppile.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/terrains/block.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 8 12 0 3 | -2.000000 -2.000000 -0.100000 4 | -2.000000 -2.000000 0.000000 5 | -2.000000 2.000000 -0.100000 6 | -2.000000 2.000000 0.000000 7 | 2.000000 -2.000000 -0.100000 8 | 2.000000 -2.000000 0.000000 9 | 2.000000 2.000000 -0.100000 10 | 2.000000 2.000000 0.000000 11 | 3 0 1 3 12 | 3 0 3 2 13 | 3 4 6 7 14 | 3 4 7 5 15 | 3 0 4 5 16 | 3 0 5 1 17 | 3 2 3 7 18 | 3 2 7 6 19 | 3 0 2 6 20 | 3 0 6 4 21 | 3 1 5 7 22 | 3 1 7 3 23 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/terrains/cube.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 8 12 0 3 | 0.000000 0.000000 0.000000 4 | 0.000000 0.000000 1.000000 5 | 0.000000 1.000000 0.000000 6 | 0.000000 1.000000 1.000000 7 | 1.000000 0.000000 0.000000 8 | 1.000000 0.000000 1.000000 9 | 1.000000 1.000000 0.000000 10 | 1.000000 1.000000 1.000000 11 | 3 0 1 3 12 | 3 0 3 2 13 | 3 4 6 7 14 | 3 4 7 5 15 | 3 0 4 5 16 | 3 0 5 1 17 | 3 2 3 7 18 | 3 2 7 6 19 | 3 0 2 6 20 | 3 0 6 4 21 | 3 1 5 7 22 | 3 1 7 3 23 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/terrains/drc_rough_terrain.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jstmn/ikflow/19634d436eb805f22578b4c5b6526f96d91d0323/ikflow/visualization_resources/klampt_resources/terrains/drc_rough_terrain.stl -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/terrains/plane.env: -------------------------------------------------------------------------------- 1 | mesh "plane.off" 2 | kFriction 0.5 3 | 4 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/terrains/plane.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 4 2 0 3 | -5.000000 -5.000000 0.000000 4 | -5.000000 5.000000 0.000000 5 | 5.000000 -5.000000 0.000000 6 | 5.000000 5.000000 0.000000 7 | 3 0 2 1 8 | 3 1 2 3 9 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/terrains/plane_for_atlas.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 4 2 0 3 | -5.000000 -5.000000 -0.930000 4 | -5.000000 5.000000 -0.93000 5 | 5.000000 -5.000000 -0.93000 6 | 5.000000 5.000000 -0.9300000 7 | 3 0 2 1 8 | 3 1 2 3 9 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/terrains/ridges.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 42 40 0 3 | -10.000000 -10.000000 0.000000 4 | -10.000000 10.000000 0.000000 5 | -9.000000 -10.000000 1.000000 6 | -9.000000 10.000000 1.000000 7 | -8.000000 -10.000000 0.000000 8 | -8.000000 10.000000 0.000000 9 | -7.000000 -10.000000 1.000000 10 | -7.000000 10.000000 1.000000 11 | -6.000000 -10.000000 0.000000 12 | -6.000000 10.000000 0.000000 13 | -5.000000 -10.000000 1.000000 14 | -5.000000 10.000000 1.000000 15 | -4.000000 -10.000000 0.000000 16 | -4.000000 10.000000 0.000000 17 | -3.000000 -10.000000 1.000000 18 | -3.000000 10.000000 1.000000 19 | -2.000000 -10.000000 0.000000 20 | -2.000000 10.000000 0.000000 21 | -1.000000 -10.000000 1.000000 22 | -1.000000 10.000000 1.000000 23 | 0.000000 -10.000000 0.000000 24 | 0.000000 10.000000 0.000000 25 | 1.000000 -10.000000 1.000000 26 | 1.000000 10.000000 1.000000 27 | 2.000000 -10.000000 0.000000 28 | 2.000000 10.000000 0.000000 29 | 3.000000 -10.000000 1.000000 30 | 3.000000 10.000000 1.000000 31 | 4.000000 -10.000000 0.000000 32 | 4.000000 10.000000 0.000000 33 | 5.000000 -10.000000 1.000000 34 | 5.000000 10.000000 1.000000 35 | 6.000000 -10.000000 0.000000 36 | 6.000000 10.000000 0.000000 37 | 7.000000 -10.000000 1.000000 38 | 7.000000 10.000000 1.000000 39 | 8.000000 -10.000000 0.000000 40 | 8.000000 10.000000 0.000000 41 | 9.000000 -10.000000 1.000000 42 | 9.000000 10.000000 1.000000 43 | 10.000000 -10.000000 0.000000 44 | 10.000000 10.000000 0.000000 45 | 3 0 2 1 46 | 3 1 2 3 47 | 3 2 4 3 48 | 3 3 4 5 49 | 3 4 6 5 50 | 3 5 6 7 51 | 3 6 8 7 52 | 3 7 8 9 53 | 3 8 10 9 54 | 3 9 10 11 55 | 3 10 12 11 56 | 3 11 12 13 57 | 3 12 14 13 58 | 3 13 14 15 59 | 3 14 16 15 60 | 3 15 16 17 61 | 3 16 18 17 62 | 3 17 18 19 63 | 3 18 20 19 64 | 3 19 20 21 65 | 3 20 22 21 66 | 3 21 22 23 67 | 3 22 24 23 68 | 3 23 24 25 69 | 3 24 26 25 70 | 3 25 26 27 71 | 3 26 28 27 72 | 3 27 28 29 73 | 3 28 30 29 74 | 3 29 30 31 75 | 3 30 32 31 76 | 3 31 32 33 77 | 3 32 34 33 78 | 3 33 34 35 79 | 3 34 36 35 80 | 3 35 36 37 81 | 3 36 38 37 82 | 3 37 38 39 83 | 3 38 40 39 84 | 3 39 40 41 85 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/terrains/smallplane.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 4 2 0 3 | -0.250000 -0.250000 0.000000 4 | -0.250000 0.250000 0.000000 5 | 0.250000 -0.250000 0.000000 6 | 0.250000 0.250000 0.000000 7 | 3 0 2 1 8 | 3 1 2 3 9 | -------------------------------------------------------------------------------- /ikflow/visualization_resources/klampt_resources/terrains/stair.off: -------------------------------------------------------------------------------- 1 | OFF 2 | 8 6 0 3 | -5.000000 -5.000000 0.000000 4 | -5.000000 5.000000 0.000000 5 | 0.500000 -5.000000 0.000000 6 | 0.500000 5.000000 0.000000 7 | 0.500000 -5.000000 0.100000 8 | 0.500000 5.000000 0.100000 9 | 5.000000 -5.000000 0.100000 10 | 5.000000 5.000000 0.100000 11 | 3 0 2 1 12 | 3 1 2 3 13 | 3 2 4 3 14 | 3 3 4 5 15 | 3 4 6 5 16 | 3 5 6 7 17 | -------------------------------------------------------------------------------- /lint: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | uv run ruff check --exclude ikflow/thirdparty/ --fix 4 | 5 | # Clear pynbs 6 | uv run jupyter nbconvert --clear-output --inplace notebooks/*.ipynb 7 | -------------------------------------------------------------------------------- /notebooks/ik_convergence_analysis.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "id": "d619eab0", 6 | "metadata": {}, 7 | "source": [ 8 | "# Overview\n", 9 | "\n", 10 | "This notebook analyzes the use of levenberg-marquardt for IK refinement" 11 | ] 12 | }, 13 | { 14 | "cell_type": "code", 15 | "execution_count": null, 16 | "id": "f124121e", 17 | "metadata": {}, 18 | "outputs": [], 19 | "source": [ 20 | "%load_ext autoreload\n", 21 | "%autoreload 2\n", 22 | "\n", 23 | "from IPython.display import display, HTML\n", 24 | "\n", 25 | "display(HTML(\"\"))" 26 | ] 27 | }, 28 | { 29 | "cell_type": "code", 30 | "execution_count": null, 31 | "id": "c10b367e", 32 | "metadata": {}, 33 | "outputs": [], 34 | "source": [ 35 | "from time import time\n", 36 | "\n", 37 | "import torch\n", 38 | "import numpy as np\n", 39 | "from jrl.utils import set_seed, evenly_spaced_colors\n", 40 | "import matplotlib.pyplot as plt\n", 41 | "\n", 42 | "from ikflow.model_loading import get_ik_solver\n", 43 | "\n", 44 | "torch.set_printoptions(linewidth=300, precision=6, sci_mode=False)" 45 | ] 46 | }, 47 | { 48 | "cell_type": "code", 49 | "execution_count": null, 50 | "id": "f7298640", 51 | "metadata": {}, 52 | "outputs": [], 53 | "source": [ 54 | "MODEL_NAME = \"panda__full__lp191_5.25m\"\n", 55 | "POS_ERROR_THRESHOLD = 0.001\n", 56 | "ROT_ERROR_THRESHOLD = 0.01\n", 57 | "\n", 58 | "ikflow_solver, _ = get_ik_solver(MODEL_NAME)\n", 59 | "robot = ikflow_solver.robot" 60 | ] 61 | }, 62 | { 63 | "cell_type": "code", 64 | "execution_count": null, 65 | "id": "626aeb35", 66 | "metadata": {}, 67 | "outputs": [], 68 | "source": [ 69 | "def debug_dist_to_jlims(robot, qs):\n", 70 | " \"\"\"test with:\n", 71 | " qs_random, _ = ikflow_solver.robot.sample_joint_angles_and_poses(5)\n", 72 | " debug_dist_to_jlims(robot, torch.tensor(qs_random))\n", 73 | " \"\"\"\n", 74 | "\n", 75 | " mins = 100 * torch.ones(len(qs))\n", 76 | " eps = 1e-5\n", 77 | " for i, (l, u) in enumerate(robot.actuated_joints_limits):\n", 78 | " assert torch.min(qs[:, i]) > l - eps\n", 79 | " assert torch.max(qs[:, i]) < u + eps, f\"{torch.max(qs[:, i])} !< {u}\"\n", 80 | " mins = torch.minimum(mins, torch.abs(qs[:, i] - l))\n", 81 | " mins = torch.minimum(mins, torch.abs(qs[:, i] - u))\n", 82 | " print(\"distances to joint limits:\", mins)\n", 83 | " return mins" 84 | ] 85 | }, 86 | { 87 | "cell_type": "code", 88 | "execution_count": null, 89 | "id": "d61c36b5", 90 | "metadata": {}, 91 | "outputs": [], 92 | "source": [ 93 | "def plot_runtime_curve(curves, labels):\n", 94 | " fig, (axl, axr) = plt.subplots(1, 2, figsize=(18, 8))\n", 95 | " fig.suptitle(\"Levenberg-Marquardt IK Convergence\")\n", 96 | "\n", 97 | " axl.set_title(\"Runtime\")\n", 98 | " axl.grid(alpha=0.2)\n", 99 | " axl.set_xlabel(\"batch size\")\n", 100 | " axl.set_ylabel(\"runtime (s)\")\n", 101 | "\n", 102 | " axr.set_title(\"Success Pct\")\n", 103 | " axr.grid(alpha=0.2)\n", 104 | " axr.set_xlabel(\"batch size\")\n", 105 | " axr.set_ylabel(\"success pct (%)\")\n", 106 | "\n", 107 | " colors = evenly_spaced_colors(int(1.5 * len(curves)))\n", 108 | "\n", 109 | " for (batch_sizes, runtimes, runtime_stds, success_pcts, success_pct_stds), label, color in zip(\n", 110 | " curves, labels, colors\n", 111 | " ):\n", 112 | " axl.plot(batch_sizes, runtimes, label=label, color=color)\n", 113 | " axl.fill_between(batch_sizes, runtimes - runtime_stds, runtimes + runtime_stds, alpha=0.15, color=color)\n", 114 | " axl.scatter(batch_sizes, runtimes, s=15, color=color)\n", 115 | "\n", 116 | " axr.plot(batch_sizes, success_pcts, label=label, color=color)\n", 117 | " axr.fill_between(\n", 118 | " batch_sizes, success_pcts - success_pct_stds, success_pcts + success_pct_stds, alpha=0.15, color=color\n", 119 | " )\n", 120 | " axr.scatter(batch_sizes, success_pcts, s=15, color=color)\n", 121 | "\n", 122 | " # plt.savefig(f\"lma_errors_{i}.pdf\", bbox_inches=\"tight\")\n", 123 | " # axl.legend()\n", 124 | " axr.legend()\n", 125 | " plt.show()" 126 | ] 127 | }, 128 | { 129 | "cell_type": "code", 130 | "execution_count": null, 131 | "id": "b6569b16", 132 | "metadata": {}, 133 | "outputs": [], 134 | "source": [ 135 | "batch_sizes = [10, 100, 500, 1000, 2000]\n", 136 | "k_retry = 3\n", 137 | "\n", 138 | "all_repeat_counts = [(1, 2, 5), (1, 3, 10), (1, 4, 10), (1, 5, 10)]\n", 139 | "\n", 140 | "device = \"cuda:0\"\n", 141 | "\n", 142 | "curves = []\n", 143 | "labels = []\n", 144 | "\n", 145 | "set_seed()\n", 146 | "\n", 147 | "for repeat_counts in all_repeat_counts:\n", 148 | " runtimes = []\n", 149 | " success_pcts = []\n", 150 | " runtime_stds = []\n", 151 | " success_pct_stds = []\n", 152 | "\n", 153 | " for batch_size in batch_sizes:\n", 154 | " sub_runtimes = []\n", 155 | " sub_success_pcts = []\n", 156 | "\n", 157 | " for _ in range(k_retry):\n", 158 | " _, target_poses = ikflow_solver.robot.sample_joint_angles_and_poses(\n", 159 | " batch_size, only_non_self_colliding=True\n", 160 | " )\n", 161 | " target_poses = torch.tensor(target_poses, device=device, dtype=torch.float32)\n", 162 | " t0 = time()\n", 163 | " _, valid_solutions = ikflow_solver.generate_exact_ik_solutions(\n", 164 | " target_poses,\n", 165 | " pos_error_threshold=POS_ERROR_THRESHOLD,\n", 166 | " rot_error_threshold=ROT_ERROR_THRESHOLD,\n", 167 | " repeat_counts=repeat_counts,\n", 168 | " )\n", 169 | " sub_runtimes.append(time() - t0)\n", 170 | " sub_success_pcts.append(100 * valid_solutions.sum().item() / batch_size)\n", 171 | "\n", 172 | " runtimes.append(np.mean(sub_runtimes))\n", 173 | " runtime_stds.append(np.std(sub_runtimes))\n", 174 | "\n", 175 | " success_pcts.append(np.mean(sub_success_pcts))\n", 176 | " success_pct_stds.append(np.std(sub_success_pcts))\n", 177 | "\n", 178 | " print(f\"batch_size: {batch_size}, runtime: {runtimes[-1]:.3f}s\")\n", 179 | "\n", 180 | " curves.append(\n", 181 | " (\n", 182 | " np.array(batch_sizes),\n", 183 | " np.array(runtimes),\n", 184 | " np.array(runtime_stds),\n", 185 | " np.array(success_pcts),\n", 186 | " np.array(success_pct_stds),\n", 187 | " )\n", 188 | " )\n", 189 | " labels.append(f\"repeat_counts: {repeat_counts}\")" 190 | ] 191 | }, 192 | { 193 | "cell_type": "code", 194 | "execution_count": null, 195 | "id": "835aa998", 196 | "metadata": {}, 197 | "outputs": [], 198 | "source": [ 199 | "plot_runtime_curve(curves, labels)" 200 | ] 201 | }, 202 | { 203 | "cell_type": "code", 204 | "execution_count": null, 205 | "id": "ac3f346f", 206 | "metadata": {}, 207 | "outputs": [], 208 | "source": [ 209 | "x = torch.zeros((1000, 8), device=\"cuda:0\")\n", 210 | "t0 = time()\n", 211 | "x_cpu = x.cpu()\n", 212 | "print(f\"cpu time: {1000 * (time() - t0):.8f}ms\")" 213 | ] 214 | } 215 | ], 216 | "metadata": { 217 | "kernelspec": { 218 | "display_name": "uv run python 3 (ipykernel)", 219 | "language": "python", 220 | "name": "python3" 221 | }, 222 | "language_info": { 223 | "codemirror_mode": { 224 | "name": "ipython", 225 | "version": 3 226 | }, 227 | "file_extension": ".py", 228 | "mimetype": "text/x-python", 229 | "name": "python", 230 | "nbconvert_exporter": "python", 231 | "pygments_lexer": "ipython3", 232 | "version": "3.8.10" 233 | } 234 | }, 235 | "nbformat": 4, 236 | "nbformat_minor": 5 237 | } 238 | -------------------------------------------------------------------------------- /notebooks/inference_optimization.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "id": "d619eab0", 6 | "metadata": {}, 7 | "source": [ 8 | "# Overview\n", 9 | "\n", 10 | "This notebook analyzes the use of levenberg-marquardt for IK refinement" 11 | ] 12 | }, 13 | { 14 | "cell_type": "code", 15 | "execution_count": null, 16 | "id": "f124121e", 17 | "metadata": {}, 18 | "outputs": [], 19 | "source": [ 20 | "%load_ext autoreload\n", 21 | "%autoreload 2\n", 22 | "\n", 23 | "from IPython.display import display, HTML\n", 24 | "display(HTML(\"\"))" 25 | ] 26 | }, 27 | { 28 | "cell_type": "code", 29 | "execution_count": null, 30 | "id": "c10b367e", 31 | "metadata": {}, 32 | "outputs": [], 33 | "source": [ 34 | "from time import time \n", 35 | "\n", 36 | "import torch\n", 37 | "import numpy as np\n", 38 | "from jrl.utils import evenly_spaced_colors\n", 39 | "import matplotlib.pyplot as plt\n", 40 | "\n", 41 | "from ikflow.model_loading import get_ik_solver\n", 42 | "\n", 43 | "torch.set_printoptions(linewidth=300, precision=6, sci_mode=False)" 44 | ] 45 | }, 46 | { 47 | "cell_type": "code", 48 | "execution_count": null, 49 | "id": "f7298640", 50 | "metadata": {}, 51 | "outputs": [], 52 | "source": [ 53 | "MODEL_NAME = \"panda__full__lp191_5.25m\"\n", 54 | "compile_1_cfg = {\"fullgraph\": False, \"mode\": \"default\", \"backend\": \"inductor\"}\n", 55 | "compile_2_cfg = {\"fullgraph\": True, \"mode\": \"default\" , \"backend\": \"inductor\"}\n", 56 | "compile_3_cfg = {\"fullgraph\": False, \"mode\": \"default\" , \"backend\": \"inductor\", \"dynamic\": True}\n", 57 | "\n", 58 | "\n", 59 | "# === Errors\n", 60 | "# {\"fullgraph\": True, \"mode\": \"default\" , \"backend\": \"cudagraphs\"}\n", 61 | "# ^ doesn't work (RuntimeError: Node '5bec40': [(7,)] -> GLOWCouplingBlock -> [(7,)] encountered an error.)\n", 62 | "# {\"fullgraph\": True, \"mode\": \"max-autotune\" , \"backend\": \"inductor\"}\n", 63 | "# ^ doesn't work (RuntimeError: Node '5bec40': [(7,)] -> GLOWCouplingBlock -> [(7,)] encountered an error.)\n", 64 | "# {\"fullgraph\": False, \"mode\": \"max-autotune\" , \"backend\": \"inductor\"}\n", 65 | "# ^ really slow\n", 66 | "# {\"fullgraph\": False, \"mode\": \"default\" , \"backend\": \"onnxrt\", \"dynamic\": True}\n", 67 | "# ^ onyx unavailable\n", 68 | "# {\"fullgraph\": True, \"mode\": \"reduce-overhead\" , \"backend\": \"inductor\", \"dynamic\": True}\n", 69 | "# ^ really slow\n", 70 | "\n", 71 | "# >>> torch._dynamo.list_backends()\n", 72 | "# ['cudagraphs', 'inductor', 'onnxrt', 'openxla', 'openxla_eval', 'tvm']\n", 73 | "\n", 74 | "solver, _ = get_ik_solver(MODEL_NAME, compile_model=None)\n", 75 | "robot = solver.robot\n", 76 | "solver_compiled_1, _ = get_ik_solver(MODEL_NAME, compile_model=compile_1_cfg)\n", 77 | "solver_compiled_2, _ = get_ik_solver(MODEL_NAME, compile_model=compile_2_cfg)\n", 78 | "solver_compiled_3, _ = get_ik_solver(MODEL_NAME, compile_model=compile_3_cfg)" 79 | ] 80 | }, 81 | { 82 | "cell_type": "code", 83 | "execution_count": null, 84 | "id": "7abae5a8", 85 | "metadata": {}, 86 | "outputs": [], 87 | "source": [ 88 | "batch_sizes = [10, 100, 500, 1000, 2500, 5000, 7500, 10000]\n", 89 | "k_retry = 10\n", 90 | "\n", 91 | "curves = []\n", 92 | "labels = []\n", 93 | "device = \"cuda:0\"\n", 94 | "\n", 95 | "\n", 96 | "def eval_runtime(solver):\n", 97 | "\n", 98 | " runtimes = []\n", 99 | " runtime_stds = []\n", 100 | "\n", 101 | " for batch_size in batch_sizes:\n", 102 | " sub_runtimes = []\n", 103 | " for _ in range(k_retry):\n", 104 | " _, target_poses = robot.sample_joint_angles_and_poses(batch_size, only_non_self_colliding=False)\n", 105 | " target_poses = torch.tensor(target_poses, device=device, dtype=torch.float32)\n", 106 | " t0 = time()\n", 107 | " solver.generate_ik_solutions(target_poses)\n", 108 | " sub_runtimes.append(time() - t0)\n", 109 | " runtimes.append(np.mean(sub_runtimes))\n", 110 | " runtime_stds.append(np.std(sub_runtimes))\n", 111 | " print(f\"batch_size: {batch_size},\\tmean runtime: {runtimes[-1]:.5f} s,\\tsub_runtimes: {sub_runtimes}\")\n", 112 | " return np.array(runtimes), np.array(runtime_stds)\n", 113 | "\n", 114 | "base_runtimes, base_runtime_stds = eval_runtime(solver)\n", 115 | "compiled_1_runtimes, compiled_1_runtime_stds = eval_runtime(solver_compiled_1)\n", 116 | "compiled_2_runtimes, compiled_2_runtime_stds = eval_runtime(solver_compiled_2)\n", 117 | "compiled_3_runtimes, compiled_3_runtime_stds = eval_runtime(solver_compiled_3)" 118 | ] 119 | }, 120 | { 121 | "cell_type": "code", 122 | "execution_count": null, 123 | "id": "d61c36b5", 124 | "metadata": {}, 125 | "outputs": [], 126 | "source": [ 127 | "def plot_runtime_curve(multi_runtimes_means, multi_runtimes_stds, labels):\n", 128 | " fig, ax = plt.subplots(1, 1, figsize=(10, 10))\n", 129 | " ax.set_title(\"Batch size vs. Inference time for default vs. compiled nn.Module\")\n", 130 | " ax.grid(alpha=0.2)\n", 131 | " ax.set_xlabel(\"Batch size\")\n", 132 | " ax.set_ylabel(\"Inference time [s]\")\n", 133 | "\n", 134 | " colors = evenly_spaced_colors(int(len(multi_runtimes_means)))\n", 135 | "\n", 136 | " for (runtimes, runtime_stds, label, color) in zip(multi_runtimes_means, multi_runtimes_stds, labels, colors):\n", 137 | " ax.plot(batch_sizes, runtimes, label=label, color=color)\n", 138 | " ax.fill_between(batch_sizes, runtimes - runtime_stds, runtimes + runtime_stds, alpha=0.15, color=color)\n", 139 | " ax.scatter(batch_sizes, runtimes, s=15, color=color)\n", 140 | "\n", 141 | " ax.legend()\n", 142 | " plt.show()\n", 143 | "\n", 144 | "plot_runtime_curve(\n", 145 | " [base_runtimes, compiled_1_runtimes, compiled_2_runtimes, compiled_3_runtimes],\n", 146 | " [base_runtime_stds, compiled_1_runtime_stds, compiled_2_runtime_stds, compiled_3_runtime_stds], \n", 147 | " [\"un-compiled\", f\"compiled_1: {compile_1_cfg}\", f\"compiled_2: {compile_2_cfg}\", f\"compiled_3: {compile_3_cfg}\"]\n", 148 | ")\n" 149 | ] 150 | } 151 | ], 152 | "metadata": { 153 | "kernelspec": { 154 | "display_name": "ikflow-FSwUzYdQ-py3.8", 155 | "language": "python", 156 | "name": "python3" 157 | }, 158 | "language_info": { 159 | "codemirror_mode": { 160 | "name": "ipython", 161 | "version": 3 162 | }, 163 | "file_extension": ".py", 164 | "mimetype": "text/x-python", 165 | "name": "python", 166 | "nbconvert_exporter": "python", 167 | "pygments_lexer": "ipython3", 168 | "version": "3.8.10" 169 | } 170 | }, 171 | "nbformat": 4, 172 | "nbformat_minor": 5 173 | } 174 | -------------------------------------------------------------------------------- /notebooks/robot_visualizations.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "id": "d619eab0", 6 | "metadata": {}, 7 | "source": [ 8 | "# Overview\n", 9 | "This notebook illustrates the impact of the latent noise distribution on the resulting error statistics of generated samples" 10 | ] 11 | }, 12 | { 13 | "cell_type": "code", 14 | "execution_count": null, 15 | "id": "f124121e", 16 | "metadata": {}, 17 | "outputs": [], 18 | "source": [ 19 | "%load_ext autoreload\n", 20 | "%autoreload 2\n", 21 | "%load_ext wurlitzer\n", 22 | "# ^^^ the wurlitzer extension is used to capture C/C++ output to be displayed in the notebook\n", 23 | "# ^^^ this is very useful for debugging, but it doesn't work on windows\n", 24 | "import os\n", 25 | "\n", 26 | "os.chdir(\"../\")" 27 | ] 28 | }, 29 | { 30 | "cell_type": "code", 31 | "execution_count": null, 32 | "id": "c10b367e", 33 | "metadata": {}, 34 | "outputs": [], 35 | "source": [ 36 | "import numpy as np\n", 37 | "from klampt.model import coordinates, trajectory\n", 38 | "from klampt import *\n", 39 | "from klampt.math import so3, se3\n", 40 | "from klampt import vis\n", 41 | "\n", 42 | "from ikflow.model_loading import get_ik_solver" 43 | ] 44 | }, 45 | { 46 | "cell_type": "code", 47 | "execution_count": null, 48 | "id": "18599224", 49 | "metadata": {}, 50 | "outputs": [], 51 | "source": [ 52 | "# PANDA\n", 53 | "model_name = \"panda__full__lp191_5.25m\"\n", 54 | "ik_solver, _ = get_ik_solver(model_name)\n", 55 | "urdf_filepath = ik_solver.robot.urdf_filepath\n", 56 | "R = so3.from_rpy([0, np.pi / 2, 0])\n", 57 | "target_T_iktarget = (R, [-0.1, 0, 0])" 58 | ] 59 | }, 60 | { 61 | "cell_type": "markdown", 62 | "id": "c67d761a", 63 | "metadata": {}, 64 | "source": [ 65 | "# Visualizing\n", 66 | "\n", 67 | "Per ieee: \"One column width: 3.5 inches, 88.9 millimeters, or 21 picas \"\n", 68 | "\n", 69 | "Note: Don't change the color here. The standalone gui has a tool to change the appearance properties of objects in the world\n", 70 | "\n", 71 | "When positioning the camera, try to fill the frame almost entirely with the robot." 72 | ] 73 | }, 74 | { 75 | "cell_type": "code", 76 | "execution_count": null, 77 | "id": "e1d923f0", 78 | "metadata": {}, 79 | "outputs": [], 80 | "source": [ 81 | "SHOW_MUG = True\n", 82 | "SOL_MODE = \"IKFLOW\"\n", 83 | "N_SOLUTIONS = 10" 84 | ] 85 | }, 86 | { 87 | "cell_type": "code", 88 | "execution_count": null, 89 | "id": "a8d966ef", 90 | "metadata": {}, 91 | "outputs": [], 92 | "source": [ 93 | "world = ik_solver.robot._klampt_world_model\n", 94 | "if SHOW_MUG:\n", 95 | " assert world.loadRigidObject(\"ikflow/visualization_resources/klampt_resources/objects/mug_klampt.obj\")\n", 96 | " mug = world.rigidObject(0)\n", 97 | " vis.setColor(vis.getItemName(mug), 0.5, 0.5, 1.0, a=1.0)\n", 98 | "robot = world.robot(0)" 99 | ] 100 | }, 101 | { 102 | "cell_type": "code", 103 | "execution_count": null, 104 | "id": "41b13968", 105 | "metadata": {}, 106 | "outputs": [], 107 | "source": [ 108 | "# Open viz in pop up window.\n", 109 | "vis.kill()\n", 110 | "vis.init([\"PyQt\"])\n", 111 | "vis.add(\"world\", world)\n", 112 | "vis.add(\"coordinates\", coordinates.manager())\n", 113 | "background_color = (1, 1, 1, 0.7)\n", 114 | "vis.setBackgroundColor(background_color[0], background_color[1], background_color[2], background_color[3])\n", 115 | "size = 5\n", 116 | "for x0 in range(-size, size + 1):\n", 117 | " for y0 in range(-size, size + 1):\n", 118 | " vis.add(\n", 119 | " f\"floor_{x0}_{y0}\",\n", 120 | " trajectory.Trajectory([1, 0], [(-size, y0, 0), (size, y0, 0)]),\n", 121 | " color=(0.75, 0.75, 0.75, 1.0),\n", 122 | " width=2.0,\n", 123 | " hide_label=True,\n", 124 | " pointSize=0,\n", 125 | " )\n", 126 | " vis.add(\n", 127 | " f\"floor_{x0}_{y0}2\",\n", 128 | " trajectory.Trajectory([1, 0], [(x0, -size, 0), (x0, size, 0)]),\n", 129 | " color=(0.75, 0.75, 0.75, 1.0),\n", 130 | " width=2.0,\n", 131 | " hide_label=True,\n", 132 | " pointSize=0,\n", 133 | " )\n", 134 | "vis.show()" 135 | ] 136 | }, 137 | { 138 | "cell_type": "code", 139 | "execution_count": null, 140 | "id": "518be1e9", 141 | "metadata": {}, 142 | "outputs": [], 143 | "source": [ 144 | "# Tf parameters\n", 145 | "R = so3.from_rpy([np.pi / 2, 0, -0.3])\n", 146 | "target_T_mug = (R, [0.065, 0.005, -0.005])\n", 147 | "\n", 148 | "# Robot parameters\n", 149 | "alpha = 1\n", 150 | "# robot_color = (1, 1, 1, alpha)\n", 151 | "robot_color = (0.5, 0.5, 0.5, alpha)\n", 152 | "\n", 153 | "\n", 154 | "def get_color(i, n):\n", 155 | " return robot_color\n", 156 | "\n", 157 | "\n", 158 | "def update_sols(_end_pose):\n", 159 | " print(f\"_end_pose = {_end_pose.tolist()}\")\n", 160 | " if SOL_MODE == \"IKFLOW\":\n", 161 | " sols = ik_solver.generate_ik_solutions(_end_pose, N_SOLUTIONS).cpu().numpy()\n", 162 | " elif SOL_MODE == \"IK\":\n", 163 | " sols = ik_solver.robot.inverse_kinematics(_end_pose, N_SOLUTIONS, debug=False)\n", 164 | " else:\n", 165 | " raise ValueError(\"Unknown mode\")\n", 166 | "\n", 167 | " qs = ik_solver.robot._x_to_qs(sols)\n", 168 | " for i in range(N_SOLUTIONS):\n", 169 | " q = qs[i]\n", 170 | " rcolor = get_color(i, N_SOLUTIONS)\n", 171 | " if i == 0:\n", 172 | " vis.setColor(vis.getItemName(robot), rcolor[0], rcolor[1], rcolor[2], a=rcolor[3])\n", 173 | " robot.setConfig(q)\n", 174 | " continue\n", 175 | " vis.add(f\"robot_{i}\", q, color=rcolor)\n", 176 | "\n", 177 | "\n", 178 | "def update_scene():\n", 179 | " target = vis.getItemConfig(\"target\")\n", 180 | " world_T_target = (target[:-3], target[-3:])\n", 181 | " print(\"world_T_target =\", world_T_target)\n", 182 | " if SHOW_MUG:\n", 183 | " world_T_mug = se3.mul(world_T_target, target_T_mug)\n", 184 | " mug.setTransform(world_T_mug[0], world_T_mug[1])\n", 185 | " world_T_iktarget = se3.mul(world_T_target, target_T_iktarget)\n", 186 | " update_sols(np.array(world_T_iktarget[1] + so3.quaternion(world_T_iktarget[0])))\n", 187 | " vis.update()\n", 188 | "\n", 189 | "\n", 190 | "def current_state():\n", 191 | " target = vis.getItemConfig(\"target\")\n", 192 | " world_T_target = (target[:-3], target[-3:])\n", 193 | " print(\"world_T_target =\", world_T_target)\n", 194 | " print(\"target_T_iktarget =\", target_T_iktarget)\n", 195 | " vp = vis.getViewport()\n", 196 | " print(\"vp = vis.getViewport()\")\n", 197 | " print(\"vp.camera.rot =\", vp.camera.rot)\n", 198 | " print(\"vp.camera.tgt =\", vp.camera.tgt)\n", 199 | " print(\"vp.camera.dist =\", vp.camera.dist)\n", 200 | " print(\"vp.clippingplanes = \", vp.clippingplanes)\n", 201 | " print(\"vis.setViewport(vp)\")" 202 | ] 203 | }, 204 | { 205 | "cell_type": "markdown", 206 | "id": "f7ffa496", 207 | "metadata": {}, 208 | "source": [ 209 | "### Set IK target" 210 | ] 211 | }, 212 | { 213 | "cell_type": "code", 214 | "execution_count": null, 215 | "id": "9aecdbae", 216 | "metadata": {}, 217 | "outputs": [], 218 | "source": [ 219 | "# Update the target to a particular pose\n", 220 | "# --\n", 221 | "pose = [0.50, 0, 0.65, 1, 0, 0, 0]\n", 222 | "world_T_target = ([1.0, 0, 0, 0, 1.0, 0, 0, 0, 1.0], [0.35, 0, 0.65])\n", 223 | "vis.add(\"target\", world_T_target, length=0.25, width=1)\n", 224 | "\n", 225 | "vis.hideLabel(\"target\")\n", 226 | "vis.update()" 227 | ] 228 | }, 229 | { 230 | "cell_type": "markdown", 231 | "id": "a5be1fae", 232 | "metadata": {}, 233 | "source": [ 234 | "### Update scene" 235 | ] 236 | }, 237 | { 238 | "cell_type": "code", 239 | "execution_count": null, 240 | "id": "6aeba360", 241 | "metadata": {}, 242 | "outputs": [], 243 | "source": [ 244 | "# SOL_MODE = \"IK\"\n", 245 | "SOL_MODE = \"IKFLOW\"\n", 246 | "update_scene()" 247 | ] 248 | }, 249 | { 250 | "cell_type": "markdown", 251 | "id": "c527c889", 252 | "metadata": {}, 253 | "source": [ 254 | "### Show / Hide IK target" 255 | ] 256 | }, 257 | { 258 | "cell_type": "code", 259 | "execution_count": null, 260 | "id": "4029abec", 261 | "metadata": {}, 262 | "outputs": [], 263 | "source": [ 264 | "# SHOW\n", 265 | "vis.edit(\"target\", doedit=True)\n", 266 | "vis.setAttribute(\"target\", \"length\", 0.2)\n", 267 | "vis.hideLabel(\"target\")" 268 | ] 269 | }, 270 | { 271 | "cell_type": "code", 272 | "execution_count": null, 273 | "id": "9ce30708", 274 | "metadata": {}, 275 | "outputs": [], 276 | "source": [ 277 | "# HIDE\n", 278 | "vis.edit(\"target\", doedit=False)\n", 279 | "vis.setAttribute(\"target\", \"length\", 0.2)\n", 280 | "vis.hideLabel(\"target\")" 281 | ] 282 | }, 283 | { 284 | "cell_type": "markdown", 285 | "id": "53acc430", 286 | "metadata": {}, 287 | "source": [ 288 | "### Printout the current vis. state" 289 | ] 290 | }, 291 | { 292 | "cell_type": "code", 293 | "execution_count": null, 294 | "id": "ab76a43b", 295 | "metadata": {}, 296 | "outputs": [], 297 | "source": [ 298 | "current_state()" 299 | ] 300 | } 301 | ], 302 | "metadata": { 303 | "kernelspec": { 304 | "display_name": "venv", 305 | "language": "python", 306 | "name": "python3" 307 | }, 308 | "language_info": { 309 | "codemirror_mode": { 310 | "name": "ipython", 311 | "version": 3 312 | }, 313 | "file_extension": ".py", 314 | "mimetype": "text/x-python", 315 | "name": "python", 316 | "nbconvert_exporter": "python", 317 | "pygments_lexer": "ipython3", 318 | "version": "3.8.10" 319 | }, 320 | "vscode": { 321 | "interpreter": { 322 | "hash": "c106ee123d7af531b7acc4f77c56e24741db8737a45cbfd984346f16de511ef0" 323 | } 324 | } 325 | }, 326 | "nbformat": 4, 327 | "nbformat_minor": 5 328 | } 329 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "ikflow" 3 | version = "0.2.0" 4 | requires-python = ">=3.10,<3.12" 5 | description = "Open source implementation of the 'IKFlow' inverse kinematics solver" 6 | authors = [ 7 | { name="Jeremy Morgan", email="jsmorgan6@gmail.com" } 8 | ] 9 | readme = "README.md" 10 | dependencies = [ 11 | "freia==0.2", 12 | "jupyter>=1.1.1", 13 | "matplotlib>=3.7.5", 14 | "nbconvert>=7.16.6", 15 | "numpy==1.24.4", 16 | "pandas>=2.2.3", 17 | "pytorch-lightning>=2.5.1", 18 | "ruff>=0.11.2", 19 | "tabulate>=0.9.0", 20 | "tqdm>=4.67.1", 21 | "wandb>=0.19.8", 22 | "jrl @ git+https://github.com/jstmn/Jrl.git@ef4c2f6eb1ba84395ff0bb01d5b7713854df6908", 23 | "pyqt5>=5.15.11", 24 | "pytest>=8.3.5", 25 | ] 26 | 27 | [tool.setuptools.packages] 28 | find = { include = ["ikflow"], exclude = ["media", "examples", "notebooks"] } 29 | 30 | [tool.setuptools] 31 | include-package-data = true 32 | 33 | [tool.setuptools.package-data] 34 | ikflow = ["model_descriptions.yaml", "**/*.xml", "**/*.off", "**/*.obj"] 35 | 36 | 37 | [tool.ruff] 38 | line-length = 120 39 | lint.ignore = ["E741"] 40 | -------------------------------------------------------------------------------- /scripts/benchmark_generate_exact_solutions.py: -------------------------------------------------------------------------------- 1 | from typing import Callable, Optional 2 | from argparse import ArgumentParser 3 | from time import time 4 | import pickle 5 | 6 | import torch 7 | import numpy as np 8 | from jrl.utils import set_seed, evenly_spaced_colors 9 | from jrl.robots import Panda 10 | import matplotlib.pyplot as plt 11 | 12 | from ikflow.model import TINY_MODEL_PARAMS 13 | from ikflow.ikflow_solver import IKFlowSolver 14 | from ikflow.model_loading import get_ik_solver 15 | 16 | set_seed() 17 | 18 | POS_ERROR_THRESHOLD = 0.001 19 | ROT_ERROR_THRESHOLD = 0.01 20 | 21 | 22 | def fn_mean_std(fn: Callable, k: int): 23 | runtimes = [] 24 | for _ in range(k): 25 | t0 = time() 26 | fn() 27 | runtimes.append(1000 * (time() - t0)) 28 | return np.mean(runtimes), np.std(runtimes) 29 | 30 | 31 | def benchmark_get_pose_error(): 32 | """Benchmarks running IKFlowSolver._calculate_pose_error() on the CPU vs. the GPU""" 33 | 34 | robot = Panda() 35 | ikflow_solver = IKFlowSolver(TINY_MODEL_PARAMS, robot) 36 | 37 | retry_k = 5 38 | batch_sizes = [1, 10, 100, 500, 750, 1000, 2000, 3000, 4000, 5000] 39 | ts_cpu = [] 40 | ts_cu = [] 41 | stds_cpu = [] 42 | stds_cu = [] 43 | 44 | print("Number of solutions, CPU, CUDA") 45 | for n_solutions in batch_sizes: 46 | qs, target_poses = ikflow_solver.robot.sample_joint_angles_and_poses(n_solutions) 47 | qs_cpu = torch.tensor(qs, device="cpu", dtype=torch.float32) 48 | qs_cu = torch.tensor(qs, device="cuda:0", dtype=torch.float32) 49 | target_poses_cpu = torch.tensor(target_poses, device="cpu", dtype=torch.float32) 50 | target_poses_cu = torch.tensor(target_poses, device="cuda:0", dtype=torch.float32) 51 | 52 | t_cpu, t_cpu_std = fn_mean_std(lambda: ikflow_solver._calculate_pose_error(qs_cpu, target_poses_cpu), retry_k) 53 | t_cu, t_cu_std = fn_mean_std(lambda: ikflow_solver._calculate_pose_error(qs_cu, target_poses_cu), retry_k) 54 | ts_cpu.append(t_cpu) 55 | ts_cu.append(t_cu) 56 | stds_cpu.append(t_cpu_std) 57 | stds_cu.append(t_cu_std) 58 | print(n_solutions, "\t", round(t_cpu, 6), "\t", round(t_cu, 6)) 59 | 60 | # Plot results 61 | fig, ax = plt.subplots(1, 1, figsize=(8, 8)) 62 | fig.suptitle("Levenberg-Marquardt IK Convergence") 63 | 64 | ax.set_title("Runtime") 65 | ax.grid(alpha=0.2) 66 | ax.set_xlabel("batch size") 67 | ax.set_ylabel("runtime (s)") 68 | 69 | colors = evenly_spaced_colors(3) 70 | ts_cpu = np.array(ts_cpu) 71 | ts_cu = np.array(ts_cu) 72 | stds_cpu = np.array(stds_cpu) 73 | stds_cu = np.array(stds_cu) 74 | 75 | ax.plot(batch_sizes, ts_cpu, label="cpu", color=colors[0]) 76 | ax.fill_between(batch_sizes, ts_cpu - stds_cpu, ts_cpu + stds_cpu, alpha=0.15, color=colors[0]) 77 | ax.scatter(batch_sizes, ts_cpu, s=15, color=colors[0]) 78 | 79 | ax.plot(batch_sizes, ts_cu, label="cuda", color=colors[1]) 80 | ax.fill_between(batch_sizes, ts_cu - stds_cu, ts_cu + stds_cu, alpha=0.15, color=colors[1]) 81 | ax.scatter(batch_sizes, ts_cu, s=15, color=colors[1]) 82 | 83 | ax.legend() 84 | plt.savefig("get_pose_error_runtime_curve.pdf", bbox_inches="tight") 85 | plt.show() 86 | 87 | 88 | def get_stats(batch_sizes, k_retry, ikflow_solver, device, repeat_counts, run_lma_on_cpu): 89 | runtimes = [] 90 | success_pcts = [] 91 | runtime_stds = [] 92 | success_pct_stds = [] 93 | n_found = [] 94 | n_found_stds = [] 95 | 96 | for batch_size in batch_sizes: 97 | 98 | sub_runtimes = [] 99 | sub_success_pcts = [] 100 | sub_n_found = [] 101 | 102 | for _ in range(k_retry): 103 | _, target_poses = ikflow_solver.robot.sample_joint_angles_and_poses( 104 | batch_size, only_non_self_colliding=True 105 | ) 106 | target_poses = torch.tensor(target_poses, device=device, dtype=torch.float32) 107 | t0 = time() 108 | solutions, valid_solution_idxs = ikflow_solver.generate_exact_ik_solutions( 109 | target_poses, 110 | pos_error_threshold=POS_ERROR_THRESHOLD, 111 | rot_error_threshold=ROT_ERROR_THRESHOLD, 112 | repeat_counts=repeat_counts, 113 | run_lma_on_cpu=run_lma_on_cpu, 114 | ) 115 | sub_runtimes.append(time() - t0) 116 | sub_success_pcts.append(100 * valid_solution_idxs.sum().item() / batch_size) 117 | sub_n_found.append(valid_solution_idxs.sum().item()) 118 | print("should not be equal, if not all successful:", solutions.shape[0], valid_solution_idxs.sum().item()) 119 | 120 | runtimes.append(np.mean(sub_runtimes)) 121 | runtime_stds.append(np.std(sub_runtimes)) 122 | 123 | success_pcts.append(np.mean(sub_success_pcts)) 124 | success_pct_stds.append(np.std(sub_success_pcts)) 125 | 126 | # TODO: this isn't getting properly plotted 127 | n_found.append(np.mean(sub_n_found)) 128 | n_found_stds.append(np.std(sub_n_found)) 129 | print(f"batch_size: {batch_size}, runtime: {runtimes[-1]:.3f}s") 130 | print(f"batch_size: {batch_size}, runtime: {runtimes[-1]:.3f}s") 131 | return ( 132 | np.array(batch_sizes), 133 | np.array(runtimes), 134 | np.array(runtime_stds), 135 | np.array(success_pcts), 136 | np.array(success_pct_stds), 137 | np.array(n_found), 138 | np.array(n_found_stds), 139 | ) 140 | 141 | 142 | def benchmark_get_exact_ik(model_name: str, comparison_data_filepath: Optional[str] = None): 143 | ikflow_solver, _ = get_ik_solver(model_name) 144 | 145 | # batch_sizes = [1, 10, 100, 500, 750, 1000, 2000, 3000, 4000, 5000] 146 | batch_sizes = [1, 10, 100, 500, 750, 1000] 147 | k_retry = 3 148 | # all_repeat_counts = [(1, 3, 10), (1, 5, 10)] 149 | # all_repeat_counts = [(1, 3, 10), (1, 3)] 150 | all_repeat_counts = [(1, 3, 10), (1,)] 151 | device = "cuda:0" 152 | 153 | curves = [] 154 | labels = [] 155 | 156 | # for run_lma_on_cpu in [False, True]: 157 | for run_lma_on_cpu in [True]: 158 | for repeat_counts in all_repeat_counts: 159 | curves.append(get_stats(batch_sizes, k_retry, ikflow_solver, device, repeat_counts, run_lma_on_cpu)) 160 | # labels.append(f"repeat_counts: {repeat_counts}, run_lma_on_cpu: {run_lma_on_cpu}") 161 | labels.append(f"ikflow + lma - {repeat_counts} iterations") 162 | 163 | # Load comparison data from a pickle file 164 | if comparison_data_filepath is not None: 165 | with open(comparison_data_filepath, "rb") as f: 166 | data = pickle.load(f) 167 | assert abs(data["position_threshold"] - POS_ERROR_THRESHOLD) < 1e-7 168 | assert abs(data["rotation_threshold"] - ROT_ERROR_THRESHOLD) < 1e-7 169 | batch_sizes = data["batch_sizes"] 170 | runtimes = data["runtimes"] 171 | runtime_stds = data["runtime_stds"] 172 | success_pcts = data["success_pcts"] 173 | success_pct_stds = data["success_pct_stds"] 174 | n_found = data["n_found"] 175 | n_found_stds = data["n_found_stds"] 176 | curves.append(( 177 | np.array(batch_sizes), 178 | np.array(runtimes), 179 | np.array(runtime_stds), 180 | np.array(success_pcts), 181 | np.array(success_pct_stds), 182 | np.array(n_found), 183 | np.array(n_found_stds), 184 | )) 185 | labels.append(data["label"]) 186 | 187 | # Plot results 188 | # 189 | fig, (axl, axr, axrr) = plt.subplots(1, 3, figsize=(24, 8)) 190 | fig.tight_layout() 191 | fig.suptitle( 192 | "Number of Solutions vs. Runtime for IKFlow + Levenberg-Marquardt IK Solution Generation" 193 | f" ({POS_ERROR_THRESHOLD} m, {ROT_ERROR_THRESHOLD} rad tol)", 194 | fontsize=16, 195 | ) 196 | 197 | axl.set_title("Runtime") 198 | axl.grid(alpha=0.2) 199 | axl.set_xlabel("Number of Solutions") 200 | axl.set_ylabel("Runtime (s)") 201 | 202 | axr.set_title("Success Pct") 203 | axr.set_xlabel("Number of Solutions") 204 | axr.set_ylabel("Success percentage (%)") 205 | 206 | axrr.set_title("Number of solutions found") 207 | axrr.grid(alpha=0.2) 208 | axrr.set_xlabel("Number of Solutions Requested") 209 | axrr.set_ylabel("Number of Solutions Returned") 210 | 211 | colors = evenly_spaced_colors(int(1.5 * len(curves))) 212 | max_runtime = 0 213 | 214 | for ( 215 | (batch_sizes, runtimes, runtime_stds, success_pcts, success_pct_stds, n_found, n_found_stds), 216 | label, 217 | color, 218 | ) in zip(curves, labels, colors): 219 | max_runtime = max(max_runtime, runtimes.max()) 220 | axl.plot(batch_sizes, runtimes, label=label, color=color) 221 | axl.fill_between(batch_sizes, runtimes - runtime_stds, runtimes + runtime_stds, alpha=0.15, color=color) 222 | axl.scatter(batch_sizes, runtimes, s=15, color=color) 223 | 224 | axr.plot(batch_sizes, success_pcts, label=label, color=color) 225 | axr.fill_between( 226 | batch_sizes, success_pcts - success_pct_stds, success_pcts + success_pct_stds, alpha=0.15, color=color 227 | ) 228 | axr.scatter(batch_sizes, success_pcts, s=15, color=color) 229 | 230 | axrr.plot(batch_sizes, n_found, label=label, color=color) 231 | axrr.fill_between(batch_sizes, n_found - n_found_stds, n_found + n_found_stds, alpha=0.15, color=color) 232 | axrr.scatter(batch_sizes, n_found, s=15, color=color) 233 | 234 | # axl.legend() 235 | axl.set_ylim(-0.01, max_runtime * 1.1) 236 | axr.legend() 237 | axrr.legend() 238 | plt.savefig(f"exact_ik_runtime__model:{model_name}.pdf", bbox_inches="tight") 239 | plt.savefig(f"exact_ik_runtime__model:{model_name}.png", bbox_inches="tight") 240 | plt.show() 241 | 242 | 243 | """ Example usage 244 | 245 | uv run python scripts/benchmark_generate_exact_solutions.py --model_name=panda__full__lp191_5.25m 246 | uv run python scripts/benchmark_generate_exact_solutions.py --model_name=panda__full__lp191_5.25m --comparison_data_filepath=ik_runtime_results__curobo.pkl 247 | 248 | """ 249 | 250 | if __name__ == "__main__": 251 | parser = ArgumentParser() 252 | parser.add_argument("--model_name", type=str, required=True) 253 | parser.add_argument("--comparison_data_filepath", type=str, required=False) 254 | args = parser.parse_args() 255 | 256 | # benchmark_get_pose_error() 257 | benchmark_get_exact_ik(args.model_name, comparison_data_filepath=args.comparison_data_filepath) 258 | -------------------------------------------------------------------------------- /scripts/benchmark_runtime.py: -------------------------------------------------------------------------------- 1 | from time import time 2 | from typing import Callable 3 | import argparse 4 | 5 | import numpy as np 6 | import torch 7 | import pandas as pd 8 | import matplotlib.pyplot as plt 9 | from jrl.utils import set_seed 10 | 11 | from ikflow.model_loading import get_ik_solver 12 | from ikflow.ikflow_solver import IKFlowSolver 13 | 14 | torch.set_printoptions(linewidth=200, precision=6, sci_mode=False) 15 | 16 | set_seed() 17 | 18 | 19 | def fn_mean_std(fn: Callable, k: int): 20 | runtimes = [] 21 | for _ in range(k): 22 | t0 = time() 23 | fn() 24 | runtimes.append(1000 * (time() - t0)) 25 | return np.mean(runtimes), np.std(runtimes) 26 | 27 | 28 | POS_ERROR_THRESHOLD = 0.001 29 | ROT_ERROR_THRESHOLD = 0.01 30 | 31 | 32 | def solve_ikflow(ikflow_solver: IKFlowSolver, target_poses: torch.Tensor): 33 | ikflow_solver.generate_ik_solutions(target_poses) 34 | 35 | 36 | def solve_klampt(ikflow_solver: IKFlowSolver, target_poses: torch.Tensor): 37 | n_failed = 0 38 | for i in range(len(target_poses)): 39 | q_i = ikflow_solver.robot.inverse_kinematics_klampt(target_poses[i]) 40 | if q_i is None: 41 | n_failed += 1 42 | 43 | 44 | def solve_klampt_ikflow_seed(ikflow_solver: IKFlowSolver, target_poses: torch.Tensor): 45 | qs = ikflow_solver.generate_ik_solutions(target_poses).cpu().numpy() 46 | target_poses = target_poses.cpu().numpy() 47 | n_failed = 0 48 | for i in range(len(target_poses)): 49 | q_i = ikflow_solver.robot.inverse_kinematics_klampt(target_poses[i], seed=qs[i, :]) 50 | if q_i is None: 51 | n_failed += 1 52 | 53 | 54 | def solve_lma(ikflow_solver: IKFlowSolver, target_poses: torch.Tensor): 55 | ikflow_solver.generate_exact_ik_solutions(target_poses, run_lma_on_cpu=True) 56 | 57 | 58 | """ Example 59 | 60 | uv run python scripts/benchmark_runtime.py --model_name=panda__full__lp191_5.25m 61 | 62 | TODO: report success pct as well. klampt / lma are going to have a different % exact solutions generated 63 | """ 64 | 65 | if __name__ == "__main__": 66 | parser = argparse.ArgumentParser(prog="evaluate.py - evaluates IK models") 67 | parser.add_argument( 68 | "--model_name", type=str, required=True, help="Name of the saved model (see ikflow/model_descriptions.yaml)" 69 | ) 70 | parser.add_argument("--k", type=int, default=3, help="Number of repeats per batch size") 71 | args = parser.parse_args() 72 | 73 | df = pd.DataFrame( 74 | columns=["method", "number of solutions", "total runtime (ms)", "runtime std", "runtime per solution (ms)"] 75 | ) 76 | method_names = ["ikflow - NOT EXACT", "klampt", "ikflow with levenberg-marquardt", "klampt with ikflow seeds"] 77 | 78 | ikflow_solver, _ = get_ik_solver(args.model_name) 79 | 80 | for batch_size in [2, 5, 10, 50, 100, 500, 1000, 2500]: 81 | print(f"\n---\nBatch size: {batch_size}") 82 | 83 | _, target_poses = ikflow_solver.robot.sample_joint_angles_and_poses( 84 | batch_size, only_non_self_colliding=True, tqdm_enabled=False 85 | ) 86 | target_poses = torch.tensor(target_poses, dtype=torch.float32, device="cuda:0") 87 | 88 | lambdas = [ 89 | lambda: solve_ikflow(ikflow_solver, target_poses.clone()), 90 | lambda: solve_klampt(ikflow_solver, target_poses.cpu().numpy()), 91 | lambda: solve_lma(ikflow_solver, target_poses.clone()), 92 | lambda: solve_klampt_ikflow_seed(ikflow_solver, target_poses.clone()), 93 | ] 94 | for lambda_, method_name in zip(lambdas, method_names): 95 | mean_runtime_ms, std_runtime = fn_mean_std(lambda_, args.k) 96 | new_row = [method_name, batch_size, mean_runtime_ms, std_runtime, mean_runtime_ms / batch_size] 97 | df.loc[len(df)] = new_row 98 | 99 | df = df.sort_values(by=["method", "number of solutions"]) 100 | 101 | print(df) 102 | 103 | # Plot 104 | max_runtime = -1 105 | fig, ax = plt.subplots(1, 1, figsize=(12, 10)) 106 | ax.grid(alpha=0.2) 107 | for method_name in method_names: 108 | df_method = df[df["method"] == method_name] 109 | n_solutions = df_method["number of solutions"] 110 | total_runtime_ms = df_method["total runtime (ms)"] 111 | std = df_method["runtime std"] 112 | ax.plot(n_solutions, total_runtime_ms, label=method_name) 113 | ax.fill_between(n_solutions, total_runtime_ms - std, total_runtime_ms + std, alpha=0.2) 114 | max_runtime = max(max_runtime, total_runtime_ms.to_numpy()[-1]) 115 | 116 | ax.set_ylim(-0.1, max_runtime + 0.5) 117 | ax.set_title("Number of solutions vs runtime for geodesic_distance_between_quaternions()") 118 | ax.set_xlabel("Number of solutions") 119 | ax.set_ylabel("Total runtime (ms)") 120 | ax.legend() 121 | fig.savefig("ik_benchmark_runtime.pdf", bbox_inches="tight") 122 | plt.show() 123 | -------------------------------------------------------------------------------- /scripts/build_dataset.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from typing import Optional, List 3 | import os 4 | from time import time 5 | 6 | from ikflow.config import DATASET_DIR, ALL_DATASET_TAGS, DATASET_TAG_NON_SELF_COLLIDING 7 | from ikflow.utils import ( 8 | get_dataset_directory, 9 | safe_mkdir, 10 | print_tensor_stats, 11 | get_sum_joint_limit_range, 12 | get_dataset_filepaths, 13 | assert_joint_angle_tensor_in_joint_limits, 14 | ) 15 | from jrl.robots import get_robot, Robot 16 | 17 | import torch 18 | 19 | TRAINING_SET_SIZE_SMALL = int(1e5) 20 | TEST_SET_SIZE = 15000 21 | 22 | 23 | def print_saved_datasets_stats(tags: List[str], robots: Optional[Robot] = []): 24 | """Printout summary statistics for each dataset. Optionally print out the default joint limits of all robots in 25 | `robots` 26 | """ 27 | 28 | print("\tSummary info on all saved datasets:") 29 | 30 | def print_joint_limits(limits, robot=""): 31 | print(f"print_joint_limits('{robot}')") 32 | sum_range = 0 33 | for idx, (l, u) in enumerate(limits): 34 | sum_range += u - l 35 | print(f" joint_{idx}: ({l},\t{u})") 36 | print(f" sum_range: {sum_range}") 37 | 38 | for robot in robots: 39 | print_joint_limits(robot.actuated_joints_limits, robot=robot) 40 | 41 | sp = "\t" 42 | print(f"\nrobot {sp} dataset_name {sp} sum_joint_range") 43 | print(f"----- {sp} ------------ {sp} ---------------") 44 | 45 | # For each child directory (for each robot) in config.DATASET_DIR: 46 | for dataset_directory, dirs, files in os.walk(DATASET_DIR): 47 | # Ignore the config.DATASET_DIR itself. os.walk returns the parent direcory followed by all child 48 | # directories 49 | if len(dirs) > 0: 50 | continue 51 | 52 | dataset_name = dataset_directory.split("/")[-1] 53 | robot_name = dataset_name.split("_")[0] 54 | 55 | try: 56 | samples_tr_file_path, _, _, _, _ = get_dataset_filepaths(dataset_directory, tags) 57 | samples_tr = torch.load(samples_tr_file_path) 58 | sum_joint_range = get_sum_joint_limit_range(samples_tr) 59 | print(f"{robot_name} {sp} {dataset_name} {sp} {sum_joint_range}") 60 | except FileNotFoundError: 61 | samples_tr_file_path, _, _, _, _ = get_dataset_filepaths(dataset_directory, ALL_DATASET_TAGS) 62 | samples_tr = torch.load(samples_tr_file_path) 63 | sum_joint_range = get_sum_joint_limit_range(samples_tr) 64 | print(f"{robot_name} {sp} {dataset_name} {sp} {sum_joint_range}") 65 | 66 | 67 | # TODO(@jeremysm): Consider adding plots from generated datasets. See `plot_dataset_samples()` in 'Dataset visualization 68 | # and validation.ipynb' 69 | def save_dataset_to_disk( 70 | robot: Robot, 71 | dataset_directory: str, 72 | training_set_size: int, 73 | test_set_size: int, 74 | only_non_self_colliding: bool, 75 | tags: List[str], 76 | joint_limit_eps: float = 1e-6, 77 | ): 78 | """ 79 | Save training & testset numpy arrays to the provided directory 80 | """ 81 | 82 | safe_mkdir(dataset_directory) 83 | if only_non_self_colliding: 84 | assert DATASET_TAG_NON_SELF_COLLIDING in tags, ( 85 | f"Error: If `only_non_self_colliding` is True, then the tag '{DATASET_TAG_NON_SELF_COLLIDING}' should be in" 86 | " `tags`" 87 | ) 88 | 89 | # Build testset 90 | samples_te, poses_te = robot.sample_joint_angles_and_poses( 91 | test_set_size, 92 | joint_limit_eps=joint_limit_eps, 93 | only_non_self_colliding=only_non_self_colliding, 94 | tqdm_enabled=True, 95 | ) 96 | samples_tr, poses_tr = robot.sample_joint_angles_and_poses( 97 | training_set_size, 98 | joint_limit_eps=joint_limit_eps, 99 | only_non_self_colliding=only_non_self_colliding, 100 | tqdm_enabled=True, 101 | ) 102 | 103 | # Save training set 104 | ( 105 | samples_tr_file_path, 106 | poses_tr_file_path, 107 | samples_te_file_path, 108 | poses_te_file_path, 109 | info_filepath, 110 | ) = get_dataset_filepaths(dataset_directory, tags) 111 | 112 | samples_tr = torch.tensor(samples_tr, dtype=torch.float32) 113 | poses_tr = torch.tensor(poses_tr, dtype=torch.float32) 114 | samples_te = torch.tensor(samples_te, dtype=torch.float32) 115 | poses_te = torch.tensor(poses_te, dtype=torch.float32) 116 | 117 | # Sanity check 118 | for arr in [samples_tr, samples_te, poses_tr, poses_te]: 119 | for i in range(arr.shape[1]): 120 | assert torch.std(arr[:, i]) > 0.001, f"Error: Column {i} in samples has zero stdev" 121 | assert_joint_angle_tensor_in_joint_limits(robot.actuated_joints_limits, samples_tr, "samples_tr", 0.0) 122 | assert_joint_angle_tensor_in_joint_limits(robot.actuated_joints_limits, samples_te, "samples_te", 0.0) 123 | 124 | with open(info_filepath, "w") as f: 125 | f.write("Dataset info") 126 | f.write(f" robot: {robot.name}\n") 127 | f.write(f" dataset_directory: {dataset_directory}\n") 128 | f.write(f" training_set_size: {training_set_size}\n") 129 | f.write(f" test_set_size: {test_set_size}\n") 130 | 131 | # Sanity check arrays. 132 | print_tensor_stats(samples_tr, writable=f, name="samples_tr") 133 | print_tensor_stats(poses_tr, writable=f, name="poses_tr") 134 | print_tensor_stats(samples_te, writable=f, name="samples_te") 135 | print_tensor_stats(poses_te, writable=f, name="poses_te") 136 | 137 | torch.save(samples_tr, samples_tr_file_path) 138 | torch.save(poses_tr, poses_tr_file_path) 139 | torch.save(samples_te, samples_te_file_path) 140 | torch.save(poses_te, poses_te_file_path) 141 | 142 | 143 | def _get_tags(args): 144 | tags = [] 145 | if args.only_non_self_colliding: 146 | tags.append(DATASET_TAG_NON_SELF_COLLIDING) 147 | else: 148 | print("====" * 10) 149 | print("WARNING: Saving dataset with self-colliding configurations. This is not recommended.") 150 | tags.sort() 151 | for tag in tags: 152 | assert tag in ALL_DATASET_TAGS 153 | return tags 154 | 155 | 156 | """ 157 | # Build dataset 158 | 159 | uv run python scripts/build_dataset.py --robot_name=fetch --training_set_size=25000000 --only_non_self_colliding 160 | uv run python scripts/build_dataset.py --robot_name=panda --training_set_size=25000000 --only_non_self_colliding 161 | uv run python scripts/build_dataset.py --robot_name=fetch_arm --training_set_size=25000000 --only_non_self_colliding 162 | """ 163 | 164 | 165 | if __name__ == "__main__": 166 | parser = argparse.ArgumentParser() 167 | parser.add_argument("--robot_name", type=str) 168 | parser.add_argument("--training_set_size", type=int, default=int(2.5 * 1e6)) 169 | parser.add_argument("--only_non_self_colliding", action="store_true") 170 | args = parser.parse_args() 171 | 172 | robot = get_robot(args.robot_name) 173 | tags = _get_tags(args) 174 | 175 | # Build dataset 176 | print(f"Building dataset for robot: {robot}") 177 | dset_directory = get_dataset_directory(robot.name) 178 | t0 = time() 179 | save_dataset_to_disk( 180 | robot, 181 | dset_directory, 182 | args.training_set_size, 183 | TEST_SET_SIZE, 184 | args.only_non_self_colliding, 185 | tags, 186 | joint_limit_eps=0.004363323129985824, # np.deg2rad(0.25) 187 | ) 188 | print(f"Saved dataset with {args.training_set_size} samples in {time() - t0:.2f} seconds") 189 | print_saved_datasets_stats(tags) 190 | -------------------------------------------------------------------------------- /scripts/download_model_from_wandb_checkpoint.py: -------------------------------------------------------------------------------- 1 | from typing import Dict 2 | import argparse 3 | import os 4 | from time import time 5 | import pickle 6 | 7 | import wandb 8 | import torch 9 | 10 | from ikflow.utils import get_wandb_project 11 | 12 | 13 | def format_state_dict(state_dict: Dict) -> Dict: 14 | """The `state_dict` saved in checkpoints will have keys with the form: 15 | "nn_model.module_list.0.M", "nn_model.module_list.0.M_inv", ... 16 | 17 | This function updates them to the expected format below. (note the 'nn_model' prefix is removed) 18 | "module_list.0.M", "module_list.0.M_inv", ... 19 | """ 20 | bad_prefix = "nn_model." 21 | len_prefix = len(bad_prefix) 22 | updated = {} 23 | for k, v in state_dict.items(): 24 | # Check that the original state dict is malformatted first 25 | assert k[0:len_prefix] == bad_prefix 26 | k_new = k[len_prefix:] 27 | updated[k_new] = v 28 | return updated 29 | 30 | 31 | """ 32 | _____________ 33 | Example usage 34 | 35 | uv run python scripts/download_model_from_wandb_checkpoint.py --wandb_run_id=2uidt835 36 | """ 37 | 38 | if __name__ == "__main__": 39 | parser = argparse.ArgumentParser(prog="Download ikflow model from Weights and Biases") 40 | 41 | # Note: WandB saves artifacts by the run ID (i.e. '34c2gimi') not the run name ('dashing-forest-33'). This is 42 | # slightly annoying because you need to click on a run to get its ID. 43 | parser.add_argument("--wandb_run_id", type=str, help="The run ID of the wandb run to load. Example: '34c2gimi'") 44 | parser.add_argument("--disable_progress_bar", action="store_true") 45 | args = parser.parse_args() 46 | 47 | wandb_entity, wandb_project = get_wandb_project() 48 | t0 = time() 49 | api = wandb.Api() 50 | artifact = api.artifact(f"{wandb_entity}/{wandb_project}/model-{args.wandb_run_id}:best_k") 51 | download_dir = artifact.download() 52 | print(f"Downloaded artifact in {round(time() - t0, 2)}s") 53 | 54 | t0 = time() 55 | run = api.run(f"/{wandb_entity}/{wandb_project}/runs/{args.wandb_run_id}") 56 | print(f"Downloaded run data in {round(time() - t0, 2)}s") 57 | 58 | run_name = run.name 59 | robot_name = run.config["robot"] 60 | ckpt_filepath = os.path.join(download_dir, "model.ckpt") 61 | checkpoint = torch.load(ckpt_filepath, map_location=lambda storage, loc: storage) 62 | state_dict = format_state_dict(checkpoint["state_dict"]) 63 | global_step = str(checkpoint["global_step"] / 1e6) + "M" 64 | 65 | # Save model's state_dict 66 | # TODO(@jstmn): Save the global_step aswell 67 | model_state_dict_filepath = os.path.join(f"{robot_name}__{run_name}__global_step={global_step}.pkl") 68 | with open(model_state_dict_filepath, "wb") as f: 69 | pickle.dump(state_dict, f) 70 | -------------------------------------------------------------------------------- /scripts/evaluate_inference_speed.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from time import time 3 | 4 | from ikflow.model import IkflowModelParameters 5 | from ikflow.ikflow_solver import IKFlowSolver 6 | from jrl.robots import get_robot 7 | 8 | import torch 9 | import numpy as np 10 | 11 | 12 | def calculate_ave_runtime(ik_solver: IKFlowSolver, n_samples: int): 13 | """_summary_ 14 | 15 | Args: 16 | ik_solver (IKFlowSolver): _description_ 17 | n_samples (int): _description_ 18 | 19 | Returns: 20 | _type_: _description_ 21 | """ 22 | 23 | # Calculate average runtime for 100 samples 24 | sample_times = [] 25 | with torch.inference_mode(): 26 | for i in range(30): 27 | pose = np.random.random(7) 28 | t0 = time() 29 | ik_solver.generate_ik_solutions(pose, n_samples) 30 | sample_times.append(time() - t0) 31 | sample_time_100 = 1000 * np.mean(sample_times) # to milleseconds 32 | sample_time_100_std = np.std(sample_times) 33 | return sample_time_100, sample_time_100_std 34 | 35 | 36 | """ Example usage 37 | 38 | uv run python scripts/runtime_eval.py \ 39 | --nb_nodes_range 4 6 8 10 12 14 40 | """ 41 | 42 | if __name__ == "__main__": 43 | parser = argparse.ArgumentParser(prog="Runtime tester") 44 | 45 | # Note: WandB saves artifacts by the run ID (i.e. '34c2gimi') not the run name ('dashing-forest-33'). This is 46 | # slightly annoying because you need to click on a run to get its ID. 47 | parser.add_argument("--wandb_run_id_to_load_checkpoint", type=str, help="Example: '34c2gimi'") 48 | parser.add_argument("--robot_name", type=str) 49 | 50 | # Model parameters 51 | parser.add_argument("--dim_latent_space", type=int, default=9) 52 | parser.add_argument("--coeff_fn_config", type=int, default=3) 53 | parser.add_argument("--coeff_fn_internal_size", type=int, default=2) 54 | parser.add_argument("--nb_nodes_range", nargs="+", type=int) 55 | 56 | parser.add_argument( 57 | "--n_samples_for_runtime", 58 | default=100, 59 | type=int, 60 | help="Check the average runtime to get this number of solutions", 61 | ) 62 | 63 | args = parser.parse_args() 64 | 65 | hparams = IkflowModelParameters() 66 | hparams.dim_latent_space = args.dim_latent_space 67 | hparams.coeff_fn_config = args.coeff_fn_config 68 | hparams.coeff_fn_internal_size = args.coeff_fn_internal_size 69 | robot = get_robot("panda") 70 | 71 | print(f"Average runtime for {args.n_samples_for_runtime} samples") 72 | 73 | # TODO(@jstmn): Format output 74 | print("\nNumber of nodes\t| Runtime (ms)") 75 | print("----------------|-------------------") 76 | for nb_nodes in args.nb_nodes_range: 77 | hparams.nb_nodes = nb_nodes 78 | ik_solver = IKFlowSolver( 79 | hparams, 80 | robot, 81 | ) 82 | 83 | ave_runtime, std = calculate_ave_runtime(ik_solver, args.n_samples_for_runtime) 84 | print(f"{nb_nodes}\t\t| {ave_runtime:.2f}") 85 | -------------------------------------------------------------------------------- /scripts/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | 4 | from jrl.robots import get_robot 5 | from pytorch_lightning.loggers import WandbLogger 6 | from pytorch_lightning.callbacks import ModelCheckpoint 7 | from pytorch_lightning.trainer import Trainer 8 | from pytorch_lightning import seed_everything 9 | import wandb 10 | import torch 11 | from ikflow.config import DATASET_TAG_NON_SELF_COLLIDING 12 | from jrl.config import GPU_IDX 13 | 14 | from ikflow import config 15 | from ikflow.model import IkflowModelParameters 16 | from ikflow.ikflow_solver import IKFlowSolver 17 | from ikflow.training.lt_model import IkfLitModel 18 | from ikflow.training.lt_data import IkfLitDataset 19 | from ikflow.training.training_utils import get_checkpoint_dir 20 | from ikflow.utils import boolean_string, non_private_dict, get_wandb_project 21 | 22 | 23 | assert GPU_IDX >= 0 24 | DEFAULT_MAX_EPOCHS = 5000 25 | SEED = 0 26 | seed_everything(SEED, workers=True) 27 | 28 | # Model parameters 29 | DEFAULT_COUPLING_LAYER = "glow" 30 | DEFAULT_RNVP_CLAMP = 2.5 31 | DEFAULT_SOFTFLOW_NOISE_SCALE = 0.001 32 | DEFAULT_SOFTFLOW_ENABLED = True 33 | DEFAULT_N_NODES = 6 34 | DEFAULT_DIM_LATENT_SPACE = 8 35 | DEFAULT_COEFF_FN_INTERNAL_SIZE = 1024 36 | DEFAULT_COEFF_FN_CONFIG = 3 37 | DEFAULT_Y_NOISE_SCALE = 1e-7 38 | DEFAULT_ZEROS_NOISE_SCALE = 1e-3 39 | DEFAULT_SIGMOID_ON_OUTPUT = False 40 | 41 | # Training parameters 42 | DEFAULT_OPTIMIZER = "adamw" 43 | DEFAULT_LR = 1e-4 44 | DEFAULT_BATCH_SIZE = 512 45 | DEFAULT_N_EPOCHS = 100 46 | DEFAULT_GAMMA = 0.9794578299341784 47 | DEFAULT_STEP_LR_EVERY = int(int(2.5 * 1e6) / 64) 48 | DEFAULT_GRADIENT_CLIP_VAL = 1 49 | 50 | 51 | # Logging stats 52 | DEFAULT_EVAL_EVERY = 20000 53 | DEFAULT_VAL_SET_SIZE = 500 54 | DEFAULT_LOG_EVERY = 20000 55 | DEFAULT_CHECKPOINT_EVERY = 250000 56 | 57 | 58 | """ 59 | _____________ 60 | Example usage 61 | 62 | ROBOT=fr3 63 | 64 | # Smoke test - wandb enabled 65 | uv run python scripts/train.py \ 66 | --robot_name=$ROBOT \ 67 | --nb_nodes=6 \ 68 | --batch_size=128 \ 69 | --learning_rate=0.0005 \ 70 | --log_every=125 \ 71 | --eval_every=50 \ 72 | --val_set_size=20 \ 73 | --checkpoint_every=250 74 | 75 | 76 | # Smoke test - wandb enabled, with sigmoid clamping 77 | uv run python scripts/train.py \ 78 | --robot_name=$ROBOT \ 79 | --nb_nodes=6 \ 80 | --batch_size=64 \ 81 | --learning_rate=0.0005 \ 82 | --log_every=250 \ 83 | --eval_every=250 \ 84 | --val_set_size=250 \ 85 | --softflow_enabled=False \ 86 | --sigmoid_on_output=True \ 87 | --dim_latent_space=10 \ 88 | --checkpoint_every=50000 89 | 90 | 91 | # Smoke test - wandb disabled 92 | uv run python scripts/train.py \ 93 | --robot_name=$ROBOT \ 94 | --log_every=25 \ 95 | --batch_size=64 \ 96 | --eval_every=100 \ 97 | --val_set_size=100 \ 98 | --checkpoint_every=500 \ 99 | --disable_wandb 100 | 101 | # Test the learning rate scheduler 102 | uv run python scripts/train.py \ 103 | --robot_name=$ROBOT \ 104 | --learning_rate=0.01 \ 105 | --gamma=0.5 \ 106 | --step_lr_every=10 \ 107 | --disable_wandb 108 | """ 109 | 110 | 111 | if __name__ == "__main__": 112 | parser = argparse.ArgumentParser(prog="IKFlow training script") 113 | 114 | # Note: WandB saves artifacts by the run ID (i.e. '34c2gimi') not the run name ('dashing-forest-33'). This is 115 | # slightly annoying because you need to click on a run to get its ID. 116 | parser.add_argument("--wandb_run_id_to_load_checkpoint", type=str, help="Example: '34c2gimi'") 117 | parser.add_argument("--robot_name", type=str) 118 | 119 | # Model parameters 120 | parser.add_argument("--coupling_layer", type=str, default=DEFAULT_COUPLING_LAYER) 121 | parser.add_argument("--rnvp_clamp", type=float, default=DEFAULT_RNVP_CLAMP) 122 | parser.add_argument("--softflow_noise_scale", type=float, default=DEFAULT_SOFTFLOW_NOISE_SCALE) 123 | # NOTE: NEVER use 'bool' type with argparse. It will cause you pain and suffering. 124 | parser.add_argument("--softflow_enabled", type=str, default=DEFAULT_SOFTFLOW_ENABLED) 125 | parser.add_argument("--nb_nodes", type=int, default=DEFAULT_N_NODES) 126 | parser.add_argument("--dim_latent_space", type=int, default=DEFAULT_DIM_LATENT_SPACE) 127 | parser.add_argument("--coeff_fn_config", type=int, default=DEFAULT_COEFF_FN_CONFIG) 128 | parser.add_argument("--coeff_fn_internal_size", type=int, default=DEFAULT_COEFF_FN_INTERNAL_SIZE) 129 | parser.add_argument("--y_noise_scale", type=float, default=DEFAULT_Y_NOISE_SCALE) 130 | parser.add_argument("--zeros_noise_scale", type=float, default=DEFAULT_ZEROS_NOISE_SCALE) 131 | # See note above about pain and suffering. 132 | parser.add_argument("--sigmoid_on_output", type=str, default=DEFAULT_SIGMOID_ON_OUTPUT) 133 | 134 | # Training parameters 135 | parser.add_argument("--optimizer", type=str, default=DEFAULT_OPTIMIZER) 136 | parser.add_argument("--batch_size", type=int, default=DEFAULT_BATCH_SIZE) 137 | parser.add_argument("--gamma", type=float, default=DEFAULT_GAMMA) 138 | parser.add_argument("--learning_rate", type=float, default=DEFAULT_LR) 139 | parser.add_argument("--n_epochs", type=int, default=DEFAULT_N_EPOCHS) 140 | parser.add_argument("--step_lr_every", type=int, default=DEFAULT_STEP_LR_EVERY) 141 | parser.add_argument("--gradient_clip_val", type=float, default=DEFAULT_GRADIENT_CLIP_VAL) 142 | parser.add_argument("--lambd", type=float, default=1) 143 | parser.add_argument("--weight_decay", type=float, default=1.8e-05) 144 | 145 | # Logging options 146 | parser.add_argument("--eval_every", type=int, default=DEFAULT_EVAL_EVERY) 147 | parser.add_argument("--val_set_size", type=int, default=DEFAULT_VAL_SET_SIZE) 148 | parser.add_argument("--log_every", type=int, default=DEFAULT_LOG_EVERY) 149 | parser.add_argument("--checkpoint_every", type=int, default=DEFAULT_CHECKPOINT_EVERY) 150 | parser.add_argument("--dataset_tags", nargs="+", type=str, default=[DATASET_TAG_NON_SELF_COLLIDING]) 151 | parser.add_argument("--run_description", type=str) 152 | parser.add_argument("--disable_progress_bar", action="store_true") 153 | parser.add_argument("--disable_wandb", action="store_true") 154 | 155 | args = parser.parse_args() 156 | print("\nArgparse arguments:") 157 | for k, v in vars(args).items(): 158 | print(f" {k}={v}") 159 | print() 160 | 161 | if args.dataset_tags is None: 162 | args.dataset_tags = [] 163 | 164 | assert DATASET_TAG_NON_SELF_COLLIDING in args.dataset_tags, ( 165 | "The 'non-self-colliding' dataset should be specified (for now)" 166 | ) 167 | assert args.optimizer in ["ranger", "adadelta", "adamw"] 168 | assert 0 <= args.lambd and args.lambd <= 1 169 | 170 | # Load model 171 | robot = get_robot(args.robot_name) 172 | base_hparams = IkflowModelParameters() 173 | base_hparams.run_description = args.run_description 174 | base_hparams.coupling_layer = args.coupling_layer 175 | base_hparams.nb_nodes = args.nb_nodes 176 | base_hparams.dim_latent_space = args.dim_latent_space 177 | base_hparams.coeff_fn_config = args.coeff_fn_config 178 | base_hparams.coeff_fn_internal_size = args.coeff_fn_internal_size 179 | base_hparams.rnvp_clamp = args.rnvp_clamp 180 | base_hparams.softflow_noise_scale = args.softflow_noise_scale 181 | base_hparams.y_noise_scale = args.y_noise_scale 182 | base_hparams.zeros_noise_scale = args.zeros_noise_scale 183 | base_hparams.softflow_enabled = boolean_string(args.softflow_enabled) 184 | base_hparams.sigmoid_on_output = boolean_string(args.sigmoid_on_output) 185 | print() 186 | print(base_hparams) 187 | 188 | torch.autograd.set_detect_anomaly(False) 189 | data_module = IkfLitDataset(robot.name, args.batch_size, args.val_set_size, args.dataset_tags) 190 | 191 | # Setup wandb logging 192 | wandb_logger = None 193 | if not args.disable_wandb: 194 | wandb_entity, wandb_project = get_wandb_project() 195 | # Call `wandb.init` before creating a `WandbLogger` object so that runs have randomized names. Without this 196 | # call, the run names are all set to the project name. See this article for further information: https://lightrun.com/answers/lightning-ai-lightning-wandblogger--use-random-name-instead-of-project-as-default-name 197 | cfg = {"robot": args.robot_name} 198 | cfg.update(non_private_dict(args.__dict__)) 199 | data_module.add_dataset_hashes_to_cfg(cfg) 200 | 201 | wandb.init( 202 | entity=wandb_entity, 203 | project=wandb_project, 204 | notes=args.run_description, 205 | config=cfg, 206 | ) 207 | wandb_logger = WandbLogger(log_model="all", save_dir=config.WANDB_CACHE_DIR) 208 | run = wandb.run 209 | run.tags = run.tags + tuple(args.dataset_tags) 210 | 211 | ik_solver = IKFlowSolver(base_hparams, robot) 212 | model = IkfLitModel( 213 | ik_solver=ik_solver, 214 | base_hparams=base_hparams, 215 | learning_rate=args.learning_rate, 216 | checkpoint_every=args.checkpoint_every, 217 | log_every=args.log_every, 218 | gradient_clip=args.gradient_clip_val, 219 | lambd=args.lambd, 220 | gamma=args.gamma, 221 | step_lr_every=args.step_lr_every, 222 | weight_decay=args.weight_decay, 223 | optimizer_name=args.optimizer, 224 | sigmoid_on_output=boolean_string(args.sigmoid_on_output), 225 | ) 226 | 227 | # Checkpoint callback 228 | checkpoint_directory = get_checkpoint_dir(args.robot_name) 229 | if not args.disable_wandb: 230 | wandb.config.update({"checkpoint_directory": checkpoint_directory}) 231 | 232 | checkpoint_callback = ModelCheckpoint( 233 | dirpath=checkpoint_directory, 234 | every_n_train_steps=args.checkpoint_every, 235 | save_on_train_epoch_end=False, 236 | # Save the last 3 checkpoints. Checkpoint files are logged as wandb artifacts so it doesn't really matter anyway 237 | save_top_k=3, 238 | monitor="global_step", 239 | mode="max", 240 | filename="ikflow-checkpoint-{step}", 241 | ) 242 | 243 | # Train 244 | trainer = Trainer( 245 | logger=wandb_logger, 246 | callbacks=[checkpoint_callback], 247 | val_check_interval=args.eval_every, 248 | # gpus=1, # deprecated 249 | devices=[GPU_IDX], 250 | accelerator="gpu", 251 | log_every_n_steps=args.log_every, 252 | max_epochs=DEFAULT_MAX_EPOCHS, 253 | enable_progress_bar=False if (os.getenv("IS_SLURM") is not None) or args.disable_progress_bar else True, 254 | ) 255 | trainer.fit(model, data_module) 256 | -------------------------------------------------------------------------------- /scripts/train_from_checkpoint.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | from time import time 4 | 5 | from jrl.config import GPU_IDX 6 | from jrl.robots import get_robot 7 | from pytorch_lightning.loggers import WandbLogger 8 | from pytorch_lightning.callbacks import ModelCheckpoint 9 | from pytorch_lightning.trainer import Trainer 10 | from pytorch_lightning import seed_everything 11 | import wandb 12 | import torch 13 | 14 | from ikflow import config 15 | from ikflow.model import IkflowModelParameters 16 | from ikflow.ikflow_solver import IKFlowSolver 17 | from ikflow.training.training_utils import get_checkpoint_dir 18 | from ikflow.training.lt_model import IkfLitModel 19 | from ikflow.training.lt_data import IkfLitDataset 20 | from ikflow.utils import get_wandb_project 21 | 22 | assert GPU_IDX >= 0 23 | DEFAULT_MAX_EPOCHS = 5000 24 | SEED = 0 25 | seed_everything(SEED, workers=True) 26 | 27 | 28 | """ 29 | _____________ 30 | Example usage 31 | 32 | uv run python scripts/train_from_checkpoint.py --wandb_run_id=1vhgo90v 33 | """ 34 | 35 | if __name__ == "__main__": 36 | parser = argparse.ArgumentParser(prog="Resume training from a checkpoint") 37 | 38 | # Note: WandB saves artifacts by the run ID (i.e. '34c2gimi') not the run name ('dashing-forest-33'). This is 39 | # slightly annoying because you need to click on a run to get its ID. 40 | parser.add_argument("--wandb_run_id", type=str, help="The run ID of the wandb run to load. Example: '34c2gimi'") 41 | parser.add_argument("--disable_progress_bar", action="store_true") 42 | args = parser.parse_args() 43 | 44 | wandb_entity, wandb_project = get_wandb_project() 45 | wandb_run = wandb.init(entity=wandb_entity, project=wandb_project, id=args.wandb_run_id, resume="must") 46 | 47 | # Note: The checkpointing call back we use (`ModelCheckpoint`) saves the k most recent checkpoints and chooses 48 | # which to save by ranking with `global_step`. This means for whatever reason that the 'latest' alias isn't 49 | # available on the wandb side. They store 'best_k' instead which is what we want (I hope <*_*>) 50 | artifact = wandb_run.use_artifact(f"model-{args.wandb_run_id}:best_k") 51 | t0 = time() 52 | artifact_dir = artifact.download() 53 | print(f"Downloaded artifact '{artifact.name}' in {round(1000 * (time() - t0), 2)} ms") 54 | ckpt_filepath = os.path.join(artifact_dir, "model.ckpt") 55 | checkpoint = torch.load(ckpt_filepath, map_location=lambda storage, loc: storage) 56 | 57 | # Get the original runs WandB config. Example config: 58 | # { 59 | # ... 60 | # "robot_name":"panda_arm", 61 | # "val_set_size":25, 62 | # "optimizer_name":"ranger", 63 | # ... 64 | # } 65 | api = wandb.Api() 66 | prev_run = api.run(f"/{wandb_entity}/{wandb_project}/runs/{args.wandb_run_id}") 67 | run_cfg = prev_run.config 68 | robot_name = run_cfg["robot"] 69 | 70 | # Create training objects 71 | ikflow_hparams: IkflowModelParameters = checkpoint["hyper_parameters"]["base_hparams"] 72 | robot = get_robot(robot_name) 73 | # TODO(@jstm): Try wandb.restore() (see https://wandb.ai/cayush/pytorchlightning/reports/Use-PyTorch-Lightning-with-Weights-Biases--Vmlldzo2NjQ1Mw#:~:text=wandb.restore(%27EarlyStoppingADam%2D32%2D0.001.pth%27)) 74 | wandb_logger = WandbLogger(log_model="all", save_dir=config.WANDB_CACHE_DIR) 75 | 76 | ik_solver = IKFlowSolver(ikflow_hparams, robot) 77 | 78 | # Checkpoint callback 79 | if "checkpoint_directory" in run_cfg and os.path.isdir(run_cfg["checkpoint_directory"]): 80 | checkpoint_directory = run_cfg["checkpoint_directory"] 81 | else: 82 | checkpoint_directory = get_checkpoint_dir(robot_name) 83 | checkpoint_callback = ModelCheckpoint( 84 | dirpath=checkpoint_directory, 85 | every_n_train_steps=run_cfg["checkpoint_every"], 86 | save_on_train_epoch_end=False, 87 | # Save the last 3 checkpoints. Checkpoint files are logged as wandb artifacts so it doesn't really matter anyway 88 | save_top_k=3, 89 | monitor="global_step", 90 | mode="max", 91 | filename="ikflow-checkpoint-{step}", 92 | ) 93 | 94 | # Train 95 | model = IkfLitModel.load_from_checkpoint(ckpt_filepath, ik_solver=ik_solver) 96 | data_module = IkfLitDataset(robot_name, run_cfg["batch_size"], run_cfg["val_set_size"], run_cfg["dataset_tags"]) 97 | 98 | trainer = Trainer( 99 | # resume_from_checkpoint=ckpt_filepath, # 'Please pass `Trainer.fit(ckpt_path=)` directly instead' 100 | logger=wandb_logger, 101 | callbacks=[checkpoint_callback], 102 | val_check_interval=run_cfg["eval_every"], 103 | devices=[GPU_IDX], 104 | accelerator="gpu", 105 | log_every_n_steps=run_cfg["log_every"], 106 | max_epochs=DEFAULT_MAX_EPOCHS, 107 | enable_progress_bar=False if (os.getenv("IS_SLURM") is not None) or args.disable_progress_bar else True, 108 | ) 109 | trainer.fit(model, data_module, ckpt_path=ckpt_filepath) 110 | -------------------------------------------------------------------------------- /scripts/visualize.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | 4 | from ikflow import visualizations 5 | from ikflow.utils import set_seed 6 | from ikflow.model_loading import get_ik_solver 7 | 8 | set_seed() 9 | 10 | _VALID_DEMO_NAMES = ["oscillate_target", "oscillate_latent"] 11 | VIS_FN_ARGUMENTS = {"oscillate_target": {"nb_sols": 10, "fixed_latent": True}, "oscillate_latent": {}} 12 | 13 | """ Example usage. Note that `model_name` should match an entry in `model_descriptions.yaml` 14 | 15 | # Panda 16 | uv run python scripts/visualize.py --model_name=panda__full__lp191_5.25m --demo_name=oscillate_latent 17 | uv run python scripts/visualize.py --model_name=panda__full__lp191_5.25m --demo_name=oscillate_target 18 | 19 | # Rizon4 20 | uv run python scripts/visualize.py --model_name rizon4__snowy-brook-208__global_step=2.75M --demo_name=oscillate_latent 21 | uv run python scripts/visualize.py --model_name rizon4__snowy-brook-208__global_step=2.75M --demo_name=oscillate_target 22 | 23 | # Fetch 24 | uv run python scripts/visualize.py --model_name=fetch__large__ns183_9.75m --demo_name=oscillate_latent 25 | uv run python scripts/visualize.py --model_name=fetch__large__ns183_9.75m --demo_name=oscillate_target 26 | 27 | # FetchArm 28 | uv run python scripts/visualize.py --model_name=fetch_arm__large__mh186_9.25m --demo_name=oscillate_latent 29 | uv run python scripts/visualize.py --model_name=fetch_arm__large__mh186_9.25m --demo_name=oscillate_target 30 | """ 31 | 32 | 33 | if __name__ == "__main__": 34 | parser = argparse.ArgumentParser(prog="evaluate.py - evaluates IK models") 35 | parser.add_argument("--model_name", type=str, help="Name of the saved model to look for in model_descriptions.yaml") 36 | parser.add_argument("--visualization_program", type=str, help="One of [---TODO---]") 37 | parser.add_argument("--demo_name", type=str, help="One of ['oscillate_latent', 'oscillate_target']") 38 | args = parser.parse_args() 39 | 40 | assert args.demo_name in _VALID_DEMO_NAMES, f"Invalid demo name. Must be one of {_VALID_DEMO_NAMES}" 41 | 42 | # Build IKFlowSolver and set weights 43 | ik_solver, hyper_parameters = get_ik_solver(args.model_name) 44 | 45 | fn = getattr(visualizations, args.demo_name) 46 | fn(ik_solver, **VIS_FN_ARGUMENTS[args.demo_name]) 47 | -------------------------------------------------------------------------------- /scripts/visualize_robot.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from ikflow.visualizations import oscillate_joints 4 | from jrl.robots import get_robot 5 | 6 | """ Example usage 7 | 8 | uv run python scripts/visualize_robot.py --robot_name=panda 9 | uv run python scripts/visualize_robot.py --robot_name=fetch 10 | """ 11 | 12 | if __name__ == "__main__": 13 | parser = argparse.ArgumentParser(prog="evaluate.py - evaluates IK models") 14 | parser.add_argument("--robot_name", type=str, help="Example: 'panda', 'baxter', ...") 15 | args = parser.parse_args() 16 | robot = get_robot(args.robot_name) 17 | oscillate_joints(robot) 18 | -------------------------------------------------------------------------------- /test: -------------------------------------------------------------------------------- 1 | uv run python -m pytest tests/ -------------------------------------------------------------------------------- /tests/evaluation_utils_test.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | 3 | import torch 4 | import numpy as np 5 | from jrl.robots import Panda 6 | from jrl.utils import set_seed 7 | 8 | from ikflow.evaluation_utils import solution_pose_errors, calculate_joint_limits_exceeded 9 | 10 | set_seed() 11 | 12 | 13 | class EvaluationUtilsTest(unittest.TestCase): 14 | def test_solution_pose_errors(self): 15 | """Check that solution_pose_errors() matches the expected value""" 16 | robot = Panda() 17 | 18 | target_pose = torch.tensor([1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0], dtype=torch.float32, device="cpu") 19 | solutions = torch.zeros((1, 7), dtype=torch.float32, device="cpu") 20 | realized_pose = robot.forward_kinematics(solutions)[0] 21 | realized_pose_gt = torch.tensor( 22 | [0.088, 0.0, 0.926, 0.0, 0.92387953, 0.38268343, 0.0], dtype=torch.float32, device="cpu" 23 | ) 24 | np.testing.assert_allclose(realized_pose, realized_pose_gt, atol=1e-5) # sanity check values 25 | 26 | # np.sqrt((1 - 0.088 )**2 + (1 - 0.0 )**2 + (1 - 0.926 )**2 ) = 1.355440887681938 27 | l2_error_expected = 1.355440887681938 28 | angular_error_expected = 3.1415927 # From https://www.andre-gaschler.com/rotationconverter/ 29 | 30 | l2_error_returned, angular_error_returned = solution_pose_errors(robot, solutions, target_pose) 31 | self.assertAlmostEqual(l2_error_returned[0].item(), l2_error_expected) 32 | self.assertAlmostEqual(angular_error_returned[0].item(), angular_error_expected, delta=5e-4) 33 | 34 | def test_calculate_joint_limits_exceeded(self): 35 | """Test that calculate_joint_limits_exceeded() is working as expected""" 36 | 37 | configs = torch.tensor( 38 | [ 39 | [0, 0, 0], 40 | [0, 0, 0], 41 | [-2, 0, 0], 42 | [0, -1.999, 0], 43 | [0, 2.0001, 0], 44 | ] 45 | ) 46 | joint_limits = [ 47 | (-1, 1), 48 | (-2, 2), 49 | (-3, 3), 50 | ] 51 | expected = torch.tensor([False, False, True, False, True], dtype=torch.bool) 52 | returned = calculate_joint_limits_exceeded(configs, joint_limits) 53 | self.assertEqual(returned.dtype, torch.bool) 54 | self.assertEqual(returned.shape, (5,)) 55 | torch.testing.assert_close(returned, expected) 56 | 57 | 58 | if __name__ == "__main__": 59 | unittest.main() 60 | -------------------------------------------------------------------------------- /tests/ikflow_solver_test.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | from time import time 3 | 4 | import torch 5 | from jrl.robots import Panda 6 | from jrl.math_utils import geodesic_distance_between_quaternions 7 | from jrl.config import DEVICE 8 | 9 | from ikflow.model_loading import get_ik_solver 10 | from ikflow.utils import set_seed 11 | from ikflow.model import TINY_MODEL_PARAMS 12 | from ikflow.ikflow_solver import IKFlowSolver 13 | 14 | set_seed() 15 | 16 | 17 | def _assert_different(t1: torch.Tensor, t2: torch.Tensor): 18 | assert t1.ndim == 2 19 | assert t2.ndim == 2 20 | for i in range(t1.shape[0]): 21 | for j in range(t2.shape[0]): 22 | v_t1 = t1[i, j] 23 | n_matching = ((t2 - v_t1).abs() < 1e-8).sum().item() 24 | assert n_matching == 0 25 | 26 | 27 | class IkflowSolverTest(unittest.TestCase): 28 | # uv run python tests/ikflow_solver_test.py IkflowSolverTest.test___calculate_pose_error 29 | def test___calculate_pose_error(self): 30 | model_name = "panda__full__lp191_5.25m" 31 | n_solutions = 1000 32 | 33 | ikflow_solver, _ = get_ik_solver(model_name) 34 | # ikflow_solver._calculate_pose_error(qs, target_poses) 35 | 36 | print("Number of solutions, CPU, CUDA") 37 | for n_solutions in [1, 10, 100, 1000, 5000]: 38 | qs, target_poses = ikflow_solver.robot.sample_joint_angles_and_poses(n_solutions) 39 | qs_cpu = torch.tensor(qs, device="cpu", dtype=torch.float32) 40 | qs_cu = torch.tensor(qs, device="cuda:0", dtype=torch.float32) 41 | target_poses_cpu = torch.tensor(target_poses, device="cpu", dtype=torch.float32) 42 | target_poses_cu = torch.tensor(target_poses, device="cuda:0", dtype=torch.float32) 43 | 44 | t0 = time() 45 | ikflow_solver._calculate_pose_error(qs_cpu, target_poses_cpu) 46 | t_cpu = time() - t0 47 | 48 | t0 = time() 49 | ikflow_solver._calculate_pose_error(qs_cu, target_poses_cu) 50 | t_cuda = time() - t0 51 | 52 | print(n_solutions, "\t", round(1000 * t_cpu, 6), "\t", round(1000 * t_cuda, 6)) 53 | 54 | # pytest tests/ikflow_solver_test.py -k 'test_generate_exact_ik_solutions' 55 | # uv run python tests/ikflow_solver_test.py IkflowSolverTest.test_generate_exact_ik_solutions 56 | def test_generate_exact_ik_solutions(self): 57 | model_name = "panda__full__lp191_5.25m" 58 | POS_ERROR_THRESHOLD = 0.001 59 | ROT_ERROR_THRESHOLD = 0.01 60 | # device = "cpu" 61 | device = DEVICE 62 | n_solutions = 1000 63 | repeat_counts = (1, 3, 10) 64 | 65 | ikflow_solver, _ = get_ik_solver(model_name) 66 | robot = ikflow_solver.robot 67 | 68 | set_seed(3) 69 | torch.set_printoptions(linewidth=300, precision=7, sci_mode=False) 70 | _, target_poses = ikflow_solver.robot.sample_joint_angles_and_poses( 71 | n_solutions, only_non_self_colliding=True, tqdm_enabled=False 72 | ) 73 | target_poses = torch.tensor(target_poses, device=device, dtype=torch.float32) 74 | solutions, valid_solutions = ikflow_solver.generate_exact_ik_solutions( 75 | target_poses, 76 | pos_error_threshold=POS_ERROR_THRESHOLD, 77 | rot_error_threshold=ROT_ERROR_THRESHOLD, 78 | repeat_counts=repeat_counts, 79 | ) 80 | 81 | # evaluate solutions 82 | pose_realized = robot.forward_kinematics(solutions) 83 | torch.testing.assert_close(pose_realized[:, 0:3], target_poses[:, 0:3], atol=POS_ERROR_THRESHOLD, rtol=1e-2) 84 | rot_errors = geodesic_distance_between_quaternions(target_poses[:, 3:], pose_realized[:, 3:]) 85 | self.assertLess(rot_errors.max().item(), ROT_ERROR_THRESHOLD) 86 | torch.testing.assert_close(solutions, robot.clamp_to_joint_limits(solutions.clone())) 87 | assert valid_solutions.sum().item() == n_solutions 88 | 89 | def test_solve_multiple_poses(self): 90 | robot = Panda() 91 | ikflow_solver = IKFlowSolver(TINY_MODEL_PARAMS, robot) 92 | 93 | # Test 1: equal inputs should have equal outputs 94 | ys = torch.tensor([[0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0]], device=DEVICE, dtype=torch.float32) 95 | latent = torch.tensor( 96 | [ 97 | [0, 0, 0, 0, 0, 0, 0, 0, 0], 98 | [0, 0, 0, 0, 0, 0, 0, 0, 0], 99 | ], 100 | device=DEVICE, 101 | dtype=torch.float32, 102 | ) 103 | ikf_sols = ikflow_solver.generate_ik_solutions(ys, None, latent=latent, refine_solutions=False, allow_uninitialized=True) 104 | torch.testing.assert_close(ikf_sols[0], ikf_sols[1]) 105 | 106 | # Test 2: different inputs should have equal outputs 107 | ys = torch.tensor([[0, 0, 0, 0, 0, 0, 0], [1, 0, 0, 0, 0, 0, 0]], device=DEVICE, dtype=torch.float32) 108 | latent = torch.tensor( 109 | [ 110 | [0, 0, 0, 0, 0, 0, 0, 0, 0], 111 | [0, 0, 0, 0, 0, 0, 0, 0, 0], 112 | ], 113 | device=DEVICE, 114 | dtype=torch.float32, 115 | ) 116 | ikf_sols = ikflow_solver.generate_ik_solutions(ys, None, latent=latent, refine_solutions=False, allow_uninitialized=True) 117 | _assert_different(ikf_sols[0][None, :], ikf_sols[1][None, :]) 118 | 119 | 120 | if __name__ == "__main__": 121 | unittest.main() 122 | -------------------------------------------------------------------------------- /tests/lt_model_test.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | 3 | 4 | from ikflow.ikflow_solver import IKFlowSolver 5 | from jrl.robots import get_robot 6 | from ikflow.training.lt_model import IkfLitModel 7 | from ikflow.model import IkflowModelParameters 8 | 9 | import torch 10 | 11 | torch.manual_seed(0) 12 | 13 | ROBOT_MODEL = get_robot("panda") 14 | 15 | 16 | class LitModelTest(unittest.TestCase): 17 | def setUp(self) -> None: 18 | self.hps = IkflowModelParameters() 19 | self.ik_solver = IKFlowSolver(self.hps, ROBOT_MODEL) 20 | self.model = IkfLitModel( 21 | self.ik_solver, self.hps, learning_rate=1.0, checkpoint_every=-1, optimizer_name="adamw" 22 | ) 23 | 24 | def test_gradient_calculated(self): 25 | batch_size = 10 26 | device = "cuda:0" 27 | batch = ( 28 | torch.randn((batch_size, self.ik_solver.robot.ndof), device=device), 29 | torch.randn((batch_size, 7), device=device), 30 | ) 31 | fixed_latent = torch.randn((5, self.ik_solver.network_width), device=device) 32 | zero_response0 = self.ik_solver.generate_ik_solutions( 33 | torch.zeros(7, device=device, dtype=torch.float32), n=5, latent=fixed_latent, allow_uninitialized=True 34 | ) 35 | optimizer = torch.optim.Adadelta(self.model.nn_model.parameters(), lr=self.model.hparams.learning_rate) 36 | self.model.zero_grad() 37 | loss = self.model.training_step(batch, 0) 38 | loss.backward() 39 | optimizer.step() 40 | zero_response_f = self.ik_solver.generate_ik_solutions( 41 | torch.zeros(7, device=device, dtype=torch.float32), n=5, latent=fixed_latent, allow_uninitialized=True 42 | ) 43 | self.assertFalse(torch.allclose(zero_response0, zero_response_f)) 44 | 45 | def test_lit_model_has_params(self): 46 | n_parameters = len(list(self.model.parameters())) 47 | self.assertGreater(n_parameters, 0) 48 | 49 | 50 | if __name__ == "__main__": 51 | unittest.main() 52 | -------------------------------------------------------------------------------- /tests/model_loading_test.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | 3 | from ikflow.model_loading import model_filename 4 | 5 | import torch 6 | 7 | torch.manual_seed(0) 8 | 9 | 10 | class ModelLoadingTest(unittest.TestCase): 11 | def test_model_filename(self): 12 | url = "https://storage.googleapis.com/ikflow_models/atlas_desert-sweep-6.pkl" 13 | filename = model_filename(url) 14 | self.assertEqual(filename, "atlas_desert-sweep-6.pkl") 15 | 16 | 17 | if __name__ == "__main__": 18 | unittest.main() 19 | -------------------------------------------------------------------------------- /tests/model_test.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | 3 | import torch 4 | from jrl.robots import Panda 5 | import FrEIA.framework as Ff 6 | 7 | from ikflow.utils import set_seed, assert_joint_angle_tensor_in_joint_limits 8 | from ikflow.model import glow_cNF_model, IkflowModelParameters, get_pre_sigmoid_scaling_node 9 | from ikflow.ikflow_solver import IKFlowSolver 10 | 11 | set_seed() 12 | 13 | _PANDA = Panda() 14 | 15 | 16 | class ModelTest(unittest.TestCase): 17 | def setUp(self) -> None: 18 | # panda_joint_lims: (-> midpoint) (-> range) 19 | # (-2.8973, 2.8973) 0.0 5.7946 20 | # (-1.7628, 1.7628) 0.0 3.5256 21 | # (-2.8973, 2.8973) 0.0 5.7946 22 | # (-3.0718, -0.0698) -1.5708 3.0020000000000002 23 | # (-2.8973, 2.8973) 0.0 5.7946 24 | # (-0.0175, 3.7525) 1.8675 3.77 25 | # (-2.8973, 2.8973) 0.0 5.7946 26 | 27 | self.panda_theta_upper = torch.tensor( 28 | [[2.8973, 1.7628, 2.8973, -0.0698, 2.8973, 3.7525, 2.8973]], device="cuda" 29 | ) 30 | 31 | # sanity check: 32 | self.assertEqual(self.panda_theta_upper.shape, (1, 7)) 33 | for i, (_, upper) in enumerate(_PANDA.actuated_joints_limits): 34 | self.assertAlmostEqual(self.panda_theta_upper[0, i].item(), upper, places=5) 35 | 36 | self.panda_theta_lower = torch.tensor( 37 | [[-2.8973, -1.7628, -2.8973, -3.0718, -2.8973, -0.0175, -2.8973]], device="cuda" 38 | ) 39 | # sanity check self.panda_theta_lower: 40 | self.assertEqual(self.panda_theta_lower.shape, (1, 7)) 41 | for i, (lower, _) in enumerate(_PANDA.actuated_joints_limits): 42 | self.assertAlmostEqual(self.panda_theta_lower[0, i].item(), lower, places=5) 43 | 44 | self.panda_mid = torch.tensor([[0.0, 0.0, 0.0, -1.5708, 0.0, 1.8675, 0.0]], device="cuda") 45 | 46 | # ================================================================================================================== 47 | # Tests 48 | # 49 | 50 | def test_IkFlowFixedLinearTransform_forward(self): 51 | """Test that get_pre_sigmoid_scaling_node is returning a correctly configured IkFlowFixedLinearTransform node 52 | during the forward pass 53 | """ 54 | ndim_tot = 7 55 | input_node = Ff.InputNode(ndim_tot, name="input") 56 | nodes = [input_node] 57 | _, scaling_node = get_pre_sigmoid_scaling_node(ndim_tot, _PANDA, nodes) 58 | 59 | # Test 1: theta is the upper limit joint limit 60 | expected = 1.0 * torch.ones((1, 7), device="cuda") 61 | output, _ = scaling_node.forward([self.panda_theta_upper - 1e-8], rev=False) 62 | output = output[0] 63 | torch.testing.assert_close(output, expected) 64 | 65 | # Test 2: theta is the lower joint limit 66 | expected = torch.zeros((1, 7), device="cuda") 67 | output, _ = scaling_node.forward([self.panda_theta_lower + 1e-8], rev=False) 68 | output = output[0] 69 | torch.testing.assert_close(output, expected) 70 | 71 | # Test 3: theta is halfway between joint limits, output should be ~0.5 72 | theta = self.panda_mid 73 | expected = 0.5 * torch.ones((1, 7), device="cuda") 74 | output, _ = scaling_node.forward([theta], rev=False) 75 | output = output[0] 76 | torch.testing.assert_close(output, expected) 77 | 78 | def test_IkFlowFixedLinearTransform_reverse(self): 79 | """Test that get_pre_sigmoid_scaling_node is returning a correctly configured IkFlowFixedLinearTransform node 80 | during the reverse pass 81 | """ 82 | ndim_tot = 7 83 | input_node = Ff.InputNode(ndim_tot, name="input") 84 | nodes = [input_node] 85 | _, scaling_node = get_pre_sigmoid_scaling_node(ndim_tot, _PANDA, nodes) 86 | 87 | # Test 1: theta is the upper limit joint limit 88 | input = 1.0 * torch.ones((1, 7), device="cuda") 89 | expected = self.panda_theta_upper 90 | output, _ = scaling_node.forward([input], rev=True) 91 | output = output[0] 92 | torch.testing.assert_close(output, expected) 93 | 94 | # Test 2: theta is the lower joint limit 95 | input = torch.zeros((1, 7), device="cuda") 96 | expected = self.panda_theta_lower 97 | output, _ = scaling_node.forward([input], rev=True) 98 | output = output[0] 99 | torch.testing.assert_close(output, expected) 100 | 101 | # Test 3: theta is halfway between joint limits, output should be ~0.5 102 | input = 0.5 * torch.ones((1, 7), device="cuda") 103 | expected = self.panda_mid 104 | output, _ = scaling_node.forward([input], rev=True) 105 | output = output[0] 106 | torch.testing.assert_close(output, expected) 107 | 108 | def test_sigmoid_on_output(self): 109 | """Test that a working model is returned with sigmoid_on_output""" 110 | params = IkflowModelParameters() 111 | params.sigmoid_on_output = True 112 | model = glow_cNF_model(params, _PANDA, dim_cond=8, ndim_tot=7) 113 | conditional = torch.randn((5, 8), device="cuda") 114 | 115 | # Test 1: Reverse pass, gaussian noise 116 | latent = torch.randn((5, 7), device="cuda") 117 | output_rev, _ = model(latent, c=conditional, rev=True) 118 | assert_joint_angle_tensor_in_joint_limits(_PANDA.actuated_joints_limits, output_rev, "reverse", eps=1e-5) 119 | 120 | # Test 2: Reverse pass, really large gaussian noise 121 | latent = 1e8 * torch.randn((5, 7), device="cuda") 122 | output_rev, _ = model(latent, c=conditional, rev=True) 123 | assert_joint_angle_tensor_in_joint_limits(_PANDA.actuated_joints_limits, output_rev, "reverse", eps=1e-5) 124 | 125 | def test_glow_cNF_model(self): 126 | """Smoke test - checks that glow_cNF_model() returns""" 127 | 128 | # Setup model 129 | params = IkflowModelParameters() 130 | params.softflow_enabled = False 131 | params.dim_latent_space = 9 132 | model = glow_cNF_model(params, _PANDA, dim_cond=7, ndim_tot=9) 133 | 134 | # Run model 135 | conditional = torch.randn((5, 7), device="cuda") 136 | latent = torch.randn((5, 9), device="cuda") 137 | output_rev, _ = model(latent, c=conditional, rev=True) 138 | loss = torch.sum(output_rev) 139 | loss.backward() 140 | 141 | def test_IKFlowSolver_smoketest(self): 142 | """Smoke test - checks that IKFlowSolver is happy""" 143 | 144 | # Test 1: works with softflow disabled 145 | params = IkflowModelParameters() 146 | params.softflow_enabled = False 147 | solver = IKFlowSolver(params, _PANDA) 148 | 149 | conditional = torch.randn((5, 7), device="cuda") 150 | latent = torch.randn((5, 9), device="cuda") 151 | output_rev, _ = solver.nn_model(latent, c=conditional, rev=True) 152 | loss = torch.sum(output_rev) 153 | loss.backward() 154 | 155 | # Test 1: works with softflow 156 | params = IkflowModelParameters() 157 | params.softflow_enabled = True 158 | solver = IKFlowSolver(params, _PANDA) 159 | 160 | conditional = torch.randn((5, 8), device="cuda") 161 | latent = torch.randn((5, 9), device="cuda") 162 | output_rev, _ = solver.nn_model(latent, c=conditional, rev=True) 163 | loss = torch.sum(output_rev) 164 | loss.backward() 165 | 166 | def test_glow_cNF_model_w_softflow(self): 167 | """Smoke test - checks that glow_cNF_model() returns""" 168 | # Setup model 169 | params = IkflowModelParameters() 170 | params.softflow_enabled = True 171 | params.dim_latent_space = 9 172 | model = glow_cNF_model(params, _PANDA, dim_cond=8, ndim_tot=9) 173 | 174 | # Run model 175 | conditional = torch.randn((5, 8), device="cuda") 176 | latent = torch.randn((5, 9), device="cuda") 177 | output_rev, _ = model(latent, c=conditional, rev=True) 178 | loss = torch.sum(output_rev) 179 | loss.backward() 180 | 181 | # TODO: Implement this test 182 | # def test_accuracy(self): 183 | # """Download a trained model, check that the measured accuracy matches the expected accuracy""" 184 | 185 | 186 | if __name__ == "__main__": 187 | unittest.main() 188 | --------------------------------------------------------------------------------