diff --git a/README.md b/README.md index 60a2cdcf..67503b04 100644 --- a/README.md +++ b/README.md @@ -1,192 +1,69 @@ -# adam +# πŸ€– adam [![adam](https://github.com/ami-iit/adam/actions/workflows/tests.yml/badge.svg?branch=main)](https://github.com/ami-iit/adam/actions/workflows/tests.yml) [![](https://img.shields.io/badge/License-BSD--3--Clause-blue.svg)](https://github.com/ami-iit/adam/blob/main/LICENSE) -**Automatic Differentiation for rigid-body-dynamics AlgorithMs** +**Automatic Differentiation for rigid-body-dynamics Algorithms** -**adam** implements a collection of algorithms for calculating rigid-body dynamics for **floating-base** robots, in _mixed_ and _body fixed representations_ (see [Traversaro's A Unified View of the Equations of Motion used for Control Design of Humanoid Robots](https://www.researchgate.net/publication/312200239_A_Unified_View_of_the_Equations_of_Motion_used_for_Control_Design_of_Humanoid_Robots)) using: +**adam** computes rigid-body dynamics for floating-base robots. Built on Featherstone's algorithms and available across multiple backends: -- [Jax](https://github.com/google/jax) -- [CasADi](https://web.casadi.org/) -- [PyTorch](https://github.com/pytorch/pytorch) -- [NumPy](https://numpy.org/) +- πŸ”₯ **JAX** – compile, vectorize, and differentiate with XLA +- 🎯 **CasADi** – symbolic computation for optimization and control +- πŸ”¦ **PyTorch** – GPU acceleration and batched operations +- 🐍 **NumPy** – simple numerical evaluation -**adam** employs the **automatic differentiation** capabilities of these frameworks to compute, if needed, gradients, Jacobian, Hessians of rigid-body dynamics quantities. This approach enables the design of optimal control and reinforcement learning strategies in robotics. +All backends share the same interface and produce numerically consistent results, letting you pick the tool that fits your use case. -**adam** is based on Roy Featherstone's Rigid Body Dynamics Algorithms. +## πŸ“¦ Installation -### Table of contents - -- [adam](#adam) - - [Table of contents](#table-of-contents) - - [🐍 Dependencies](#-dependencies) - - [πŸ’Ύ Installation](#-installation) - - [🐍 Installation with pip](#-installation-with-pip) - - [πŸ“¦ Installation with conda](#-installation-with-conda) - - [Installation from conda-forge package](#installation-from-conda-forge-package) - - [πŸ”¨ Installation from repo](#-installation-from-repo) - - [πŸš€ Usage](#-usage) - - [Jax interface](#jax-interface) - - [CasADi interface](#casadi-interface) - - [PyTorch interface](#pytorch-interface) - - [PyTorch Batched interface](#pytorch-batched-interface) - - [MuJoCo interface](#mujoco-interface) - - [Inverse Kinematics](#inverse-kinematics) - - [πŸ¦Έβ€β™‚οΈ Contributing](#️-contributing) - -## 🐍 Dependencies - -- [`python3`](https://wiki.python.org/moin/BeginnersGuide) - -Other requisites are: - -- [`urdfdom-py`](https://pypi.org/project/urdfdom-py/) Python package, that exposes the `urdf_parser_py` Python module -- `jax` -- `casadi` -- `pytorch` -- `numpy` -- `array-api-compat` - -They will be installed in the installation step! - -## πŸ’Ύ Installation - -The installation can be done either using the Python provided by apt (on Debian-based distros) or via conda (on Linux and macOS). - -### 🐍 Installation with pip - -Install `python3`, if not installed, for example on **Ubuntu**: +### With pip ```bash -sudo apt install python3 python3-pip python3-venv -``` - -Create a [virtual environment](https://docs.python.org/3/library/venv.html#venv-def), if you prefer. For example: - -```bash -pip install virtualenv -python3 -m venv your_virtual_env -source your_virtual_env/bin/activate -``` - -Inside the virtual environment, install the library from pip: - -- Install **Jax** interface: - - ```bash - pip install adam-robotics[jax] - ``` - -- Install **CasADi** interface: - - ```bash - pip install adam-robotics[casadi] - ``` - -- Install **PyTorch** interface: +# JAX backend +pip install adam-robotics[jax] - ```bash - pip install adam-robotics[pytorch] - ``` +# CasADi backend +pip install adam-robotics[casadi] -- Install **ALL** interfaces: +# PyTorch backend +pip install adam-robotics[pytorch] - ```bash - pip install adam-robotics[all] - ``` - -If you want the last version: - -```bash -pip install adam-robotics[selected-interface]@git+https://github.com/ami-iit/adam +# All backends +pip install adam-robotics[all] ``` -or clone the repo and install: +### With conda ```bash -git clone https://github.com/ami-iit/adam.git -cd adam -pip install .[selected-interface] -``` - -### πŸ“¦ Installation with conda - -#### Installation from conda-forge package - -- Install **CasADi** interface: - - ```bash - conda create -n adamenv -c conda-forge adam-robotics-casadi - ``` - -- Install **Jax** interface (warning: not available on Windows): - - ```bash - conda create -n adamenv -c conda-forge adam-robotics-jax - ``` - -- Install **PyTorch** interface (warning: not available on Windows): - - ```bash - conda create -n adamenv -c conda-forge adam-robotics-pytorch - ``` - -- Install **ALL** interfaces (warning: not available on Windows): - - ```bash - conda create -n adamenv -c conda-forge adam-robotics-all - ``` - -> [!NOTE] -> Check also the conda JAX installation guide [here](https://jax.readthedocs.io/en/latest/installation.html#conda-community-supported) - -### πŸ”¨ Installation from repo - -Install in a conda environment the required dependencies: +# CasADi backend +conda create -n adamenv -c conda-forge adam-robotics-casadi -- **Jax** interface dependencies: +# JAX backend (Linux/macOS only) +conda create -n adamenv -c conda-forge adam-robotics-jax - ```bash - conda create -n adamenv -c conda-forge jax numpy lxml prettytable matplotlib urdfdom-py array-api-compat - ``` +# PyTorch backend (Linux/macOS only) +conda create -n adamenv -c conda-forge adam-robotics-pytorch -- **CasADi** interface dependencies: - - ```bash - conda create -n adamenv -c conda-forge casadi numpy lxml prettytable matplotlib urdfdom-py array-api-compat - ``` - -- **PyTorch** interface dependencies: - - ```bash - conda create -n adamenv -c conda-forge pytorch numpy lxml prettytable matplotlib urdfdom-py array-api-compat - ``` - -- **ALL** interfaces dependencies: - - ```bash - conda create -n adamenv -c conda-forge jax casadi pytorch numpy lxml prettytable matplotlib urdfdom-py array-api-compat - ``` +# All backends (Linux/macOS only) +conda create -n adamenv -c conda-forge adam-robotics-all +``` -Activate the environment, clone the repo and install the library: +### From source ```bash -conda activate adamenv git clone https://github.com/ami-iit/adam.git cd adam -pip install --no-deps . +pip install .[jax] # or [casadi], [pytorch], [all] ``` -## πŸš€ Usage +## πŸš€ Quick Start -The following are small snippets of the use of **adam**. More examples are arriving! -Have also a look at the `tests` folder. +Load a robot model and compute dynamics quantities: -### Jax interface +### JAX > [!NOTE] -> Check also the Jax installation guide [here](https://jax.readthedocs.io/en/latest/installation.html#) +> Check the [JAX installation guide](https://jax.readthedocs.io/en/latest/installation.html) ```python import adam @@ -208,17 +85,21 @@ joints_name_list = [ ] kinDyn = KinDynComputations(model_path, joints_name_list) -# choose the representation, if you want to use the body fixed representation -kinDyn.set_frame_velocity_representation(adam.Representations.BODY_FIXED_REPRESENTATION) -# or, if you want to use the mixed representation (that is the default) +# Set velocity representation (3 options available): +# 1. MIXED_REPRESENTATION (default) - time derivative of base origin position (expressed in world frame) + world-frame angular velocity kinDyn.set_frame_velocity_representation(adam.Representations.MIXED_REPRESENTATION) +# 2. BODY_FIXED_REPRESENTATION - both linear & angular velocity in body frame +# kinDyn.set_frame_velocity_representation(adam.Representations.BODY_FIXED_REPRESENTATION) +# 3. INERTIAL_FIXED_REPRESENTATION - world-frame linear & angular velocity +# kinDyn.set_frame_velocity_representation(adam.Representations.INERTIAL_FIXED_REPRESENTATION) + w_H_b = np.eye(4) joints = np.ones(len(joints_name_list)) M = kinDyn.mass_matrix(w_H_b, joints) print(M) w_H_f = kinDyn.forward_kinematics('frame_name', w_H_b, joints) -# IMPORTANT! The Jax Interface function execution can be slow! We suggest to jit them. +# JAX functions can also be jitted! # For example: def frame_forward_kinematics(w_H_b, joints): @@ -228,22 +109,16 @@ def frame_forward_kinematics(w_H_b, joints): jitted_frame_fk = jit(frame_forward_kinematics) w_H_f = jitted_frame_fk(w_H_b, joints) -# In the same way, the functions can be also vmapped -vmapped_frame_fk = vmap(frame_forward_kinematics, in_axes=(0, 0)) -# which can be also jitted -jitted_vmapped_frame_fk = jit(vmapped_frame_fk) -# and called on a batch of data +# JAX natively supports batching joints_batch = jnp.tile(joints, (1024, 1)) w_H_b_batch = jnp.tile(w_H_b, (1024, 1, 1)) -w_H_f_batch = jitted_vmapped_frame_fk(w_H_b_batch, joints_batch) - - +w_H_f_batch = kinDyn.forward_kinematics('frame_name', w_H_b_batch, joints_batch) ``` > [!NOTE] > The first call of the jitted function can be slow, since JAX needs to compile the function. Then it will be faster! -### CasADi interface +### CasADi ```python import casadi as cs @@ -264,10 +139,14 @@ joints_name_list = [ ] kinDyn = KinDynComputations(model_path, joints_name_list) -# choose the representation you want to use the body fixed representation -kinDyn.set_frame_velocity_representation(adam.Representations.BODY_FIXED_REPRESENTATION) -# or, if you want to use the mixed representation (that is the default) +# Set velocity representation (3 options available): +# 1. MIXED_REPRESENTATION (default) - time derivative of position + world-frame angular velocity kinDyn.set_frame_velocity_representation(adam.Representations.MIXED_REPRESENTATION) +# 2. BODY_FIXED_REPRESENTATION - both linear & angular velocity in body frame +# kinDyn.set_frame_velocity_representation(adam.Representations.BODY_FIXED_REPRESENTATION) +# 3. INERTIAL_FIXED_REPRESENTATION - world-frame linear & angular velocity +# kinDyn.set_frame_velocity_representation(adam.Representations.INERTIAL_FIXED_REPRESENTATION) + w_H_b = np.eye(4) joints = np.ones(len(joints_name_list)) M = kinDyn.mass_matrix_fun() @@ -284,10 +163,9 @@ w_H_b = cs.MX.eye(4) joints = cs.MX.sym('joints', len(joints_name_list)) M = kinDyn.mass_matrix_fun() print(M(w_H_b, joints)) - ``` -### PyTorch interface +### PyTorch ```python import adam @@ -317,12 +195,17 @@ M = kinDyn.mass_matrix(w_H_b, joints) print(M) ``` -### PyTorch Batched interface +### PyTorch Batched + +Use `pytorch.KinDynComputations` to process also multiple configurations. + +> [!NOTE] +> There is a class `pytorch.KinDynComputationsBatch` that has the functionality of `pytorch.KinDynComputations`. It exists to avoid API changes in existing code. New users should prefer `pytorch.KinDynComputations` for both single and batched computations. ```python import adam -from adam.pytorch import KinDynComputationsBatch +from adam.pytorch import KinDynComputations import icub_models # if you want to icub-models @@ -336,7 +219,7 @@ joints_name_list = [ 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch', 'r_ankle_roll' ] -kinDyn = KinDynComputationsBatch(model_path, joints_name_list) +kinDyn = KinDynComputations(model_path, joints_name_list) # choose the representation you want to use the body fixed representation kinDyn.set_frame_velocity_representation(adam.Representations.BODY_FIXED_REPRESENTATION) # or, if you want to use the mixed representation (that is the default) @@ -352,7 +235,7 @@ M = kinDyn.mass_matrix(w_H_b_batch, joints_batch) w_H_f = kinDyn.forward_kinematics('frame_name', w_H_b_batch, joints_batch) ``` -### MuJoCo interface +### MuJoCo adam supports loading models directly from [MuJoCo](https://mujoco.org/) `MjModel` objects. This is useful when working with MuJoCo simulations or models from [robot_descriptions](https://github.com/robot-descriptions/robot_descriptions.py). @@ -407,12 +290,6 @@ J = kinDyn.jacobian('frame_name', w_H_b, joints) ### Inverse Kinematics -adam provides an interface for solving inverse kinematics problems using CasADi. The solver supports - -- position, orientation, and full pose constraints -- frame-to-frame constraints (ball, fixed) -- optional joint limit constraints - ```python import casadi as cs import numpy as np @@ -420,36 +297,61 @@ import adam from adam.casadi import KinDynComputations from adam.casadi.inverse_kinematics import InverseKinematics, TargetType -# Load your robot model -import icub_models -model_path = icub_models.get_model_file("iCubGazeboV2_5") -# The joint list -joints_name_list = ... +# Load model +model_path = ... +joints_name_list = [...] + # Create IK solver -ik = InverseKinematics(model_path, joints) -# Add a pose target on a frame (e.g., the left sole) +ik = InverseKinematics(model_path, joints_name_list) ik.add_target("l_sole", target_type=TargetType.POSE, as_soft_constraint=True, weight=1.0) -ik.add_ball_constraint(frame_1, frame_2, as_soft_constraint=True) -# Update the target to a desired pose +# Update target and solve desired_position = np.array([0.3, 0.2, 1.0]) desired_orientation = np.eye(3) ik.update_target("l_sole", (desired_position, desired_orientation)) - -# Solve ik.solve() -# Retrieve solution +# Get solution w_H_b_sol, q_sol = ik.get_solution() print("Base pose:\n", w_H_b_sol) print("Joint values:\n", q_sol) ``` -## πŸ¦Έβ€β™‚οΈ Contributing +## πŸ“š Features + +- **Kinematics**: Forward kinematics, Jacobians (frame and base) +- **Dynamics**: Mass matrix, Coriolis/centrifugal forces and gravity, Articulated body algorithm +- **Centroidal**: Centroidal momentum matrix and derivatives +- **Differentiation**: Get gradients, Jacobians, and Hessians automatically +- **Symbolic**: Build computation graphs with CasADi for optimization +- **Batched**: Process multiple configurations in parallel with PyTorch + +## πŸ“– Documentation + +See the [full documentation](https://adam-robotics.readthedocs.io/en/latest/) for detailed API reference, more examples, and theory. + +## πŸ§ͺ Testing + +Run tests to verify installation: + +```bash +pip install .[test] # Install test dependencies +pytest tests/ +``` + +See `tests/` folder for comprehensive examples across all backends. -**adam** is an open-source project. Contributions are very welcome! +## 🀝 Contributing -Open an issue with your feature request or if you spot a bug. Then, you can also proceed with a Pull-requests! :rocket: +Found a bug or have a feature idea? Open an [issue](https://github.com/ami-iit/adam/issues) or submit a [pull request](https://github.com/ami-iit/adam/pulls)! πŸš€ > [!WARNING] -> REPOSITORY UNDER DEVELOPMENT! We cannot guarantee stable API +> This is a project under active development. API may change. + +## πŸ“„ License + +BSD 3-Clause License – see [LICENSE](LICENSE) file. + +## πŸ™ Acknowledgments + +Built on [Roy Featherstone's Rigid Body Dynamics Algorithms](https://link.springer.com/book/10.1007/978-1-4899-7560-7) and references like [Traversaro's A Unified View of the Equations of Motion](https://traversaro.github.io/traversaro-phd-thesis/traversaro-phd-thesis.pdf). diff --git a/docs/guides/backend_selection.rst b/docs/guides/backend_selection.rst new file mode 100644 index 00000000..a812f020 --- /dev/null +++ b/docs/guides/backend_selection.rst @@ -0,0 +1,140 @@ +Choosing a Backend +=================== + +adam supports multiple backends, each optimized for different use cases. This guide helps you choose the right one for your application. + +NumPy +^^^^^ + +**Best for:** Quick prototyping, validation, simplicity + +- Direct numerical computation +- No compilation or setup overhead +- Easiest to debug and understand +- Good for initial model verification + +**When to use:** + +- Validating robot models +- Quick experiments + +**Example:** + +.. code-block:: python + + from adam.numpy import KinDynComputations + kinDyn = KinDynComputations(model_path, joints_list) + M = kinDyn.mass_matrix(w_H_b, joints) + + +JAX +^^^ + +**Best for:** Research, optimization, vectorized computations + +- JIT compilation for speed +- Automatic differentiation (grad, jacobian, hessian) +- Native batch support + +**When to use:** + +- Gradient-based optimization +- Computing derivatives and Jacobians +- Processing batches of configurations +- Control design with auto-differentiation + +**Gotchas:** + +- First call with jit is slow (compilation) +- Can't use Python branching on traced values + +**Example:** + +.. code-block:: python + + from adam.jax import KinDynComputations + from jax import jit, grad + + kinDyn = KinDynComputations(model_path, joints_list) + + @jit + def compute_and_grad(w_H_b, joints): + M = kinDyn.mass_matrix(w_H_b, joints) + grad_fn = grad(lambda q: kinDyn.mass_matrix(w_H_b, q).sum()) + return M, grad_fn(joints) + + +CasADi +^^^^^^ + +**Best for:** Optimal control, trajectory optimization, symbolic formulation + +- Symbolic computation for optimization +- Function generation for code generation +- Both numeric and symbolic evaluation +- Constraint handling + +**When to use:** + +- Nonlinear model predictive control (NMPC) +- Trajectory optimization +- Building optimization problems +- When you need symbolic expressions +- Code generation for embedded systems + +**Example:** + +.. code-block:: python + + import casadi as cs + from adam.casadi import KinDynComputations + + kinDyn = KinDynComputations(model_path, joints_list) + + # Numeric evaluation + M_fun = kinDyn.mass_matrix_fun() + M = M_fun(w_H_b, joints) + + # Symbolic computation + w_H_b_sym = cs.SX.sym('H', 4, 4) + joints_sym = cs.SX.sym('q', n_dof) + M_sym = M_fun(w_H_b_sym, joints_sym) + + +PyTorch +^^^^^^^ + +**Best for:** Machine learning, GPU acceleration, batched computation + +- GPU usage out of the box +- Native batch support +- Integration with neural networks +- Automatic differentiation + +**When to use:** + +- Learning-based control +- Large-scale batch processing +- GPU-accelerated robotics +- Integration with ML frameworks + +**Gotchas:** + +- GPU memory can be limiting +- Requires careful dtype/device handling + +**Example:** + +.. code-block:: python + + import torch + from adam.pytorch import KinDynComputations + + kinDyn = KinDynComputations(model_path, joints_list) + + # Batched computation + batch_size = 1024 + w_H_b_batch = torch.randn(batch_size, 4, 4) + joints_batch = torch.randn(batch_size, n_dof) + + M_batch = kinDyn.mass_matrix(w_H_b_batch, joints_batch) # Shape: (batch, 6+n, 6+n) diff --git a/docs/guides/concepts.rst b/docs/guides/concepts.rst new file mode 100644 index 00000000..eb33bc01 --- /dev/null +++ b/docs/guides/concepts.rst @@ -0,0 +1,220 @@ +Core Concepts +============== + +This guide explains the fundamental concepts in adam. + +Floating-Base Representation +----------------------------- + +adam models robots as a **floating base** articulated system. The state consists of: + +- **Base Transform** ``w_H_b``: The 4Γ—4 homogeneous transformation matrix from world frame to base frame +- **Joint Positions** ``joints``: Scalar values for each actuated joint +- **Base Velocity** ``base_vel``: A 6D vector representing the linear and angular velocity of the base +- **Joint Velocities** ``joints_vel``: Scalar velocities for each actuated joint + +Robot dynamics +---------------- + +The robot dynamics are described by the standard equations of motion for floating-base systems: + +.. math:: + + M(q) \dot{\nu} + C(q, \nu) + g(q) = \tau + J^T(q) f_{ext} + +where: + +- :math:`M(q)` is the mass matrix +- :math:`C(q, \nu)` are Coriolis and centrifugal effects +- :math:`g(q)` is the gravity term +- :math:`\tau` are the generalized forces/torques +- :math:`J(q)` is the Jacobian matrix of the frame on which external wrenches act +- :math:`f_{ext}` are external wrenches (forces/torques) +- :math:`q` = [w_H_b, joints] is the configuration +- :math:`\nu` = [base_vel, joints_vel] is the configuration velocity +- :math:`\dot{\nu}` is the configuration acceleration + + +Velocity Representations +------------------------ + +There are three 6D velocity representations for the floating base (check also useful `iDynTree theory `_). +Let :math:`A` denote the world (inertial) frame and :math:`B` denote the base frame. The 6D velocity is stacked as :math:`\begin{bmatrix} v \\ \omega \end{bmatrix}` with linear velocity first. + +**Mixed (Default)** + +Expressed as ``Representations.MIXED_REPRESENTATION`` (default in adam and iDynTree): + +.. math:: + + {}^{A[B]}\mathrm{v}_{A,B} = \begin{bmatrix} {}^{A}\dot{o}_{B} \\ {}^{A}\omega_{A,B} \end{bmatrix} = \begin{bmatrix} {}^{A}\dot{o}_{B} \\ (\dot{{}^{A}R_{B}} {}^{A}R_{B}^{\top})^{\vee} \end{bmatrix} + +Linear velocity is the time derivative of the base origin position (expressed in world frame :math:`A`), and angular velocity is expressed in the **world frame**. +This hybrid representation is commonly used in humanoid robotics and is the default in adam. + + +**Left-Trivialized (Body-Fixed)** + +Expressed as ``Representations.BODY_FIXED_REPRESENTATION``: + +.. math:: + + {}^{B}\mathrm{v}_{A,B} = \begin{bmatrix} {}^{B}v_{B} \\ {}^{B}\omega_{A,B} \end{bmatrix} = \begin{bmatrix} {}^{B}R_{A} {}^{A}\dot{o}_{B} \\ ({}^{A}R_{B}^{\top} \dot{{}^{B}R_{A}})^{\vee} \end{bmatrix} + +Both linear and angular velocities are expressed in the **base frame** coordinates. +This is the "body-fixed" frame representation. + +**Right-Trivialized (Inertial-Fixed)** + +Expressed as ``Representations.INERTIAL_FIXED_REPRESENTATION``: + +.. math:: + + {}^{A}\mathrm{v}_{A,B} = \begin{bmatrix} {}^{A}v_{B} \\ {}^{A}\omega_{A,B} \end{bmatrix} = \begin{bmatrix} {}^{A}\dot{o}_{B} - \dot{{}^{A}R_{B}} {}^{A}R_{B}^{\top} {}^{A}o_{B} \\ (\dot{{}^{A}R_{B}} {}^{A}R_{B}^{\top})^{\vee} \end{bmatrix} + +Linear velocity is expressed in the **world frame**, angular velocity in the **world frame**. +This is the "inertial-fixed" frame representation. + + +**Setting the Representation** + +Always set this early when creating a ``KinDynComputations`` object: + +.. code-block:: python + + kinDyn.set_frame_velocity_representation(adam.Representations.MIXED_REPRESENTATION) + # or + kinDyn.set_frame_velocity_representation(adam.Representations.BODY_FIXED_REPRESENTATION) + # or + kinDyn.set_frame_velocity_representation(adam.Representations.INERTIAL_FIXED_REPRESENTATION) + +The representation affects: + +- **Jacobian computations** – Jacobians transform between representations via adjoint matrices +- **Inertia matrix structure** – Mass matrix transforms according to representation +- **Centroidal momentum calculations** – Momentum definitions differ by representation +- **Any velocity-dependent quantities** – Coriolis forces, bias forces, etc. + +See the `iDynTree theory documentation `_ for mathematical details on 6D velocity representations. + +State Format +^^^^^^^^^^^^ + +Velocities are always stacked as: + +.. code-block:: python + + [base_linear_vel, base_angular_vel, joint_velocities] + # Shape: (3,) + (3,) + (n_dof,) = (6 + n_dof,) + + +Key Matrices and Quantities +---------------------------- + +**Mass Matrix** ``M`` + +Shape: ``(6 + n_dof, 6 + n_dof)`` + +The inertia matrix relating generalized forces to generalized accelerations: + +.. math:: + + M(q) \dot{\nu} = \tau - h(q, \nu) - g(q) + +where: + +- :math:`M` is symmetric and positive definite +- Describes how joint positions affect inertia +- Used in inverse dynamics, control design + +**Centroidal Momentum Matrix** ``Ag`` + +Shape: ``(6, 6 + n_dof)`` + +Relates generalized velocities to centroidal momentum: + +.. math:: + + \begin{bmatrix} L \\ h \end{bmatrix} = A_g(q) \nu + +where: + +- ``L``: Linear momentum +- ``h``: Angular momentum about CoM + +**Jacobian** ``J`` + +Shape: ``(6, 6 + n_dof)`` + +Relates end-effector velocity to generalized velocities: + +.. math:: + + v_{ee} = J(q) \nu + +Available as: + +- ``jacobian(frame_name, ...)`` – frame Jacobian +- ``CoM_jacobian(...)`` – CoM Jacobian +- ``jacobian_dot(...)`` – time derivative + +**Forward Kinematics** ``w_H_frame`` + +Shape: ``(4, 4)`` + +The transformation from world to any frame: + +.. code-block:: python + + w_H_frame = kinDyn.forward_kinematics('frame_name', w_H_b, joints) + + +Dynamics Algorithms +------------------- + +adam uses **Featherstone's Recursive Algorithms** under the hood, which compute all quantities efficiently in O(n) time where n is the number of joints. + +**Articulated Body Algorithm (ABA)** + +Computes accelerations given forces: + +.. math:: + + \ddot{q} = \text{ABA}(q, \nu, \tau, f_{ext}) + +Used in forward dynamics. + + +**Recursive Newton-Euler Algorithm (RNEA)** + +Computes Coriolis, centrifugal, and gravity effects: + +.. math:: + + h(q, \nu) = C(q,\nu) + g(q) = \text{RNEA}(q, \nu) + +**Composite Rigid Body Algorithm (CRBA)** + +Computes the mass matrix and centroidal momentum matrix: + +.. math:: + + M(q), A_g(q) = \text{CRBA}(q) + +External Wrenches +----------------- + +When computing system acceleration with external forces (e.g., contact forces), pass them as a dictionary: + +.. code-block:: python + + external_wrenches = { + 'l_sole': np.array([fx, fy, fz, mx, my, mz]), # Contact force at left foot + 'r_sole': np.array([fx, fy, fz, mx, my, mz]), # Contact force at right foot + } + + acc = kinDyn.aba( + w_H_b, joints, base_vel, joints_vel, tau, + external_wrenches=external_wrenches + ) + diff --git a/docs/guides/index.rst b/docs/guides/index.rst new file mode 100644 index 00000000..26cab91e --- /dev/null +++ b/docs/guides/index.rst @@ -0,0 +1,12 @@ +Guides +====== + +Detailed guides for understanding and using adam effectively. + +.. toctree:: + :maxdepth: 2 + + concepts + backend_selection + mujoco + troubleshooting diff --git a/docs/guides/mujoco.rst b/docs/guides/mujoco.rst new file mode 100644 index 00000000..7c6b032b --- /dev/null +++ b/docs/guides/mujoco.rst @@ -0,0 +1,185 @@ +MuJoCo Integration +================== + +adam supports loading robot models directly from `MuJoCo `_ ``MjModel`` objects. This enables seamless integration with MuJoCo simulations and models from `robot_descriptions `_. + +Loading MuJoCo Models +--------------------- + +Use the ``from_mujoco_model()`` class method to create a ``KinDynComputations`` instance: + +.. code-block:: python + + import mujoco + import numpy as np + from adam import Representations + from adam.numpy import KinDynComputations + + # Load a MuJoCo model (e.g., from robot_descriptions) + from robot_descriptions.loaders.mujoco import load_robot_description + mj_model = load_robot_description("g1_mj_description") + + # Create KinDynComputations directly from MuJoCo model + kinDyn = KinDynComputations.from_mujoco_model(mj_model) + + # Set velocity representation + kinDyn.set_frame_velocity_representation(Representations.MIXED_REPRESENTATION) + + # Set gravity to match MuJoCo settings + kinDyn.g = np.concatenate([mj_model.opt.gravity, np.zeros(3)]) + +Joint Extraction Options +------------------------- + +By default, ``from_mujoco_model()`` extracts joint names from the MuJoCo model. You can customize this behavior: + +**Using Joint Names (Default)** + +.. code-block:: python + + kinDyn = KinDynComputations.from_mujoco_model(mj_model) + # Uses joint names from mj_model + +**Using Actuator Names** + +.. code-block:: python + + kinDyn = KinDynComputations.from_mujoco_model( + mj_model, + use_mujoco_actuators=True + ) + # Uses actuator names instead of joint names + +This is useful when your MuJoCo model has actuators defined and you want to work with actuator coordinates. + +Working with MuJoCo State +-------------------------- + +To compute dynamics quantities, you need to extract the robot state from MuJoCo and format it for adam: + +.. code-block:: python + + # Create MuJoCo data and set state + mj_data = mujoco.MjData(mj_model) + mj_data.qpos[:] = your_qpos # Your configuration + mj_data.qvel[:] = your_qvel # Your velocities + mujoco.mj_forward(mj_model, mj_data) + + # Extract base transform from MuJoCo state (for floating-base robots) + from scipy.spatial.transform import Rotation as R + + base_rot = R.from_quat(mj_data.qpos[3:7], scalar_first=True).as_matrix() + base_pos = mj_data.qpos[0:3] + w_H_b = np.eye(4) + w_H_b[:3, :3] = base_rot + w_H_b[:3, 3] = base_pos + + # Joint positions (excluding free joint) + # Ensure the serialization between MuJoCo and adam is the same + joints = mj_data.qpos[7:] + + # Compute dynamics quantities + M = kinDyn.mass_matrix(w_H_b, joints) + com_pos = kinDyn.CoM_position(w_H_b, joints) + J = kinDyn.jacobian('frame_name', w_H_b, joints) + +Velocity Representation Differences +------------------------------------ + +.. warning:: + + MuJoCo uses a different velocity representation for the floating object. + See MuJoCo documentation for details: `MuJoCo Floating objects `_. + +**MuJoCo Free Joint Velocity** + +The free joint velocity in MuJoCo is structured as: + +.. math:: + + \begin{bmatrix} {}^{I}\dot{p}_{B} \\ {}^{B}\omega_{I, B} \end{bmatrix} + +where the linear velocity is the **time derivative of the world frame position** and the angular velocity is in the **body frame**. + +**adam Mixed Representation** + +adam's default mixed representation (``MIXED_REPRESENTATION``) is: + +.. math:: + + \begin{bmatrix} {}^{I}\dot{p}_{B} \\ {}^{I}\omega_{I, B} \end{bmatrix} + +where linear velocity is the **time derivative of the world frame position** and angular velocity is in the **world frame**. + +**Transformation** + +To convert MuJoCo velocities to adam's mixed representation: + +.. code-block:: python + + # Transform angular velocity from body to world frame + R = ... # Rotation matrix + adam_base_vel = mujoco_base_vel.copy() + adam_base_vel[3:6] = R @ mujoco_base_vel[3:6] + + +Backend Support +--------------- + +All adam backends support MuJoCo model loading. For example: + +.. code-block:: python + + from adam.numpy import KinDynComputations + from adam.jax import KinDynComputations + from adam.pytorch import KinDynComputations + from adam.casadi import KinDynComputations + ... + kinDyn = KinDynComputations.from_mujoco_model(mj_model) + + +Example: Neural IK with MuJoCo Model +------------------------------------- + +See the `neural_ik.py example `_ for a complete demonstration of using adam to train a neural network for inverse kinematics. + +Common Pitfalls +--------------- + +**Joint Ordering** + +Ensure that the joint ordering in adam matches your expectations. Use: + +.. code-block:: python + + print(kinDyn.rbdalgos.model.joint_names) + +to verify the joint names and their order extracted from the MuJoCo model. + +**Free Joint Handling** + +MuJoCo's free joint (for floating-base robots) is automatically detected and excluded from the actuated joints list. The base transform and base velocity must be provided separately to adam methods. + +**Gravity Vector** + +Remember to set the gravity vector consistently between MuJoCo and adam: + +.. code-block:: python + + # MuJoCo gravity is 3D + kinDyn.g = np.concatenate([mj_model.opt.gravity, np.zeros(3)]) + # adam expects 6D: [linear_acceleration, angular_acceleration] + + +Testing with MuJoCo Models +--------------------------- + +The test suite includes MuJoCo-specific tests in ``tests/test_mujoco.py``. +These tests verify that quantities computed in adam and in mujoco correspond, given the velocity representation differences (see above). +See Also +-------- + +- :doc:`concepts` - Core concepts including velocity representations +- :doc:`backend_selection` - Choosing the right backend for your use case +- `MuJoCo Documentation `_ +- `robot_descriptions `_ - Repository of robot models diff --git a/docs/guides/troubleshooting.rst b/docs/guides/troubleshooting.rst new file mode 100644 index 00000000..e48f7afc --- /dev/null +++ b/docs/guides/troubleshooting.rst @@ -0,0 +1,179 @@ +Troubleshooting & FAQ +==================== + +Common Issues and Solutions +--------------------------- + +JAX: Slow First Execution +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +**Problem:** The first call to a JAX function is very slow. + +**Cause:** JAX needs to compile the function with XLA. + +**Solution:** Use ``jit`` to compile once: + +.. code-block:: python + + from jax import jit + + @jit + def compute_mass_matrix(w_H_b, joints): + return kinDyn.mass_matrix(w_H_b, joints) + + M = compute_mass_matrix(w_H_b, joints) # Slow (compilation) + M = compute_mass_matrix(w_H_b, joints) # Fast (cached) + +**See also:** :doc:`backend_selection` + + +JAX: Array Value in Traced Context +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +**Problem:** ``ValueError: array value being traced cannot be used in branching`` + +**Cause:** You're using Python conditionals on JAX arrays inside jitted code. + +**Solution:** Use JAX control flow: + +.. code-block:: python + + # ❌ Wrong + @jit + def bad_function(x): + if x > 0: # Can't do this! + return x + return -x + + # βœ… Correct + @jit + def good_function(x): + return jnp.where(x > 0, x, -x) + + +JAX: Loss of Precision +^^^^^^^^^^^^^^^^^^^^^^ + +**Problem:** Results don't match NumPy, differences larger than expected. + +**Cause:** JAX defaults to 32-bit floats. + +**Solution:** Enable 64-bit precision: + +.. code-block:: python + + from jax import config + config.update("jax_enable_x64", True) + + +PyTorch: Device or Dtype Mismatch +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +**Problem:** ``RuntimeError: Expected all tensors to be on the same device`` + +**Cause:** Mixing CPU and GPU tensors, or different dtypes. + +**Solution:** Ensure consistent device and dtype: + +.. code-block:: python + + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + w_H_b = torch.tensor(w_H_b, device=device, dtype=torch.float64) + joints = torch.tensor(joints, device=device, dtype=torch.float64) + + kinDyn = KinDynComputations(model_path, joints_list) + M = kinDyn.mass_matrix(w_H_b, joints) + + +PyTorch Batch: Wrong Shapes +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +**Problem:** Batched operations fail with shape errors. + +**Cause:** The unified ``KinDynComputations`` API expects batch dimension first for batched operations. + +**Solution:** Shape inputs as ``(batch_size, ...)``: + +.. code-block:: python + + batch_size = 32 + + # ❌ Wrong: (4, 4, batch) or (batch_size, 4) and (4) + w_H_b = np.random.randn(4, 4, batch_size) + joints = np.random.randn(batch_size) + + # βœ… Correct: (batch, 4, 4) and (batch, n_dof) + w_H_b = torch.tensor(np.tile(np.eye(4), (batch_size, 1, 1)), dtype=torch.float64) + joints = torch.randn(batch_size, n_dof, dtype=torch.float64) + + M = kinDyn.mass_matrix(w_H_b, joints) # Shape: (batch, 6+n, 6+n) + + +CasADi: Numeric vs Symbolic +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +**Problem:** Unsure when to use ``_fun()`` methods vs direct methods. + +**Solution:** + +- Use direct methods for immediate numeric evaluation: + + .. code-block:: python + + M = kinDyn.mass_matrix(w_H_b, joints) + +- Use ``_fun()`` for reusable compiled functions or symbolic computation: + + .. code-block:: python + + M_fun = kinDyn.mass_matrix_fun() + M_numeric = M_fun(w_H_b_numeric, joints_numeric) + M_symbolic = M_fun(w_H_b_symbolic, joints_symbolic) + + + +Numerical Issues +---------------- + +Results Don't Match Simulation +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +**Problem:** adam results differ from physics simulator (Gazebo, MuJoCo, etc). + +**Likely causes:** + +1. **Different gravity direction** – Set explicitly: + + .. code-block:: python + + kinDyn.g = np.array([0, 0, -9.81, 0, 0, 0]) + +2. **Different velocity representation** – Ensure consistency: + + .. code-block:: python + + kinDyn.set_frame_velocity_representation(adam.Representations.MIXED_REPRESENTATION) + +3. **Joint ordering mismatch** – Verify ``joints_name_list`` order matches simulator + +4. **Model differences** – Ensure same URDF or model definition + + +Accuracy Tolerance +^^^^^^^^^^^^^^^^^^ + +**Question:** How accurate are the computations? + +**Answer:** adam is validated against `iDynTree `_ with tolerance ``1e-5`` to ``1e-4``. + +For your use case: + + +Getting Help +------------ + +1. **Check the tests** – Browse `tests/ `_ for working examples +2. **Read some useful theory** – See idyntree `theory.md `_ +3. **Open an issue** – Report bugs on `GitHub `_ +4. **Review examples** – Look at `examples/ `_ diff --git a/docs/index.rst b/docs/index.rst index e7c84e0f..85d7e5bf 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -4,47 +4,87 @@ contain the root `toctree` directive. -adam ----- +adam - Rigid-Body Dynamics for Floating-Base Robots +======================================================== **Automatic Differentiation for rigid-body-dynamics AlgorithMs** -**adam** implements a collection of algorithms for calculating rigid-body dynamics for **floating-base** robots, in mixed and body fixed representations using: +**adam** computes rigid-body dynamics for floating-base robots. Choose from multiple backends depending on your use case: -.. create rst list with links to the following libraries +- πŸ”₯ `JAX `_ – JIT compilation, batched operations, and differentiation with XLA +- 🎯 `CasADi `_ – Symbolic computation for optimization and control +- πŸ”¦ `PyTorch `_ – GPU acceleration, batched operations and differentiation +- 🐍 `NumPy `_ – Simple numerical evaluation -- `Jax `_ -- `CasADi `_ -- `PyTorch `_ -- `NumPy `_ +All backends share the same API and produce numerically consistent results, letting you pick the tool that fits your workflow. +**Model Loading** + - URDF files – standard robot description format (see :doc:`quickstart/index`) + - MuJoCo models – direct integration with ``MjModel`` objects (see :doc:`guides/mujoco`) -**adam** employs the automatic differentiation capabilities of these frameworks to compute, if needed, gradients, Jacobian, Hessians of rigid-body dynamics quantities. This approach enables the design of optimal control and reinforcement learning strategies in robotics. +Core Features +------------- +**Kinematics & Geometry** + - Forward kinematics for any frame + - Jacobians for any frame + - Jacobian time derivatives -**adam** is based on **Roy Featherstone's Rigid Body Dynamics Algorithms**. +**Dynamics** + - Mass matrix computed via Composite Rigid Body Algorithm (CRBA) + - Bias forces (Coriolis and centrifugal forces + gravity term) computed via Recursive Newton-Euler Algorithm (RNEA) + - Articulated Body Algorithm (ABA) -Examples --------- +**Centroidal Dynamics** + - Centroidal momentum matrix via the Composite Rigid Body Algorithm + - Center of mass position and Jacobian -Have a look at the examples `folder in the repository `_! +**Automatic Differentiation** + - Gradients with JAX and PyTorch + - Symbolic computation with CasADi +**Advanced Features** + - Parametric models for shape/inertia optimization + - Inverse kinematics (CasADi) + - MuJoCo integration + - Batch processing (PyTorch and JAX) + +Philosophy +---------- + +Built on **Roy Featherstone's Rigid Body Dynamics Algorithms**, adam provides a composable interface across multiple backends. Consistency is guaranteed through extensive testing against `iDynTree `_. + +Resources +--------- + +- **Examples**: See the `examples folder `_ for notebooks and scripts +- **Tests**: The `tests folder `_ contains comprehensive usage patterns License ------- -`BSD-3-Clause `_ +BSD 3-Clause License – `view license `_ + .. toctree:: :maxdepth: 2 - :caption: Getting started: + :caption: Getting Started: installation quickstart/index .. toctree:: :maxdepth: 2 - :caption: API: + :caption: Guides: + + guides/concepts + guides/backend_selection + guides/mujoco + guides/troubleshooting + +.. toctree:: + :maxdepth: 2 + :caption: API Reference: modules/index diff --git a/docs/installation.rst b/docs/installation.rst index 48fbf469..eb7ed763 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -1,34 +1,56 @@ Installation ============ -This is the installation guide for the project. +adam requires Python 3.10 or later. Install it using your preferred package manager. -Prerequisites -------------- +Quickstart +---------- -adam requires python 3.7 or later. +Install with your preferred backend(s): +.. code-block:: bash -🐍 Conda installation ---------------------- + # Single backend + pip install adam-robotics[jax] # JAX backend + pip install adam-robotics[casadi] # CasADi backend + pip install adam-robotics[pytorch] # PyTorch backend + + # All backends + pip install adam-robotics[all] -We suggest to install adam using `conda `_: +Or with conda: .. code-block:: bash - conda install adam-robotics -c conda-forge + conda install -c conda-forge adam-robotics-casadi + conda install -c conda-forge adam-robotics-jax + conda install -c conda-forge adam-robotics-pytorch + +Which Backend? +-------------- + +See the :doc:`guides/backend_selection` guide to choose the right backend for your use case. -πŸ“¦ Pip installation +Development Install -------------------- -You can also install adam using `pip `_: +Clone the repository and install in editable mode: .. code-block:: bash - pip install adam-robotics + git clone https://github.com/ami-iit/adam.git + cd adam + pip install -e .[all,test] # Install all backends + test dependencies + +This enables development mode with all backends and testing tools. + +GPU Support (JAX & PyTorch) +--------------------------- + +**JAX with GPU** -If you want to install all the dependencies install ``adam-robotics[all]``. +Follow the `JAX GPU installation guide `_ for GPU support. -.. note:: +**PyTorch with GPU** - If the GPU support for ``JAX`` is needed, follow the instructions in the `Jax documentation `_. +PyTorch is automatically GPU-compatible. Check also `PyTorch installation guide `_. diff --git a/docs/modules/jax.rst b/docs/modules/jax.rst index 2aeb700a..37229480 100644 --- a/docs/modules/jax.rst +++ b/docs/modules/jax.rst @@ -9,7 +9,7 @@ This module provides the Jax implementation of the Rigid Body Dynamics algorithm .. tip:: - The functions in this module can be also ``jax.vmap``-ed to run on batches of inputs. + The functions in this module natively support batching via broadcasting, or can be explicitly vectorized with ``jax.vmap`` for advanced patterns. .. note:: diff --git a/docs/modules/pytorch.rst b/docs/modules/pytorch.rst index 4870c48f..0b48b374 100644 --- a/docs/modules/pytorch.rst +++ b/docs/modules/pytorch.rst @@ -3,10 +3,6 @@ PyTorch interface This module provides the PyTorch implementation of the Rigid Body Dynamics algorithms. -.. tip:: - - For the batched version of the algorithms, see :mod:`adam.pytorch.computation_batch`. - ---------------- .. currentmodule:: adam.pytorch diff --git a/docs/modules/pytorch_batched.rst b/docs/modules/pytorch_batched.rst deleted file mode 100644 index 73ac8e06..00000000 --- a/docs/modules/pytorch_batched.rst +++ /dev/null @@ -1,13 +0,0 @@ -PyTorch Batched interface -========================= - -This module implements the batched version of the Rigid Body Dynamics algorithms using PyTorch. - ----------------- - -.. currentmodule:: adam.pytorch.computation_batch - -.. automodule:: adam.pytorch.computation_batch - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/quickstart/casadi.rst b/docs/quickstart/casadi.rst index 1f089a98..8e83d291 100644 --- a/docs/quickstart/casadi.rst +++ b/docs/quickstart/casadi.rst @@ -1,18 +1,28 @@ -CasADi usage -============ +CasADi Backend +============== -The following example shows how to call an instance of the ``adam.casadi.KinDynComputations`` class and use it to compute the mass matrix and forward dynamics of a floating-base robot. +CasADi excels at symbolic computation and optimization formulation. Use it for optimal control, trajectory optimization, and code generation. + +Key Features +------------ + +- **Symbolic Computation** – Build optimization problems symbolically +- **Both Function Types** – Direct and ``_fun()`` variants +- **Code Generation** – Export to C/Python for embedded systems + +Basic Usage +----------- .. code-block:: python + import numpy as np + import casadi as cs import adam from adam.casadi import KinDynComputations import icub_models - import numpy as np - - # if you want to icub-models https://github.com/robotology/icub-models to retrieve the urdf + + # Load model model_path = icub_models.get_model_file("iCubGazeboV2_5") - # The joint list joints_name_list = [ 'torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch', 'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch', @@ -20,28 +30,198 @@ The following example shows how to call an instance of the ``adam.casadi.KinDynC 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll', 'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch', 'r_ankle_roll' ] - + kinDyn = KinDynComputations(model_path, joints_name_list) - # choose the representation you want to use the body fixed representation - kinDyn.set_frame_velocity_representation(adam.Representations.BODY_FIXED_REPRESENTATION) - # or, if you want to use the mixed representation (that is the default) kinDyn.set_frame_velocity_representation(adam.Representations.MIXED_REPRESENTATION) + + # Define state w_H_b = np.eye(4) - joints = np.ones(len(joints_name_list)) - M = kinDyn.mass_matrix_fun() - print(M(w_H_b, joints)) + joints = np.ones(len(joints_name_list)) * 0.1 + + # Numeric evaluation + M = kinDyn.mass_matrix(w_H_b, joints) + print(f"Mass matrix:\n{M}") + + +Function Variants +----------------- + +CasADi provides two ways to call each function: + +**Direct method (for immediate evaluation):** + +.. code-block:: python + + M = kinDyn.mass_matrix(w_H_b, joints) + J = kinDyn.jacobian('l_sole', w_H_b, joints) + +**Function factory (for reusable compiled functions):** + +.. code-block:: python + + M_fun = kinDyn.mass_matrix_fun() + J_fun = kinDyn.jacobian_fun('l_sole') + + # Evaluate multiple times (efficient) + M1 = M_fun(w_H_b, joints) + M2 = M_fun(w_H_b, joints + 0.1) + + +Symbolic Computation +-------------------- + +Build optimization problems symbolically: + +.. code-block:: python + + import casadi as cs + + # Create symbolic variables + w_H_b_sym = cs.SX.sym('H', 4, 4) + joints_sym = cs.SX.sym('q', len(joints_name_list)) + + # Symbolic mass matrix + M_fun = kinDyn.mass_matrix_fun() + M_sym = M_fun(w_H_b_sym, joints_sym) + + # Build symbolic expressions + M_det = cs.det(M_sym) + M_trace = cs.trace(M_sym) + + print(f"Determinant (symbolic):\n{M_det}") + + +Optimization Problem +-------------------- + +Setup a trajectory optimization problem: + +.. code-block:: python + + import casadi as cs + + # Variables + q_sym = cs.SX.sym('q', len(joints_name_list)) + q_dot_sym = cs.SX.sym('q_dot', len(joints_name_list)) + + # Objective: minimize energy + M_fun = kinDyn.mass_matrix_fun() + M = M_fun(np.eye(4), q_sym) + + kinetic_energy = 0.5 * cs.mtimes(q_dot_sym.T, cs.mtimes(M[:6, :6], q_dot_sym)) + + # Create optimization problem + opti = cs.Opti() + opti.minimize(kinetic_energy) + + # Add constraints... + # Solve... + - # If you want to use the symbolic version +Inverse Kinematics +------------------ + +adam provides an IK solver for CasADi: + +.. code-block:: python + + from adam.casadi.inverse_kinematics import InverseKinematics, TargetType + + # Create IK solver + ik = InverseKinematics(model_path, joints_name_list) + + # Add target + ik.add_target('l_sole', target_type=TargetType.POSE, as_soft_constraint=True) + + # Set desired pose + desired_pos = np.array([0.3, 0.2, 1.0]) + desired_rot = np.eye(3) + ik.update_target('l_sole', (desired_pos, desired_rot)) + + # Solve + ik.solve() + w_H_b_sol, q_sol = ik.get_solution() + + +Tips and Tricks +--------------- + +**Use DM for numeric results:** + +.. code-block:: python + + M_fun = kinDyn.mass_matrix_fun() + M_result = cs.DM(M_fun(w_H_b, joints)) + print(f"Numeric result: {float(M_result[0, 0])}") + +**Mix symbolic and numeric:** + +.. code-block:: python + + # Symbolic Jacobian + J_sym = J_fun(w_H_b_sym, joints_sym) + + # Create callable function + J_callable = cs.Function('J', [w_H_b_sym, joints_sym], [J_sym]) + + # Evaluate numerically + J_numeric = J_callable(w_H_b, joints) + +**Generate C code:** + +.. code-block:: python + + M_fun = kinDyn.mass_matrix_fun() + M_fun.generate('mass_matrix') # Generates mass_matrix.c + + +Loading from MuJoCo +------------------- + +Load models from MuJoCo for symbolic optimization: + +.. code-block:: python + + import mujoco + from robot_descriptions.loaders.mujoco import load_robot_description + from adam.casadi import KinDynComputations + import casadi as cs + + # Load MuJoCo model + mj_model = load_robot_description("g1_mj_description") + + # Create KinDynComputations from MuJoCo model + kinDyn = KinDynComputations.from_mujoco_model(mj_model) + kinDyn.set_frame_velocity_representation(adam.Representations.MIXED_REPRESENTATION) + + # Use in optimization + opti = cs.Opti() + joints = opti.variable(kinDyn.NDoF) + w_H_b = cs.SX.eye(4) - joints = cs.SX.sym('joints', len(joints_name_list)) - M = kinDyn.mass_matrix_fun() - print(M(w_H_b, joints)) - - # This is usable also with casadi.MX - w_H_b = cs.MX.eye(4) - joints = cs.MX.sym('joints', len(joints_name_list)) - M = kinDyn.mass_matrix_fun() - print(M(w_H_b, joints)) - - w_H_f = kinDyn.forward_kinematics_fun() - print(w_H_f('frame_name', w_H_b, joints)) + M_fun = kinDyn.mass_matrix_fun() + M = M_fun(w_H_b, joints) + + # Add constraints using dynamics + opti.minimize(cs.trace(M)) + opti.subject_to(cs.fabs(joints) <= 1.0) + opti.solver('ipopt') + sol = opti.solve() + +See :doc:`../guides/mujoco` for more details on MuJoCo integration. + +When to Use CasADi +------------------ + +βœ… **Good for:** +- Optimal control (MPC, trajectory optimization) +- Symbolic formulation +- Code generation for embedded systems +- Building constraint expressions +- Trajectory planning + +❌ **Not ideal for:** +- Simple numeric evaluation (NumPy is easier) +- Batching (not supported) + +**See also:** :doc:`../guides/backend_selection` diff --git a/docs/quickstart/index.rst b/docs/quickstart/index.rst index c0195c2e..6ac8946d 100644 --- a/docs/quickstart/index.rst +++ b/docs/quickstart/index.rst @@ -1,14 +1,18 @@ -Quickstart -=========== +Quickstart Examples +=================== -This section contains some snippets of code that show how to use the library. +This section contains practical examples for each backend. Choose based on your use case: -.. _quickstart: +- **NumPy** – Simple numerical computation, validation +- **JAX** – Automatic differentiation, GPU, native batching +- **CasADi** – Symbolic computation, optimization formulation, automatic differentiation +- **PyTorch** – GPU acceleration, native batching, ML integration, automatic differentiation .. toctree:: :maxdepth: 2 - casadi + numpy jax + casadi pytorch - pytorch_batched + diff --git a/docs/quickstart/jax.rst b/docs/quickstart/jax.rst index 01b84c74..ecf8eeaa 100644 --- a/docs/quickstart/jax.rst +++ b/docs/quickstart/jax.rst @@ -1,26 +1,30 @@ -Jax usage -========= +JAX Backend +=========== -The following example shows how to call an instance of the ``adam.jax.KinDynComputations`` class and use it to compute the mass matrix and forward dynamics of a floating-base robot. +The JAX backend excels at automatic differentiation and batching. Use it for gradient-based optimization and research prototypes. -.. tip:: - We suggest to ``jax.jit`` the functions as it will make them run faster! +Key Features +------------ -.. note:: - When the functions are ``jax.jit``-ed, the first time you run them, it will take a bit longer to execute as they are being compiled by Jax. +- **JIT Compilation** – First call is slow, subsequent calls are fast +- **Automatic Differentiation** – Compute gradients, Jacobians, Hessians +- **Native Batching** – Process batches of configurations +- **GPU Support** – Runs on GPU with proper JAX installation + +Basic Usage +----------- .. code-block:: python + import numpy as np + import jax.numpy as jnp import adam from adam.jax import KinDynComputations + from jax import jit, grad import icub_models - import numpy as np - import jax.numpy as jnp - from jax import jit, vmap - - # if you want to icub-models https://github.com/robotology/icub-models to retrieve the urdf + + # Load model model_path = icub_models.get_model_file("iCubGazeboV2_5") - # The joint list joints_name_list = [ 'torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch', 'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch', @@ -28,33 +32,213 @@ The following example shows how to call an instance of the ``adam.jax.KinDynComp 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll', 'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch', 'r_ankle_roll' ] - + kinDyn = KinDynComputations(model_path, joints_name_list) - # choose the representation, if you want to use the body fixed representation - kinDyn.set_frame_velocity_representation(adam.Representations.BODY_FIXED_REPRESENTATION) - # or, if you want to use the mixed representation (that is the default) kinDyn.set_frame_velocity_representation(adam.Representations.MIXED_REPRESENTATION) - w_H_b = np.eye(4) - joints = np.ones(len(joints_name_list)) + + # Define state + w_H_b = jnp.eye(4) + joints = jnp.ones(len(joints_name_list)) * 0.1 + + # Compute (slow, no compilation) M = kinDyn.mass_matrix(w_H_b, joints) - print(M) - w_H_f = kinDyn.forward_kinematics('frame_name', w_H_b, joints) - - # IMPORTANT! The Jax Interface function execution can be slow! We suggest to jit them. - # For example: - - def frame_forward_kinematics(w_H_b, joints): - # This is needed since str is not a valid JAX type - return kinDyn.forward_kinematics('frame_name', w_H_b, joints) - - jitted_frame_fk = jit(frame_forward_kinematics) - w_H_f = jitted_frame_fk(w_H_b, joints) - - # In the same way, the functions can be also vmapped - vmapped_frame_fk = vmap(frame_forward_kinematics, in_axes=(0, 0)) - # which can be also jitted - jitted_vmapped_frame_fk = jit(vmapped_frame_fk) - # and called on a batch of data - joints_batch = jnp.tile(joints, (1024, 1)) - w_H_b_batch = jnp.tile(w_H_b, (1024, 1, 1)) - w_H_f_batch = jitted_vmapped_frame_fk(w_H_b_batch, joints_batch) + print(f"Mass matrix shape: {M.shape}") + + +JIT Compilation +--------------- + +Wrap your functions with ``@jit`` for speed: + +.. code-block:: python + + from jax import jit + + @jit + def compute(w_H_b, joints): + M = kinDyn.mass_matrix(w_H_b, joints) + J = kinDyn.jacobian('l_sole', w_H_b, joints) + return M, J + + # First call: slow (compilation) + M, J = compute(w_H_b, joints) + + # Subsequent calls: fast (cached) + M, J = compute(w_H_b, joints) + +.. warning:: + + Frame names must remain as strings (not traced). Wrap them in a closure: + + .. code-block:: python + + # βœ… Correct + def make_fk_fn(frame_name): + @jit + def fk(w_H_b, joints): + return kinDyn.forward_kinematics(frame_name, w_H_b, joints) + return fk + + fk_l_sole = make_fk_fn('l_sole') + +Automatic Differentiation +-------------------------- + +Compute gradients easily: + +.. code-block:: python + + from jax import grad + + # Gradient of mass matrix trace w.r.t. joint positions + def mass_matrix_trace(w_H_b, joints): + M = kinDyn.mass_matrix(w_H_b, joints) + return jnp.trace(M) + + grad_fn = grad(mass_matrix_trace, argnums=1) # Gradient w.r.t. joints + grad_joints = grad_fn(w_H_b, joints) + print(f"Gradient shape: {grad_joints.shape}") + +**Higher-order derivatives:** + +.. code-block:: python + + from jax import grad, hessian + + hess_fn = hessian(mass_matrix_trace, argnums=1) + hess_joints = hess_fn(w_H_b, joints) + print(f"Hessian shape: {hess_joints.shape}") + + +Native Batching +---------------- + +JAX automatically broadcasts batched operations: + +.. code-block:: python + + # Batch size 1024 + batch_size = 1024 + w_H_b_batch = jnp.tile(jnp.eye(4), (batch_size, 1, 1)) # Shape: (1024, 4, 4) + joints_batch = jnp.tile(joints, (batch_size, 1)) # Shape: (1024, n_dof) + + # Just pass batched tensors - JAX handles batching automatically + M_batch = kinDyn.mass_matrix(w_H_b_batch, joints_batch) # Shape: (1024, 6+n, 6+n) + J_batch = kinDyn.jacobian('l_sole', w_H_b_batch, joints_batch) # Shape: (1024, 6, 6+n) + print(f"Mass matrix shape: {M_batch.shape}") + +**Combine JIT with Native Batching:** + +.. code-block:: python + + # JIT for maximum speed + @jit + def jit_batched_compute(w_H_b_batch, joints_batch): + M = kinDyn.mass_matrix(w_H_b_batch, joints_batch) + J = kinDyn.jacobian('l_sole', w_H_b_batch, joints_batch) + return M, J + + # First call compiles, subsequent calls are fast + M_batch, J_batch = jit_batched_compute(w_H_b_batch, joints_batch) + + +Optimization Example +-------------------- + +Use gradients for optimization: + +.. code-block:: python + + import optax # pip install optax + from jax import grad, jit + + def objective(joints): + """Minimize mass matrix trace""" + M = kinDyn.mass_matrix(w_H_b, joints) + return jnp.trace(M) + + # Setup optimizer + optimizer = optax.adam(learning_rate=0.01) + opt_state = optimizer.init(joints) + + # JIT the step + @jit + def step(joints, opt_state): + loss, grads = jax.value_and_grad(objective)(joints) + updates, opt_state = optimizer.update(grads, opt_state) + joints = optax.apply_updates(joints, updates) + return joints, opt_state, loss + + # Optimize + for i in range(100): + joints, opt_state, loss = step(joints, opt_state) + if i % 10 == 0: + print(f"Step {i}, Loss: {loss:.6f}") + + +Tips and Tricks +--------------- + +**Enable 64-bit precision** (recommended for robustness): + +.. code-block:: python + + from jax import config + config.update("jax_enable_x64", True) + +**Disable JIT temporarily for debugging:** + +.. code-block:: python + + from jax import config + config.update("jax_disable_jit", True) + +Loading from MuJoCo +------------------- + +Load models from MuJoCo and leverage JAX's JIT and autodiff: + +.. code-block:: python + + import mujoco + from robot_descriptions.loaders.mujoco import load_robot_description + from adam.jax import KinDynComputations + import jax.numpy as jnp + from jax import jit, grad + + # Load MuJoCo model + mj_model = load_robot_description("g1_mj_description") + + # Create KinDynComputations from MuJoCo model + kinDyn = KinDynComputations.from_mujoco_model(mj_model) + kinDyn.set_frame_velocity_representation(adam.Representations.MIXED_REPRESENTATION) + + # Use with JIT and autodiff + @jit + def compute_mass_trace(w_H_b, joints): + M = kinDyn.mass_matrix(w_H_b, joints) + return jnp.trace(M) + + w_H_b = jnp.eye(4) + joints = jnp.zeros(kinDyn.NDoF) + + trace_val = compute_mass_trace(w_H_b, joints) + grad_fn = grad(compute_mass_trace, argnums=1) + grad_val = grad_fn(w_H_b, joints) + +See :doc:`../guides/mujoco` for more details on MuJoCo integration. + +When to Use JAX +--------------- + +βœ… **Good for:** +- Gradient-based optimization +- Computing Jacobians and Hessians +- Processing batches with native batching +- GPU acceleration + +❌ **Not ideal for:** +- One-off computations (NumPy is faster) +- Symbolic manipulation (use CasADi) + +**See also:** :doc:`../guides/backend_selection` diff --git a/docs/quickstart/numpy.rst b/docs/quickstart/numpy.rst new file mode 100644 index 00000000..fde10dc6 --- /dev/null +++ b/docs/quickstart/numpy.rst @@ -0,0 +1,116 @@ +NumPy Backend +============== + +The NumPy backend is the simplest choice for direct computation. Use it for model validation, quick experiments, and debugging. + +Basic Usage +----------- + +.. code-block:: python + + import numpy as np + import adam + from adam.numpy import KinDynComputations + import icub_models + + # Load model + model_path = icub_models.get_model_file("iCubGazeboV2_5") + joints_name_list = [ + 'torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch', + 'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch', + 'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow', 'l_hip_pitch', 'l_hip_roll', + 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll', 'r_hip_pitch', + 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch', 'r_ankle_roll' + ] + + # Create KinDynComputations instance + kinDyn = KinDynComputations(model_path, joints_name_list) + kinDyn.set_frame_velocity_representation(adam.Representations.MIXED_REPRESENTATION) + + # Define state + w_H_b = np.eye(4) # Base at origin + joints = np.ones(len(joints_name_list)) * 0.1 # All joints at 0.1 rad + + # Compute dynamics + M = kinDyn.mass_matrix(w_H_b, joints) + print(f"Mass matrix shape: {M.shape}") + print(f"Mass matrix:\n{M}") + + # Forward kinematics + w_H_f = kinDyn.forward_kinematics('l_sole', w_H_b, joints) + print(f"Foot transform:\n{w_H_f}") + + # Jacobian + J = kinDyn.jacobian('l_sole', w_H_b, joints) + print(f"Jacobian shape: {J.shape}") + + +Common Operations +----------------- + +**Centroidal Momentum Matrix** + +.. code-block:: python + + Ag = kinDyn.centroidal_momentum_matrix(w_H_b, joints) + print(f"Centroidal momentum matrix shape: {Ag.shape}") + +**Center of Mass** + +.. code-block:: python + + com = kinDyn.CoM_position(w_H_b, joints) + print(f"CoM position: {com}") + + J_com = kinDyn.CoM_jacobian(w_H_b, joints) + print(f"CoM Jacobian shape: {J_com.shape}") + +**Bias Forces (Coriolis + Gravity)** + +.. code-block:: python + + base_vel = np.zeros(6) + joints_vel = np.zeros(len(joints_name_list)) + + h = kinDyn.bias_force(w_H_b, joints, base_vel, joints_vel) + print(f"Bias force: {h}") + +Loading from MuJoCo +------------------- + +You can also load models directly from MuJoCo: + +.. code-block:: python + + import mujoco + from robot_descriptions.loaders.mujoco import load_robot_description + from adam.numpy import KinDynComputations + + # Load MuJoCo model + mj_model = load_robot_description("g1_mj_description") + + # Create KinDynComputations from MuJoCo model + kinDyn = KinDynComputations.from_mujoco_model(mj_model) + kinDyn.set_frame_velocity_representation(adam.Representations.MIXED_REPRESENTATION) + + # Use as normal + w_H_b = np.eye(4) + joints = np.zeros(kinDyn.NDoF) + M = kinDyn.mass_matrix(w_H_b, joints) + +See :doc:`../guides/mujoco` for detailed MuJoCo integration guide. + +When to Use NumPy +----------------- + +βœ… **Good for:** +- Debugging and understanding results +- Code that doesn't need gradients or batching +- Educational examples + +❌ **Not ideal for:** +- Computing gradients (use JAX or PyTorch) +- Symbolic formulations (use CasADi) +- Large-scale optimization + +**Next steps:** If you need gradients and batching, see :doc:`jax` and :doc:`pytorch`. For optimization, see :doc:`casadi`. diff --git a/docs/quickstart/pytorch.rst b/docs/quickstart/pytorch.rst index 077ea9fc..dd9d899f 100644 --- a/docs/quickstart/pytorch.rst +++ b/docs/quickstart/pytorch.rst @@ -1,18 +1,30 @@ -PyTorch usage -============= +PyTorch Backend +=============== -The following example shows how to call an instance of the ``adam.pytorch.KinDynComputations`` class and use it to compute the mass matrix and forward dynamics of a floating-base robot. +The PyTorch backend offers GPU acceleration and automatic differentiation. Use it for learning-based control, large-scale optimization, and ML integration. + +Key Features +------------ + +- **GPU Acceleration** – Compute on GPUs +- **Automatic Differentiation** – Compute gradients naturally +- **Native Batching** – Efficiently process multiple configurations +- **Device Flexibility** – Easy CPU/GPU device management +- **ML Integration** – Works seamlessly with PyTorch models + +Basic Usage +----------- .. code-block:: python + import numpy as np + import torch import adam from adam.pytorch import KinDynComputations import icub_models - import numpy as np - - # if you want to icub-models https://github.com/robotology/icub-models to retrieve the urdf + + # Load model model_path = icub_models.get_model_file("iCubGazeboV2_5") - # The joint list joints_name_list = [ 'torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch', 'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch', @@ -20,13 +32,297 @@ The following example shows how to call an instance of the ``adam.pytorch.KinDyn 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll', 'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch', 'r_ankle_roll' ] + + kinDyn = KinDynComputations(model_path, joints_name_list) + kinDyn.set_frame_velocity_representation(adam.Representations.MIXED_REPRESENTATION) + + # Define state + w_H_b = torch.tensor(np.eye(4), dtype=torch.float64) + joints = torch.ones(len(joints_name_list), dtype=torch.float64) * 0.1 + + # Compute + M = kinDyn.mass_matrix(w_H_b, joints) + print(f"Mass matrix shape: {M.shape}") + + +GPU Support +----------- + +Move computations to GPU easily: + +.. code-block:: python + + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + # Move tensors to GPU + w_H_b = torch.tensor(np.eye(4), dtype=torch.float64, device=device) + joints = torch.ones(len(joints_name_list), dtype=torch.float64, device=device) * 0.1 + + # Compute on GPU + M = kinDyn.mass_matrix(w_H_b, joints) + print(f"Computation device: {M.device}") + + +Automatic Differentiation +-------------------------- + +Compute gradients naturally: + +.. code-block:: python + + # Enable gradient computation + w_H_b.requires_grad = True + joints.requires_grad = True + + # Forward pass + M = kinDyn.mass_matrix(w_H_b, joints) + loss = M.trace() + + # Backward pass + loss.backward() + + print(f"Gradient w.r.t. joints:\n{joints.grad}") + + +Common Operations +----------------- + +**Jacobian:** + +.. code-block:: python + J = kinDyn.jacobian('l_sole', w_H_b, joints) + print(f"Jacobian shape: {J.shape}") + +**Forward Kinematics:** + +.. code-block:: python + + w_H_f = kinDyn.forward_kinematics('l_sole', w_H_b, joints) + print(f"End-effector transform:\n{w_H_f}") + +**Center of Mass:** + +.. code-block:: python + + com = kinDyn.CoM_position(w_H_b, joints) + J_com = kinDyn.CoM_jacobian(w_H_b, joints) + + +Learning-Based Control +---------------------- + +Integration with neural networks: + +.. code-block:: python + + import torch.nn as nn + + class ControlPolicy(nn.Module): + def __init__(self, n_dof): + super().__init__() + self.net = nn.Sequential( + nn.Linear(n_dof, 64), + nn.ReLU(), + nn.Linear(64, n_dof) + ) + + def forward(self, joints): + return self.net(joints) + + # Use with adam + policy = ControlPolicy(len(joints_name_list)) + action = policy(joints) + + # Compute acc + acc = kinDyn.aba( + w_H_b, joints, base_vel, + joints_vel, action + ) + +Batch Processing +---------------- + +Use ``pytorch.KinDynComputations`` to process multiple configurations. + +.. note:: There is a class ``pytorch.KinDynComputationsBatch`` that has the functionality of ``pytorch.KinDynComputations``. + It exists to avoid API changes in existing code. New users should prefer ``pytorch.KinDynComputations`` for both single and batched computations. + +.. code-block:: python + + from adam.pytorch import KinDynComputations + kinDyn = KinDynComputations(model_path, joints_name_list) - # choose the representation you want to use the body fixed representation - kinDyn.set_frame_velocity_representation(adam.Representations.BODY_FIXED_REPRESENTATION) - # or, if you want to use the mixed representation (that is the default) kinDyn.set_frame_velocity_representation(adam.Representations.MIXED_REPRESENTATION) - w_H_b = np.eye(4) - joints = np.ones(len(joints_name_list)) + + # Create batch (batch dimension is FIRST) + batch_size = 1024 + w_H_b_batch = torch.tile(torch.eye(4, dtype=torch.float64), (batch_size, 1, 1)) + joints_batch = torch.randn(batch_size, len(joints_name_list), dtype=torch.float64) + + # Compute for all configurations at once + M_batch = kinDyn.mass_matrix(w_H_b_batch, joints_batch) + print(f"Output shape: {M_batch.shape}") # (batch_size, 6+n_dof, 6+n_dof) + + +**Batch Computations:** + +.. code-block:: python + + J_batch = kinDyn.jacobian('l_sole', w_H_b_batch, joints_batch) + # Shape: (batch_size, 6, 6+n) + + w_H_f_batch = kinDyn.forward_kinematics('l_sole', w_H_b_batch, joints_batch) + # Shape: (batch_size, 4, 4) + + Ag_batch = kinDyn.centroidal_momentum_matrix(w_H_b_batch, joints_batch) + # Shape: (batch_size, 6, 6+n) + +**Batch Gradients:** + +.. code-block:: python + + w_H_b_batch.requires_grad = True + joints_batch.requires_grad = True + + # Forward pass + M_batch = kinDyn.mass_matrix(w_H_b_batch, joints_batch) + loss = M_batch.mean() + + # Backward pass through entire batch + loss.backward() + + print(f"Gradient shape: {joints_batch.grad.shape}") + + +Batch Use Cases +--------------- + +**Trajectory Evaluation** + +.. code-block:: python + + # Evaluate entire trajectory + trajectory_configs = torch.randn(trajectory_length, n_dof) + + # Compute Jacobians for all configurations + J_trajectory = kinDyn.jacobian('l_sole', w_H_b_batch, trajectory_configs) + +**Simulation Rollouts** + +.. code-block:: python + + # Simulate multiple trajectories in parallel + batch_size = 256 + trajectory_steps = 100 + + for step in range(trajectory_steps): + # Compute dynamics for all parallel simulations + M_batch = kinDyn.mass_matrix(w_H_b_batch, joints_batch) + # Update joints... + + +Batch Tips and Tricks +--------------------- + +**Shape Requirements** + +Always remember: batch dimension is **first** + +.. code-block:: python + + # βœ… Correct + w_H_b_batch.shape # (batch, 4, 4) + joints_batch.shape # (batch, n_dof) + + # ❌ Wrong + w_H_b_batch.shape # (4, 4, batch) + joints_batch.shape # (n_dof, batch) + +**Consistent dtypes** + +.. code-block:: python + + # All tensors should have same dtype + assert w_H_b_batch.dtype == joints_batch.dtype + assert w_H_b_batch.dtype == torch.float64 # Recommended + +**Memory usage for large batches** + +.. code-block:: python + + # For very large batches, use smaller mini-batches + mini_batch_size = 256 # Adjust based on GPU memory + + # Process in chunks + for i in range(0, total_configs, mini_batch_size): + batch = configs[i:i+mini_batch_size] + result = kinDyn_batch.mass_matrix(w_H_b_batch[:len(batch)], batch) + + +Performance Tips +---------------- + +**Use float64 for stability (recommended):** + +.. code-block:: python + + torch.set_default_dtype(torch.float64) + +**Pin memory for faster CPU-GPU transfer:** + +.. code-block:: python + + w_H_b = torch.tensor(np.eye(4), dtype=torch.float64, pin_memory=True) + +**Use in-place operations carefully:** + +.. code-block:: python + + # Be careful with requires_grad when using in-place ops + joints = joints.clone().detach().requires_grad_(True) + + +Loading from MuJoCo +------------------- + +Load models from MuJoCo and use with PyTorch's autodiff: + +.. code-block:: python + + import mujoco + from robot_descriptions.loaders.mujoco import load_robot_description + from adam.pytorch import KinDynComputations + import torch + + # Load MuJoCo model + mj_model = load_robot_description("g1_mj_description") + + kinDyn = KinDynComputations.from_mujoco_model(mj_model) + kinDyn.set_frame_velocity_representation(adam.Representations.MIXED_REPRESENTATION) + + # Use with PyTorch + w_H_b = torch.eye(4, dtype=torch.float64) + joints = torch.zeros(kinDyn.NDoF, dtype=torch.float64, requires_grad=True) + M = kinDyn.mass_matrix(w_H_b, joints) - print(M) + loss = M.trace() + loss.backward() # Gradients in joints.grad + +See :doc:`../guides/mujoco` for more on MuJoCo integration. + +When to Use PyTorch +------------------- + +βœ… **Good for:** +- GPU computation +- Learning-based control +- ML pipeline integration +- Large-scale batch +- Gradient-based optimization + +❌ **Not ideal for:** +- Simple numeric evaluation (NumPy is easier) +- Symbolic computation (use CasADi instead) + +**See also:** :doc:`../guides/backend_selection` diff --git a/docs/quickstart/pytorch_batched.rst b/docs/quickstart/pytorch_batched.rst deleted file mode 100644 index 9809e8e2..00000000 --- a/docs/quickstart/pytorch_batched.rst +++ /dev/null @@ -1,40 +0,0 @@ -PyTorch batched usage -===================== - -The following example shows how to call an instance of the ``adam.pytorch.KinDynComputationsBatch`` class and use it to compute the mass matrix and forward dynamics of a floating-base robot. - -.. note:: - The first time you run a function from this module, it will take a bit longer to execute as they are being compiled by JAX. - -.. code-block:: python - - - import adam - from adam.pytorch import KinDynComputationsBatch - import icub_models - - # if you want to icub-models - model_path = icub_models.get_model_file("iCubGazeboV2_5") - # The joint list - joints_name_list = [ - 'torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch', - 'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch', - 'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow', 'l_hip_pitch', 'l_hip_roll', - 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll', 'r_hip_pitch', - 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch', 'r_ankle_roll' - ] - - kinDyn = KinDynComputationsBatch(model_path, joints_name_list) - # choose the representation you want to use the body fixed representation - kinDyn.set_frame_velocity_representation(adam.Representations.BODY_FIXED_REPRESENTATION) - # or, if you want to use the mixed representation (that is the default) - kinDyn.set_frame_velocity_representation(adam.Representations.MIXED_REPRESENTATION) - w_H_b = np.eye(4) - joints = np.ones(len(joints_name_list)) - - num_samples = 1024 - w_H_b_batch = torch.tensor(np.tile(w_H_b, (num_samples, 1, 1)), dtype=torch.float32) - joints_batch = torch.tensor(np.tile(joints, (num_samples, 1)), dtype=torch.float32) - - M = kinDyn.mass_matrix(w_H_b_batch, joints_batch) - w_H_f = kinDyn.forward_kinematics('frame_name', w_H_b_batch, joints_batch) diff --git a/examples/neural_ik.py b/examples/neural_ik.py new file mode 100644 index 00000000..9400fd98 --- /dev/null +++ b/examples/neural_ik.py @@ -0,0 +1,535 @@ +""" +Dual-hand neural inverse-kinematics with ADAM (PyTorch backend) on a MuJoCo model. + +What this does +-------------- +- Loads a humanoid robot from `robot_descriptions` using ADAM's `from_mujoco_model`. +- Automatically detects both left and right hand bodies. +- Trains a small MLP that maps (current joints, desired hand positions) -> joint configuration. +- Uses ADAM's differentiable forward kinematics for the loss; MuJoCo is only for + visualization (no MuJoCo gradients). + +Training uses synthetic data: +- Sample random goal configurations within joint limits. +- Compute corresponding hand positions via ADAM FK. +- Train the network to predict joint configurations that reach both target hand positions. +- Regularization keeps uninvolved joints near zero. + +Usage +----- + python3 examples/neural_ik.py \\ + --description g1_mj_description \\ + --epochs 200 --batch-size 512 \\ + --visualize + +Dependencies +------------ +pip install mujoco robot_descriptions torch +""" + +from __future__ import annotations + +import argparse +import math +import sys +import time +from dataclasses import dataclass +from typing import Iterable, Sequence + +import numpy as np +import torch + +try: + import mujoco +except Exception as exc: # pragma: no cover - optional dependency + raise SystemExit( + "This example needs 'mujoco'. Install with: pip install mujoco" + ) from exc + +try: + from robot_descriptions.loaders.mujoco import load_robot_description +except Exception as exc: # pragma: no cover - optional dependency + raise SystemExit( + "This example needs 'robot_descriptions[mujoco]'. " + "Install with: pip install 'robot_descriptions[mujoco]'" + ) from exc + +import adam +from adam.pytorch import KinDynComputations + + +DEFAULT_DESCRIPTION = "g1_mj_description" +DEFAULT_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +@dataclass +class JointLimits: + lower: torch.Tensor # shape (n_dof,) + upper: torch.Tensor # shape (n_dof,) + names: list[str] + + +def _load_model(description: str) -> mujoco.MjModel: + model = load_robot_description(description) + if model is None: + raise RuntimeError( + f"Could not load '{description}'. " + "Ensure the description is installed and available." + ) + return model + + +def _detect_hand_bodies(model: mujoco.MjModel) -> list[str]: + """Detect left and right hand/wrist bodies for dual-hand IK. + + Returns list of hand body names [left_hand, right_hand]. + """ + body_names = [ + mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, i) + for i in range(model.nbody) + ] + print("All body names:", body_names) + + left_hand = next( + ( + n + for n in body_names + if "left" in n.lower() and ("hand" in n.lower() or "wrist" in n.lower()) + ), + None, + ) + right_hand = next( + ( + n + for n in body_names + if "right" in n.lower() and ("hand" in n.lower() or "wrist" in n.lower()) + ), + None, + ) + + hands = [] + if left_hand: + hands.append(left_hand) + if right_hand: + hands.append(right_hand) + + if not hands: + raise RuntimeError("Could not detect hand bodies in the model") + + return hands + + +def _joint_permutation( + model: mujoco.MjModel, kd: KinDynComputations +) -> tuple[np.ndarray, list[str], list[str]]: + mj_joint_names = [ + mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, j_id) + for j_id in range(model.njnt) + if model.jnt_type[j_id] != mujoco.mjtJoint.mjJNT_FREE + ] + adam_joint_names = kd.rbdalgos.model.actuated_joints + perm = [mj_joint_names.index(name) for name in adam_joint_names] + return np.eye(len(mj_joint_names))[:, perm], mj_joint_names, adam_joint_names + + +def _joint_limits( + model: mujoco.MjModel, + permutation: np.ndarray, + joint_names: Sequence[str], + device: torch.device, + dtype: torch.dtype, +) -> JointLimits: + # MuJoCo jnt_range has shape (njnt, 2). Unbounded joints often have zeros. + rng = model.jnt_range[ + [ + j_id + for j_id in range(model.njnt) + if model.jnt_type[j_id] != mujoco.mjtJoint.mjJNT_FREE + ] + ] + # permutation maps MuJoCo joint order -> ADAM order; pull per-joint limits accordingly. + perm_idx = permutation.argmax(axis=0) + rng = rng[perm_idx] + lower = [] + upper = [] + for i, name in enumerate(joint_names): + lo, hi = rng[i] + if math.isclose(lo, 0.0) and math.isclose(hi, 0.0): + lo, hi = -math.pi, math.pi + lower.append(lo) + upper.append(hi) + return JointLimits( + lower=torch.as_tensor(lower, device=device, dtype=dtype), + upper=torch.as_tensor(upper, device=device, dtype=dtype), + names=list(joint_names), + ) + + +def _sample_joints( + limits: JointLimits, batch_size: int, generator: torch.Generator | None = None +) -> torch.Tensor: + u = torch.rand( + (batch_size, len(limits.names)), + device=limits.lower.device, + dtype=limits.lower.dtype, + generator=generator, + ) + return limits.lower + (limits.upper - limits.lower) * u + + +class NeuralIK(torch.nn.Module): + """Tiny MLP: (q, target_xyz...) -> q_out. Supports multiple targets.""" + + def __init__(self, n_dof: int, hidden: int = 256, n_targets: int = 1): + super().__init__() + self.n_targets = n_targets + input_dim = n_dof + 3 * n_targets # q + all target positions (xyz each) + self.net = torch.nn.Sequential( + torch.nn.Linear(input_dim, hidden), + torch.nn.ReLU(), + torch.nn.Linear(hidden, hidden), + torch.nn.ReLU(), + torch.nn.Linear(hidden, n_dof), + ) + + def forward(self, q: torch.Tensor, *target_positions: torch.Tensor) -> torch.Tensor: + """Predict target joint configuration given current q and one or more target XYZ positions.""" + # Concatenate all target positions + if len(target_positions) == 1: + targets = target_positions[0] + else: + targets = torch.cat(target_positions, dim=-1) + x = torch.cat([q, targets], dim=-1) + return self.net(x) + + +def train_neural_ik( + kd: KinDynComputations, + target_bodies: list[str], + limits: JointLimits, + permutation: np.ndarray, + mj_model: mujoco.MjModel, + epochs: int = 200, + batch_size: int = 512, + lr: float = 3e-3, + seed: int = 0, + visualize: bool = False, + viz_fps: float = 60.0, + joint_regularization: float = 0.01, +) -> None: + torch.manual_seed(seed) + device = limits.lower.device + dtype = limits.lower.dtype + n_dof = len(limits.names) + n_targets = len(target_bodies) + + model = NeuralIK(n_dof=n_dof, n_targets=n_targets).to(device=device, dtype=dtype) + opt = torch.optim.Adam( + model.parameters(), lr=lr, betas=(0.9, 0.999), weight_decay=1e-5 + ) + base = torch.eye(4, device=device, dtype=dtype).expand(batch_size, 4, 4) + + for epoch in range(1, epochs + 1): + q_curr = _sample_joints(limits, batch_size) + q_goal = _sample_joints(limits, batch_size) + + with torch.no_grad(): + # Compute all target positions + goal_positions = [] + for body in target_bodies: + goal_H = kd.forward_kinematics(body, base, q_goal) + goal_positions.append(goal_H[:, :3, 3]) + + # Single-step prediction for stable training + q_pred = model(q_curr, *goal_positions) + q_pred = torch.clamp(q_pred, limits.lower, limits.upper) + + # Compute all predicted positions + pred_positions = [] + for body in target_bodies: + pred_H = kd.forward_kinematics(body, base, q_pred) + pred_positions.append(pred_H[:, :3, 3]) + + # Primary task loss: reach all target positions + task_loss = sum( + torch.nn.functional.mse_loss(pred_pos, goal_pos) + for pred_pos, goal_pos in zip(pred_positions, goal_positions) + ) / len(target_bodies) + + # Regularization: penalize joint positions to keep uninvolved joints near zero + reg_loss = joint_regularization * torch.mean(q_pred**2) + + loss = task_loss + reg_loss + opt.zero_grad() + loss.backward() + opt.step() + + if epoch % max(1, epochs // 10) == 0: + # Compute average error across all targets + total_err = sum( + torch.linalg.norm(pred_pos - goal_pos, dim=1).mean().item() + for pred_pos, goal_pos in zip(pred_positions, goal_positions) + ) / len(target_bodies) + print( + f"[{epoch:04d}/{epochs}] loss={loss.item():.6f} " + f"task={task_loss.item():.6f} reg={reg_loss.item():.6f} " + f"pos_err={total_err:.5f}" + ) + + if visualize: + _visualize_solution( + kd=kd, + model=model, + mujoco_model=mj_model, + permutation=permutation, + target_bodies=target_bodies, + limits=limits, + base=base[:1], + fps=viz_fps, + ) + + +def _visualize_solution( + kd: KinDynComputations, + model: torch.nn.Module, + mujoco_model: mujoco.MjModel, + permutation: np.ndarray, + target_bodies: list[str], + limits: JointLimits, + base: torch.Tensor, + fps: float, +) -> None: + """Visualize dual-hand IK convergence in a MuJoCo viewer.""" + try: + import mujoco.viewer + except Exception as exc: # pragma: no cover - optional dependency + print(f"Viewer unavailable: {exc}") + return + + with torch.no_grad(): + # Start from zero joint configuration for debugging + q_curr = torch.zeros_like(limits.lower).unsqueeze(0) + + # Generate target positions based on zero configuration + # This places targets near where end effectors are at zero config + target_positions = [] + for body_name in target_bodies: + body_fk = kd.forward_kinematics(body_name, base, q_curr) + target_positions.append(body_fk[:, :3, 3]) + + # Print initial state + print("Starting IK convergence visualization from zero configuration.") + print(f"Tracking {len(target_bodies)} targets at end-effector positions:") + for i, (name, pos) in enumerate(zip(target_bodies, target_positions)): + print(f" [{i}] {name}: {pos.squeeze().cpu().numpy()}") + + data = mujoco.MjData(mujoco_model) + frame_time = 1.0 / max(fps, 1e-3) + + try: + with mujoco.viewer.launch_passive(mujoco_model, data) as viewer: + # Convert targets to numpy and store centers for motion + target_centers = [ + pos.squeeze().cpu().numpy().copy() for pos in target_positions + ] + target_np = [center.copy() for center in target_centers] + + # Run continuously until user closes the window + ik_step = 0 + print("Visualization running. Close the MuJoCo window to exit.") + + while viewer.is_running(): + # Update targets with small circular motion for debugging + t = (ik_step / 100.0) * 2 * np.pi # One full circle per 100 frames + radius = 0.25 + + # Generate small circular motions around end-effector positions + # Target order: [right_hand, left_foot, right_foot, left_hand] + for i, body_name in enumerate(target_bodies): + # Phase shift to desynchronize motions + phase = t + (i * np.pi / 2) + + # Small vertical circle motion + target_np[i] = np.array( + [ + 0.3, + 0.0 + radius * np.cos(phase), + 0.2 + radius * np.sin(phase), + ] + ) + + # Convert to tensors for model + target_tensors = [ + torch.tensor( + tnp, dtype=limits.lower.dtype, device=limits.lower.device + ).unsqueeze(0) + for tnp in target_np + ] + + # Perform IK iteration to track moving targets + with torch.no_grad(): + + # Apply model multiple times per frame for smooth convergence + for _ in range(3): + q_pred = model(q_curr, *target_tensors) + q_curr = torch.clamp(q_pred, limits.lower, limits.upper) + + # Compute primary target error for logging + curr_fk = kd.forward_kinematics(target_bodies[0], base, q_curr) + curr_pos = curr_fk[:, :3, 3] + pos_err = torch.linalg.norm( + curr_pos - target_tensors[0], dim=1 + ).item() + + # Get current positions for all tracked bodies + current_positions_np = [] + for body_name in target_bodies: + body_fk = kd.forward_kinematics(body_name, base, q_curr) + current_positions_np.append( + body_fk[:, :3, 3].squeeze().cpu().numpy() + ) + + # Update MuJoCo visualization + q_mj = permutation @ q_curr.detach().cpu().numpy().squeeze() + data.qpos[:] = 0.0 + if mujoco_model.nq >= 7: # free base assumed + data.qpos[3] = 1.0 # identity quaternion + data.qpos[7 : 7 + len(q_mj)] = q_mj + else: + data.qpos[: len(q_mj)] = q_mj + mujoco.mj_forward(mujoco_model, data) + + if ik_step % 10 == 0: + print( + f"[IK step {ik_step:4d}] " + f"pos_error: {pos_err:.6f} m | " + f"targets: {len(target_bodies)}" + ) + + # Add visual markers by modifying the scene directly + # Define colors for targets and current positions + target_colors = [ + np.array([1.0, 0.0, 0.0, 0.6]), # Red + np.array([0.0, 1.0, 1.0, 0.5]), # Cyan + np.array([1.0, 0.0, 1.0, 0.5]), # Magenta + np.array([1.0, 0.5, 0.0, 0.6]), # Orange + ] + current_colors = [ + np.array([0.0, 1.0, 0.0, 0.8]), # Green + np.array([0.0, 0.0, 1.0, 0.7]), # Blue + np.array([0.5, 0.0, 0.5, 0.7]), # Purple + np.array([1.0, 1.0, 0.0, 0.8]), # Yellow + ] + + with viewer.lock(): + viewer.user_scn.ngeom = len(target_bodies) * 2 + + geom_idx = 0 + for i in range(len(target_bodies)): + # Target sphere (larger, semi-transparent) + mujoco.mjv_initGeom( + viewer.user_scn.geoms[geom_idx], + mujoco.mjtGeom.mjGEOM_SPHERE, + np.array([0.05, 0, 0]), + target_np[i], + np.eye(3).flatten(), + target_colors[i % len(target_colors)], + ) + geom_idx += 1 + + # Current position sphere (smaller, more opaque) + mujoco.mjv_initGeom( + viewer.user_scn.geoms[geom_idx], + mujoco.mjtGeom.mjGEOM_SPHERE, + np.array([0.03, 0, 0]), + current_positions_np[i], + np.eye(3).flatten(), + current_colors[i % len(current_colors)], + ) + geom_idx += 1 + + viewer.sync() + time.sleep(frame_time) + + # Increment and loop + ik_step = (ik_step + 1) % 200 # Loop motion every 200 steps + + except Exception as exc: # pragma: no cover - viewer/GL issues + print(f"Viewer error: {exc}") + + +def parse_args(argv: Iterable[str]) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Neural IK with ADAM (PyTorch)") + parser.add_argument( + "--description", + type=str, + default=DEFAULT_DESCRIPTION, + help=f"robot_descriptions MuJoCo name (default: {DEFAULT_DESCRIPTION})", + ) + parser.add_argument("--epochs", type=int, default=200, help="Training epochs.") + parser.add_argument("--batch-size", type=int, default=512, help="Batch size.") + parser.add_argument("--lr", type=float, default=1e-3, help="Learning rate.") + parser.add_argument("--seed", type=int, default=0, help="RNG seed.") + parser.add_argument( + "--dtype", + type=str, + default="float32", + choices=("float32", "float64"), + help="Torch dtype for ADAM computations and the network.", + ) + parser.add_argument( + "--visualize", + action="store_true", + help="Launch a MuJoCo viewer after training.", + ) + parser.add_argument( + "--viz-fps", type=float, default=60.0, help="Viewer refresh rate." + ) + parser.add_argument( + "--joint-reg", + type=float, + default=0.001, + help="Regularization weight to keep non-task joints near zero.", + ) + return parser.parse_args(list(argv)) + + +def main(argv: Iterable[str] | None = None) -> None: + args = parse_args(sys.argv[1:] if argv is None else argv) + dtype = torch.float64 if args.dtype == "float64" else torch.float32 + + mj_model = _load_model(args.description) + print(f"Loaded '{args.description}'.") + + # Detect hand bodies for dual-hand IK + target_bodies = _detect_hand_bodies(mj_model) + print(f"Target bodies: {', '.join(target_bodies)}") + + kd = KinDynComputations.from_mujoco_model( + mj_model, + device=DEFAULT_DEVICE, + dtype=dtype, + ) + kd.set_frame_velocity_representation(adam.Representations.MIXED_REPRESENTATION) + + perm, mj_joint_names, adam_joint_names = _joint_permutation(mj_model, kd) + limits = _joint_limits(mj_model, perm, adam_joint_names, DEFAULT_DEVICE, dtype) + + train_neural_ik( + kd=kd, + target_bodies=target_bodies, + limits=limits, + permutation=perm, + mj_model=mj_model, + epochs=args.epochs, + batch_size=args.batch_size, + lr=args.lr, + seed=args.seed, + visualize=args.visualize, + viz_fps=args.viz_fps, + joint_regularization=args.joint_reg, + ) + + +if __name__ == "__main__": # pragma: no cover - example entry point + main() diff --git a/src/adam/model/std_factories/std_joint.py b/src/adam/model/std_factories/std_joint.py index c78f4bc5..78892df9 100644 --- a/src/adam/model/std_factories/std_joint.py +++ b/src/adam/model/std_factories/std_joint.py @@ -45,6 +45,8 @@ def _set_origin(self, origin: Pose) -> Pose: Returns: Pose: set the origin """ + if origin is None: + return Pose.zero(self.math) return Pose.build(xyz=origin.xyz, rpy=origin.rpy, math=self.math) def _set_limits(self, limit: Limits) -> Limits: diff --git a/src/adam/pytorch/__init__.py b/src/adam/pytorch/__init__.py index b427f55a..1c24e0bc 100644 --- a/src/adam/pytorch/__init__.py +++ b/src/adam/pytorch/__init__.py @@ -1,6 +1,6 @@ # Copyright (C) Istituto Italiano di Tecnologia (IIT). All rights reserved. -from .computation_batch import KinDynComputationsBatch from .computations import KinDynComputations +from .computations import KinDynComputations as KinDynComputationsBatch from .torch_like import TorchLike diff --git a/src/adam/pytorch/computation_batch.py b/src/adam/pytorch/computation_batch.py deleted file mode 100644 index 768a0412..00000000 --- a/src/adam/pytorch/computation_batch.py +++ /dev/null @@ -1,316 +0,0 @@ -# Copyright (C) Istituto Italiano di Tecnologia (IIT). All rights reserved. - -import warnings - -import numpy as np -import torch - -from adam.core.constants import Representations -from adam.core.rbd_algorithms import RBDAlgorithms -from adam.pytorch.torch_like import SpatialMath -from adam.model import Model, build_model_factory -from adam.model.kindyn_mixin import KinDynFactoryMixin -from adam.core.array_api_math import spec_from_reference - - -class KinDynComputationsBatch(KinDynFactoryMixin): - """ - A PyTorch-based class for batch kinematic and dynamic computations on robotic systems. - - Provides efficient processing of robot kinematics and dynamics calculations using PyTorch tensors. - Supports GPU acceleration and automatic differentiation for robotic applications. - """ - - def __init__( - self, - urdfstring: str, - joints_name_list: list = None, - device: torch.device = ( - torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") - ), - dtype: torch.dtype = torch.float32, - root_link: str = None, - gravity: torch.Tensor = torch.as_tensor([0, 0, -9.80665, 0, 0, 0]), - ) -> None: - """ - Args: - urdfstring (str): path/string of a URDF or a MuJoCo MjModel. - NOTE: The parameter name `urdfstring` is deprecated and will be renamed to `model` in a future release. - joints_name_list (list): list of the actuated joints - root_link (str, optional): Deprecated. The root link is automatically chosen as the link with no parent in the URDF. Defaults to None. - """ - ref = torch.tensor(0.0, dtype=dtype, device=device) - math = SpatialMath(spec_from_reference(ref)) - factory = build_model_factory(description=urdfstring, math=math) - model = Model.build(factory=factory, joints_name_list=joints_name_list) - self.rbdalgos = RBDAlgorithms(model=model, math=math) - self.NDoF = self.rbdalgos.NDoF - self.g = gravity.to(dtype=dtype, device=device) - self.funcs = {} - if root_link is not None: - warnings.warn( - "The root_link argument is not used. The root link is automatically chosen as the link with no parent in the URDF", - DeprecationWarning, - stacklevel=2, - ) - - def set_frame_velocity_representation( - self, representation: Representations - ) -> None: - """Sets the representation of the velocity of the frames - - Args: - representation (Representations): The representation of the velocity - """ - self.rbdalgos.set_frame_velocity_representation(representation) - - def mass_matrix( - self, base_transform: torch.Tensor, joint_positions: torch.Tensor - ) -> torch.Tensor: - """Returns the Mass Matrix functions computed the CRBA - - Args: - base_transform (torch.Tensor): The batch of homogenous transforms from base to world frame - joint_positions (torch.Tensor): The batch of joints position - - Returns: - M (torch.Tensor): The batch Mass Matrix - """ - - M, _ = self.rbdalgos.crba(base_transform, joint_positions) - return M.array - - def centroidal_momentum_matrix( - self, base_transform: torch.Tensor, joint_positions: torch.Tensor - ) -> torch.Tensor: - """Returns the Centroidal Momentum Matrix functions computed the CRBA - - Args: - base_transform (torch.Tensor): The homogenous transform from base to world frame - joint_positions (torch.Tensor): The joints position - - Returns: - Jcc (torch.Tensor): Centroidal Momentum matrix - """ - - _, Jcm = self.rbdalgos.crba(base_transform, joint_positions) - return Jcm.array - - def relative_jacobian( - self, frame: str, joint_positions: torch.Tensor - ) -> torch.Tensor: - - return self.rbdalgos.relative_jacobian(frame, joint_positions).array - - def jacobian_dot( - self, - frame: str, - base_transform: torch.Tensor, - joint_positions: torch.Tensor, - base_velocity: torch.Tensor, - joint_velocities: torch.Tensor, - ) -> torch.Tensor: - """Returns the Jacobian derivative relative to the specified frame - - Args: - frame (str): The frame to which the jacobian will be computed - base_transform (torch.Tensor): The homogenous transform from base to world frame - joint_positions (torch.Tensor): The joints position - base_velocity (torch.Tensor): The base velocity - joint_velocities (torch.Tensor): The joint velocities - - Returns: - Jdot (torch.Tensor): The Jacobian derivative relative to the frame - """ - - return self.rbdalgos.jacobian_dot( - frame, base_transform, joint_positions, base_velocity, joint_velocities - ).array - - def forward_kinematics( - self, frame: str, base_transform: torch.Tensor, joint_positions: torch.Tensor - ) -> torch.Tensor: - """Computes the forward kinematics between the root and the specified frame - - Args: - frame (str): The frame to which the fk will be computed - - Returns: - H (torch.Tensor): The fk represented as Homogenous transformation matrix - """ - - return self.rbdalgos.forward_kinematics( - frame, base_transform, joint_positions - ).array - - def jacobian( - self, frame: str, base_transform: torch.Tensor, joint_positions: torch.Tensor - ) -> torch.Tensor: - """Returns the Jacobian relative to the specified frame - - Args: - frame (str): The frame to which the jacobian will be computed - base_transform (torch.Tensor): The homogenous transform from base to world frame - joint_positions (torch.Tensor): The joints position - - Returns: - J (torch.Tensor): The Jacobian between the root and the frame - """ - return self.rbdalgos.jacobian(frame, base_transform, joint_positions).array - - def bias_force( - self, - base_transform: torch.Tensor, - joint_positions: torch.Tensor, - base_velocity: torch.Tensor, - joint_velocities: torch.Tensor, - ) -> torch.Tensor: - """Returns the bias force of the floating-base dynamics equation, - using a reduced RNEA (no acceleration and external forces) - - Args: - base_transform (torch.Tensor): The homogenous transform from base to world frame - joint_positions (torch.Tensor): The joints position - base_velocity (torch.Tensor): The base velocity - joint_velocities (torch.Tensor): The joints velocity - - Returns: - h (torch.Tensor): the bias force - """ - gravity = self.g - return self.rbdalgos.rnea( - base_transform, - joint_positions, - base_velocity, - joint_velocities, - gravity, - ).array.squeeze() - - def coriolis_term( - self, - base_transform: torch.Tensor, - joint_positions: torch.Tensor, - base_velocity: torch.Tensor, - joint_velocities: torch.Tensor, - ) -> torch.Tensor: - """Returns the coriolis term of the floating-base dynamics equation, - using a reduced RNEA (no acceleration and external forces) - - Args: - base_transform (torch.Tensor): The homogenous transform from base to world frame - joint_positions (torch.Tensor): The joints position - base_velocity (torch.Tensor): The base velocity - joint_velocities (torch.Tensor): The joints velocity - - Returns: - C (torch.Tensor): the Coriolis term - """ - gravity = torch.zeros( - (6,), - dtype=base_transform.dtype, - device=base_transform.device, - ) - return self.rbdalgos.rnea( - base_transform, - joint_positions, - base_velocity, - joint_velocities, - gravity, - ).array.squeeze() - - def gravity_term( - self, base_transform: torch.Tensor, joint_positions: torch.Tensor - ) -> torch.Tensor: - """Returns the gravity term of the floating-base dynamics equation, - using a reduced RNEA (no acceleration and external forces) - - Args: - base_transform (torch.Tensor): The homogenous transform from base to world frame - joint_positions (torch.Tensor): The joints position - - Returns: - G (jnp.array): the gravity term - """ - batch_size = base_transform.shape[:-2] - base_velocity = torch.zeros( - batch_size + (6,), - dtype=base_transform.dtype, - device=base_transform.device, - ) - joint_velocities = torch.zeros_like(joint_positions) - return self.rbdalgos.rnea( - base_transform, - joint_positions, - base_velocity, - joint_velocities, - self.g, - ).array.squeeze() - - def CoM_position( - self, base_transform: torch.Tensor, joint_positions: torch.Tensor - ) -> torch.Tensor: - """Returns the CoM position - - Args: - base_transform (torch.Tensor): The homogenous transform from base to world frame - joint_positions (torch.Tensor): The joints position - - Returns: - CoM (torch.Tensor): The CoM position - """ - return self.rbdalgos.CoM_position(base_transform, joint_positions).array - - def CoM_jacobian( - self, base_transform: torch.Tensor, joint_positions: torch.Tensor - ) -> torch.Tensor: - """Returns the CoM Jacobian - - Args: - base_transform (torch.Tensor): The homogenous transform from base to world frame - joint_positions (torch.Tensor): The joints position - - Returns: - Jcom (torch.Tensor): The CoM Jacobian - """ - return self.rbdalgos.CoM_jacobian(base_transform, joint_positions).array - - def get_total_mass(self) -> float: - """Returns the total mass of the robot - - Returns: - mass: The total mass - """ - return self.rbdalgos.get_total_mass() - - def aba( - self, - base_transform: torch.Tensor, - joint_positions: torch.Tensor, - base_velocity: torch.Tensor, - joint_velocities: torch.Tensor, - joint_torques: torch.Tensor, - external_wrenches: dict[str, torch.Tensor] | None = None, - ) -> torch.Tensor: - """Featherstone Articulated-Body Algorithm (floating base, O(n)). - - Args: - base_transform (torch.Tensor): The batch of homogenous transforms from base to world frame - joint_positions (torch.Tensor): The batch of joints position - base_velocity (torch.Tensor): The batch of base velocity - joint_velocities (torch.Tensor): The batch of joint velocities - joint_torques (torch.Tensor): The batch of joint torques - external_wrenches (dict[str, torch.Tensor], optional): External wrenches applied to the robot. Defaults to None. - - Returns: - torch.Tensor: The batch of base acceleration and the joint accelerations - """ - - return self.rbdalgos.aba( - base_transform, - joint_positions, - base_velocity, - joint_velocities, - joint_torques, - self.g, - external_wrenches, - ).array.squeeze() diff --git a/src/adam/pytorch/computations.py b/src/adam/pytorch/computations.py index c476d1c0..25da6b5a 100644 --- a/src/adam/pytorch/computations.py +++ b/src/adam/pytorch/computations.py @@ -26,7 +26,7 @@ def __init__( ), dtype: torch.dtype = torch.float32, root_link: str = None, - gravity: np.array = torch.tensor([0, 0, -9.80665, 0, 0, 0]), + gravity: torch.Tensor = torch.tensor([0, 0, -9.80665, 0, 0, 0]), ) -> None: """ Args: