Add Spot and Spot with Arm

mpx/config/config_spot.py

@@ -0,0 +1,118 @@
+import os
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+import mpx.utils.models as mpc_dyn_model
+import mpx.utils.objectives as mpc_objectives
+
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+model_path = os.path.abspath(os.path.join(dir_path, "..")) + "/data/boston_dynamics_spot/scene.xml"
+
+# Contact frame names and body names for the Spot feet / lower legs.
+contact_frame = ["FL", "FR", "HL", "HR"]
+body_name = ["fl_lleg", "fr_lleg", "hl_lleg", "hr_lleg"]
+# contact_frame = ["FL", "FR", "RL", "RR"]
+# body_name = ["FL_calf", "FR_calf", "RL_calf", "RR_calf"]
+
+# Time and stage parameters.
+dt = 0.02
+N = 25
+mpc_frequency = 50
+
+# Gait parameters.
+timer_t = jnp.array([0.5, 0.0, 0.0, 0.5])
+duty_factor = 0.65
+step_freq = 1.35
+step_height = 0.12
+initial_height = 0.42
+robot_height = 0.42
+
+# Initial base state and nominal joint posture.
+p0 = jnp.array([0.0, 0.0, initial_height])
+quat0 = jnp.array([1.0, 0.0, 0.0, 0.0])
+q0 = jnp.array([0.0, 1.04, -1.8, 0.0, 1.04, -1.8, 0.0, 1.04, -1.8, 0.0, 1.04, -1.8])
+q0_init = q0
+
+# Nominal foot positions in the body frame at the home posture.
+p_legs0 = jnp.array([
+    0.34, 0.175, 0.0,
+    0.34, -0.175, 0.0,
+    -0.34, 0.175, 0.0,
+    -0.34, -0.175, 0.0,
+])
+
+# Dimensions.
+n_joints = 12
+n_contact = len(contact_frame)
+n = 13 + 2 * n_joints + 6 * n_contact
+m = n_joints
+grf_as_state = True
+
+# Reference controls.
+u_ref = jnp.zeros(m)
+
+# Cost weights.
+# Qp = jnp.diag(jnp.array([0.0, 0.0, 1e4]))
+# Qrot = jnp.diag(jnp.array([1000.0, 1000.0, 0.0])) * 1
+# Qq = jnp.diag(jnp.ones(n_joints)) * 1e0
+# Qdp = jnp.diag(jnp.array([1.0, 1.0, 1.0])) * 1e3
+# Qomega = jnp.diag(jnp.array([1.0, 1.0, 1.0])) * 1e2
+# Qdq = jnp.diag(jnp.ones(n_joints)) * 1e-1
+# Qtau = jnp.diag(jnp.ones(n_joints)) * 1e-2
+# Q_grf = jnp.diag(jnp.ones(3 * n_contact)) * 1e-3
+# Qleg = jnp.diag(jnp.tile(jnp.array([1e4, 1e4, 1e5]), n_contact))
+
+# Qp    = jnp.diag(jnp.array([0, 0, 1e4]))  # Cost matrix for position
+# Qrot  = jnp.diag(jnp.array([1000, 1000, 0]))  # Cost matrix for rotation
+# Qq    = jnp.diag(jnp.ones(n_joints)) * 1e-1 # Cost matrix for joint angles
+# Qdp   = jnp.diag(jnp.array([1, 1, 1])) * 5e3  # Cost matrix for position derivatives
+# Qomega= jnp.diag(jnp.array([1, 1, 1])) * 1e2  # Cost matrix for angular velocity
+# Qdq   = jnp.diag(jnp.ones(n_joints)) * 1e-1  # Cost matrix for joint angle derivatives
+# Qtau  = jnp.diag(jnp.ones(n_joints)) * 1e-1  # Cost matrix for torques
+# Q_grf = jnp.diag(jnp.ones(3*n_contact)) * 1e-2  # Cost matrix for ground reaction forces
+
+# # For the leg contact cost, repeat the unit cost for each contact point.
+# Qleg = jnp.diag(jnp.tile(jnp.array([1e4,1e4,1e5]),n_contact))
+
+# Cost matrices (diagonal matrices created using jnp.diag)
+Qp    = jnp.diag(jnp.array([0, 0, 1e4]))  # Cost matrix for position
+Qrot  = jnp.diag(jnp.array([1000, 1000, 0]))  # Cost matrix for rotation
+Qq    = jnp.diag(jnp.ones(n_joints)) * 1e-1 # Cost matrix for joint angles
+Qdp   = jnp.diag(jnp.array([1, 1, 1])) * 5e3  # Cost matrix for position derivatives
+Qomega= jnp.diag(jnp.array([1, 1, 1])) * 1e2  # Cost matrix for angular velocity
+Qdq   = jnp.diag(jnp.ones(n_joints)) * 1e-1  # Cost matrix for joint angle derivatives
+Qtau  = jnp.diag(jnp.ones(n_joints)) * 1e-2  # Cost matrix for torques
+Q_grf = jnp.diag(jnp.ones(3*n_contact)) * 1e-3 # Cost matrix for ground reaction forces
+
+# For the leg contact cost, repeat the unit cost for each contact point.
+Qleg = jnp.diag(jnp.tile(jnp.array([1e4,1e4,1e5]),n_contact))
+
+W = jax.scipy.linalg.block_diag(Qp, Qrot, Qq, Qdp, Qomega, Qdq, Qleg, Qtau, Q_grf)
+
+use_terrain_estimation = False
+_state_extra = n - (13 + 2 * n_joints + 3 * n_contact)
+initial_state = jnp.concatenate(
+    [p0, quat0, q0, jnp.zeros(6 + n_joints), p_legs0, jnp.zeros(_state_extra)]
+)
+
+cost = partial(mpc_objectives.quadruped_wb_obj, True, n_joints, n_contact, N)
+hessian_approx = None
+
+def dynamics(model, mjx_model, contact_id, body_id):
+    return partial(
+        mpc_dyn_model.quadruped_wb_dynamics,
+        model,
+        mjx_model,
+        contact_id,
+        body_id,
+        n_joints,
+        dt,
+    )
+
+# Torque bounds used by the MPC cost / clipping.
+max_torque = 500
+min_torque = -500
+solver_mode = "primal_dual"  # Solver mode for the optimization problem

mpx/config/config_spot_arm.py

@@ -0,0 +1,140 @@
+import os
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+import mpx.utils.models as mpc_dyn_model
+import mpx.utils.objectives as mpc_objectives
+
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+model_path = os.path.abspath(os.path.join(dir_path, "..")) + "/data/boston_dynamics_spot/scene_arm.xml"
+
+# Contact frame names and body names for the Spot feet / lower legs.
+contact_frame = ["FL", "FR", "HL", "HR"]
+body_name = ["fl_lleg", "fr_lleg", "hl_lleg", "hr_lleg"]
+# contact_frame = ["FL", "FR", "RL", "RR"]
+# body_name = ["FL_calf", "FR_calf", "RL_calf", "RR_calf"]
+
+# Time and stage parameters.
+dt = 0.02
+N = 25
+mpc_frequency = 50
+
+# Gait parameters.
+timer_t = jnp.array([0.5, 0.0, 0.0, 0.5])
+duty_factor = 0.65
+step_freq = 1.35
+step_height = 0.12
+# initial_height = 0.40
+# robot_height = 0.40
+initial_height = 0.46
+robot_height = 0.46
+
+# Initial base state and nominal joint posture.
+p0 = jnp.array([0.0, 0.0, initial_height])
+quat0 = jnp.array([1.0, 0.0, 0.0, 0.0])
+q0 = jnp.array([0, -2.14, 2.06, 0, 0, 0, 0, 0.0, 1.04, -1.8, 0.0, 1.04, -1.8, 0.0, 1.04, -1.8, 0.0, 1.04, -1.8])
+q0_init = q0
+
+# Nominal foot positions in the body frame at the home posture.
+p_legs0 = jnp.array([
+    0.34, 0.175, 0.0,
+    0.34, -0.175, 0.0,
+    -0.34, 0.175, 0.0,
+    -0.34, -0.175, 0.0,
+])
+# p_legs0 = jnp.array([
+#         0.26635823,  0.1658,      0.03254209,
+#         0.26635823, -0.1658,      0.03254209,
+#         -0.32934177,  0.1658,      0.03254209,
+#         -0.32934177, -0.1658,      0.03254209])
+
+# Dimensions.
+n_joints = 12 + 7
+n_contact = len(contact_frame)
+n = 13 + 2 * n_joints + 6 * n_contact
+m = n_joints
+grf_as_state = True
+
+# Reference controls.
+u_ref = jnp.zeros(m)
+
+# Cost weights.
+# Qp = jnp.diag(jnp.array([0.0, 0.0, 1e4]))
+# Qrot = jnp.diag(jnp.array([1000.0, 1000.0, 0.0])) * 10
+# # Qq = jnp.diag(jnp.ones(n_joints)) * 1e0
+# Qq = jnp.diag(jnp.concatenate([jnp.ones(7) * 1e3, jnp.ones(12) * 1e-1]))
+# Qdp = jnp.diag(jnp.array([1.0, 1.0, 1.0])) * 1e3
+# Qomega = jnp.diag(jnp.array([1.0, 1.0, 10.0])) * 1e2
+# Qdq = jnp.diag(jnp.ones(n_joints)) * 1e-1
+# Qtau = jnp.diag(jnp.ones(n_joints)) * 1e-2
+# Q_grf = jnp.diag(jnp.ones(3 * n_contact)) * 1e-2
+# Qleg = jnp.diag(jnp.tile(jnp.array([1e4, 1e4, 1e5]), n_contact))
+
+# Values for go2
+Qp    = jnp.diag(jnp.array([0, 0, 1e4]))  # Cost matrix for position
+Qrot  = jnp.diag(jnp.array([1000, 1000, 0]))  # Cost matrix for rotation
+Qq    = jnp.diag(jnp.concatenate([jnp.ones(7) * 5e2, jnp.ones(12) * 1e-1])) # Cost matrix for joint angles
+Qdp   = jnp.diag(jnp.array([1, 1, 1])) * 5e3  # Cost matrix for position derivatives
+Qomega= jnp.diag(jnp.array([1, 1, 1])) * 1e2  # Cost matrix for angular velocity
+Qdq   = jnp.diag(jnp.ones(n_joints)) * 1e-1  # Cost matrix for joint angle derivatives
+Qtau  = jnp.diag(jnp.ones(n_joints)) * 1e-2  # Cost matrix for torques
+Q_grf = jnp.diag(jnp.ones(3*n_contact)) * 1e-3 # Cost matrix for ground reaction forces
+
+# For the leg contact cost, repeat the unit cost for each contact point.
+Qleg = jnp.diag(jnp.tile(jnp.array([1e4,1e4,1e5]),n_contact))
+
+W = jax.scipy.linalg.block_diag(Qp, Qrot, Qq, Qdp, Qomega, Qdq, Qleg, Qtau, Q_grf)
+
+use_terrain_estimation = False
+_state_extra = n - (13 + 2 * n_joints + 3 * n_contact)
+initial_state = jnp.concatenate(
+    [p0, quat0, q0, jnp.zeros(6 + n_joints), p_legs0, jnp.zeros(_state_extra)]
+)
+
+cost = partial(mpc_objectives.quadruped_wb_obj, True, n_joints, n_contact, N)
+hessian_approx = None
+
+def dynamics(model, mjx_model, contact_id, body_id):
+    return partial(
+        mpc_dyn_model.quadruped_wb_dynamics,
+        model,
+        mjx_model,
+        contact_id,
+        body_id,
+        n_joints,
+        dt,
+    )
+
+# Torque bounds used by the MPC cost / clipping.
+max_torque = 500
+min_torque = -500
+solver_mode = "primal_dual"  # Solver mode for the optimization problem
+
+extra_qref_data = {
+    "amp": jnp.array([2.0, 0.5, 0.4, 0.6, 0.7, 0.8, 0.5]),
+    "freq": jnp.array([0.2, 1.0, 1.2, 0.8, 0.9, 1.0, 0.5]),
+    "joint_index": jnp.arange(7, dtype=jnp.int32),
+}
+
+def extra_qref_fn(q_ref, current_time, data):
+    if data is None:
+        return q_ref
+
+    arm_amp_ref = data["amp"]
+    arm_freq_ref = data["freq"]
+    arm_joint_index = data["joint_index"]
+
+    def arm_fn(t, carry):
+        q_ref = carry
+        #
+        time_n = t * dt + current_time
+        arm_pos = arm_amp_ref * jnp.sin(2 * jnp.pi * arm_freq_ref * time_n) + q0[arm_joint_index]
+
+        q_ref = q_ref.at[t,arm_joint_index].set(arm_pos)
+        return (q_ref)
+    init_carry = q_ref
+    q_ref = jax.lax.fori_loop(0, N+1, arm_fn, init_carry)
+    return q_ref

mpx/data/boston_dynamics_spot/CHANGELOG.md

@@ -0,0 +1,6 @@
+# Changelog – Boston Dynamics Spot Description
+
+All notable changes to this model will be documented in this file.
+
+## [2024-05-31]
+- Initial release.

mpx/data/boston_dynamics_spot/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2021, Clearpath Robotics Inc.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

mpx/data/boston_dynamics_spot/README.md

@@ -0,0 +1,38 @@
+# Boston Dynamics Spot Description (MJCF)
+
+> [!IMPORTANT]
+> Requires MuJoCo 3.1.3 or later.
+
+## Changelog
+
+See [CHANGELOG.md](./CHANGELOG.md) for a full history of changes.
+
+## Overview
+
+This package contains a simplified robot description (MJCF) of the [Spot
+Quadruped](https://bostondynamics.com/products/spot/) developed by [Boston
+Dynamics](https://bostondynamics.com/). It is derived from the [publicly
+available URDF description](https://github.com/bdaiinstitute/spot_ros2).
+
+<p float="left">
+  <img src="spot.png" width="400">
+</p>
+
+## URDF → MJCF Conversion
+
+1. Processed `.obj` files with [`obj2mjcf`](https://github.com/kevinzakka/obj2mjcf).
+2. Added `<mujoco> <compiler discardvisual="false" strippath="false" fusestatic="false"/> </mujoco>` to the URDF's
+   `<robot>` clause in order to preserve visual geometries.
+3. Loaded the URDF into MuJoCo and saved a corresponding MJCF.
+    * The spot XML was derived from [spot_simple.urdf.xacro](https://github.com/bdaiinstitute/spot_ros2/blob/main/spot_description/urdf/spot_simple.urdf.xacro).
+    * The arm XML was derived from [spot_arm_macro.urdf](https://github.com/bdaiinstitute/spot_ros2/blob/main/spot_description/urdf/spot_arm_macro.urdf).
+4. Added a `<freejoint/>` to the base, and a tracking light.
+5. Manually edited the MJCF to extract common properties into the `<default>` section.
+6. Added `<exclude>` clauses to prevent collisions between the base and the upper legs.
+7. Added position actuators (not tuned).
+8. Added a home keyframe.
+9. Added `scene.xml` which includes the robot, with a textured groundplane, skybox, and haze.
+
+## License
+
+This model is released under a [BSD-3-Clause License](LICENSE).