refactor for a more readable cost function, ref and W are now dicts

Author: lorenzo96-cmd

mpx/config/config_acrobot_swingup.py

@@ -25,7 +25,7 @@
 I2 = m2 * l2 * l2 / 12.0
 
 parameter = jnp.zeros(N + 1)
-reference = jnp.zeros(N + 1)
+reference = {"dummy": jnp.zeros(N + 1)}
 
 
 def dynamics(x, u, t, parameter):
@@ -70,7 +70,7 @@ def dynamics(x, u, t, parameter):
 Q = jnp.diag(jnp.array([1e-5 / dt, 1e-5 / dt, 1e-5 / dt, 1e-5 / dt]))
 R = jnp.diag(jnp.array([1e-4 / dt]))
 Q_f = jnp.diag(jnp.array([10.0, 10.0, 100.0, 100.0]))
-W = jnp.zeros((N, 1))
+W = {"dummy": jnp.zeros((N, 1))}
 
 
 def cost(W, reference, x, u, t):

mpx/config/config_aliengo.py

@@ -1,6 +1,7 @@
 import jax.numpy as jnp
 import jax 
 import mpx.utils.models as mpc_dyn_model
+import mpx.utils.mpc_utils as mpc_utils
 import mpx.utils.objectives as mpc_objectives
 import os 
 from functools import partial
@@ -23,6 +24,7 @@
 step_height = 0.2 # Step height in meters
 initial_height = 0.1  # Initial height of the robot's base in meters
 robot_height = 0.36  # Height of the robot's base in meters
+clearance_speed = 0.2
 
 # Initial positions, orientations, and joint angles
 p0 = jnp.array([0, 0, robot_height])  # Initial position of the robot's base
@@ -42,6 +44,8 @@
 n =  13 + 2*n_joints + 6*n_contact  # Number of states (theta1, theta1_dot, theta2, theta2_dot)
 m = n_joints  # Number of controls (F)
 grf_as_state = True
+foot_slice = slice(13 + 2 * n_joints, 13 + 2 * n_joints + 3 * n_contact)
+leg_slice = foot_slice
 # Reference torques and controls (using n_joints)
 u_ref = jnp.zeros(m)  # Reference controls (concatenated torques)
 
@@ -58,8 +62,7 @@
 # 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))
 
-# Combine all cost matrices into a block diagonal matrix
-W = jax.scipy.linalg.block_diag(Qp, Qrot, Qq, Qdp, Qomega, Qdq, Qleg, Qtau,Q_grf)
+W = {"pos": Qp, "rot": Qrot, "q": Qq, "vel": Qdp, "omega": Qomega, "dq": Qdq, "contact": Qleg, "tau": Qtau, "grf": Q_grf}
 
 use_terrain_estimation = True  # Flag to use terrain estimation
 
@@ -68,8 +71,19 @@
     [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)
+cost = partial(mpc_objectives.quadruped_wb_obj, True, n_joints, n_contact, n_contact, N)
 hessian_approx = None
+reference_generator = partial(
+    mpc_utils.reference_generator,
+    use_terrain_estimation,
+    N,
+    dt,
+    n_joints,
+    n_contact,
+    foot0=p_legs0,
+    q0=q0,
+    clearence_speed=clearance_speed,
+)
 
 def dynamics(model, mjx_model, contact_id, body_id):
     return partial(

mpx/config/config_aliengo_trot_two_step.py

@@ -20,6 +20,7 @@
 step_height = 0.08
 initial_height = base.initial_height
 robot_height = base.robot_height
+clearance_speed = base.clearance_speed
 
 p0 = base.p0
 quat0 = base.quat0
@@ -31,6 +32,8 @@
 n = base.n
 m = base.m
 grf_as_state = base.grf_as_state
+foot_slice = base.foot_slice
+leg_slice = base.leg_slice
 u_ref = base.u_ref
 
 Qp = base.Qp
@@ -47,9 +50,20 @@
 use_terrain_estimation = False
 initial_state = base.initial_state
 
-cost = partial(mpc_objectives.quadruped_wb_obj, True, n_joints, n_contact, N)
+cost = partial(mpc_objectives.quadruped_wb_obj, True, n_joints, n_contact, n_contact, N)
 hessian_approx = base.hessian_approx
 dynamics = base.dynamics
+reference_generator = partial(
+    mpc_utils.reference_generator,
+    use_terrain_estimation,
+    N,
+    dt,
+    n_joints,
+    n_contact,
+    foot0=p_legs0,
+    q0=q0,
+    clearence_speed=clearance_speed,
+)
 
 reference = partial(
     mpc_utils.reference_quadruped_trot_two_step,

mpx/config/config_barrel_roll.py

@@ -25,6 +25,7 @@
 step_height = 0.12 # Step height in meters
 initial_height = 0.1  # Initial height of the robot's base in meters
 robot_height = 0.33  # Height of the robot's base in meters
+clearance_speed = 0.2
 grf_as_state = True
 
 # Initial positions, orientations, and joint angles
@@ -57,6 +58,8 @@
 n =  13 + 2*n_joints + 6*n_contact  # Number of states (theta1, theta1_dot, theta2, theta2_dot)
 m = n_joints  # Number of controls (F)
 grf_as_state = True
+foot_slice = slice(13 + 2 * n_joints, 13 + 2 * n_joints + 3 * n_contact)
+leg_slice = foot_slice
 # Reference torques and controls (using n_joints)
 tau_ref = jnp.zeros(n_joints)  # Reference torques (all zeros)
 # tau_ref = jnp.array([7.2171830e-02, -2.1473727e+00,  5.8485503e+00,  2.6923120e-03,
@@ -77,8 +80,7 @@
 # For the leg contact cost, repeat the unit cost for each contact point.
 Qleg = jnp.diag(jnp.tile(jnp.array([1e3,1e3,5e4]),n_contact))
 
-# Combine all cost matrices into a block diagonal matrix
-W = jax.scipy.linalg.block_diag(Qp, Qrot, Qq, Qdp, Qomega, Qdq, Qleg, Qtau,Q_grf)
+W = {"pos": Qp, "rot": Qrot, "q": Qq, "vel": Qdp, "omega": Qomega, "dq": Qdq, "contact": Qleg, "tau": Qtau, "grf": Q_grf}
 
 use_terrain_estimation = True  # Flag to use terrain estimation
 
@@ -87,8 +89,19 @@
     [p0, quat0, q0, jnp.zeros(6 + n_joints), p_legs0, jnp.zeros(_state_extra)]
 )
 
-cost = partial(mpc_objectives.quadruped_wb_obj, False, n_joints, n_contact, N)
+cost = partial(mpc_objectives.quadruped_wb_obj, False, n_joints, n_contact, n_contact, N)
 hessian_approx = None
+reference_generator = partial(
+    mpc_utils.reference_generator,
+    use_terrain_estimation,
+    N,
+    dt,
+    n_joints,
+    n_contact,
+    foot0=p_legs0,
+    q0=q0,
+    clearence_speed=clearance_speed,
+)
 
 def dynamics(model, mjx_model, contact_id, body_id):
     return partial(
@@ -105,4 +118,4 @@ def dynamics(model, mjx_model, contact_id, body_id):
 # dynamics = mpc_dyn_model.quadruped_wb_dynamics_learned_contact_model
 # dynamics = mpc_dyn_model.quadruped_wb_dynamics_explicit_contact
 max_torque = 40
-min_torque = -40
+min_torque = -40

mpx/config/config_go2.py

@@ -1,6 +1,7 @@
 import jax.numpy as jnp
 import jax 
 import mpx.utils.models as mpc_dyn_model
+import mpx.utils.mpc_utils as mpc_utils
 import mpx.utils.objectives as mpc_objectives
 import os 
 from functools import partial
@@ -23,6 +24,7 @@
 step_height = 0.065 # Step height in meters
 initial_height = 0.1  # Initial height of the robot's base in meters
 robot_height = 0.27  # Height of the robot's base in meters
+clearance_speed = 0.2
 
 # Initial positions, orientations, and joint angles
 p0 = jnp.array([0, 0, robot_height])  # Initial position of the robot's base
@@ -43,6 +45,8 @@
 n =  13 + 2*n_joints + 6*n_contact  # Number of states (theta1, theta1_dot, theta2, theta2_dot)
 m = n_joints  # Number of controls (F)
 grf_as_state = True
+foot_slice = slice(13 + 2 * n_joints, 13 + 2 * n_joints + 3 * n_contact)
+leg_slice = foot_slice
 # Reference torques and controls (using n_joints)
 u_ref = jnp.zeros(m)  # Reference controls (concatenated torques)
 
@@ -59,8 +63,7 @@
 # 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))
 
-# Combine all cost matrices into a block diagonal matrix
-W = jax.scipy.linalg.block_diag(Qp, Qrot, Qq, Qdp, Qomega, Qdq, Qleg, Qtau,Q_grf)
+W = {"pos": Qp, "rot": Qrot, "q": Qq, "vel": Qdp, "omega": Qomega, "dq": Qdq, "contact": Qleg, "tau": Qtau, "grf": Q_grf}
 
 use_terrain_estimation = True  # Flag to use terrain estimation
 
@@ -69,8 +72,19 @@
     [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)
+cost = partial(mpc_objectives.quadruped_wb_obj, True, n_joints, n_contact, n_contact, N)
 hessian_approx = None
+reference_generator = partial(
+    mpc_utils.reference_generator,
+    use_terrain_estimation,
+    N,
+    dt,
+    n_joints,
+    n_contact,
+    foot0=p_legs0,
+    q0=q0,
+    clearence_speed=clearance_speed,
+)
 
 def dynamics(model, mjx_model, contact_id, body_id):
     return partial(
@@ -85,4 +99,4 @@ def dynamics(model, mjx_model, contact_id, body_id):
 # dynamics = mpc_dyn_model.quadruped_wb_dynamics_learned_contact_model
 # dynamics = mpc_dyn_model.quadruped_wb_dynamics_explicit_contact
 max_torque = 25
-min_torque = -25
+min_torque = -25

mpx/config/config_h1.py

@@ -1,6 +1,7 @@
 import jax.numpy as jnp
 import jax 
 import mpx.utils.models as mpc_dyn_model
+import mpx.utils.mpc_utils as mpc_utils
 import mpx.utils.objectives as mpc_objectives
 import os 
 import sys 
@@ -25,6 +26,7 @@
 step_height = 0.08 # Step height in meters
 initial_height = 0.9 # Initial height of the robot's base in meters
 robot_height = 0.9  # Height of the robot's base in meters
+clearance_speed = 0.2
 
 # Initial positions, orientations, and joint angles
 p0 = jnp.array([0, 0, 0.9])  # Initial position of the robot's base
@@ -48,6 +50,8 @@
 n =  13 + 2*n_joints + 3*n_contact + 3*n_contact # Number of states
 m = n_joints # Number of controls (F)
 grf_as_state = True
+foot_slice = slice(13 + 2 * n_joints, 13 + 2 * n_joints + 3 * n_contact)
+leg_slice = foot_slice
 # Reference torques and controls (using n_joints)
 u_ref = jnp.zeros(m)  # Reference controls (concatenated torques)
 
@@ -71,8 +75,7 @@
 Qleg = jnp.diag(jnp.tile(jnp.array([1e3,1e3,1e5]),n_contact))
 Qgrf = jnp.diag(jnp.ones(3*n_contact))*1e-3
 
-# Combine all cost matrices into a block diagonal matrix
-W = jax.scipy.linalg.block_diag(Qp, Qrot, Qq, Qdp, Qomega, Qdq, Qleg, Qtau, Qgrf)
+W = {"pos": Qp, "rot": Qrot, "q": Qq, "vel": Qdp, "omega": Qomega, "dq": Qdq, "contact": Qleg, "tau": Qtau, "grf": Qgrf}
 
 use_terrain_estimation = False  # Flag to use terrain estimation
 
@@ -83,6 +86,17 @@
 
 cost = partial(mpc_objectives.h1_wb_obj, n_joints, n_contact, N)
 hessian_approx = None
+reference_generator = partial(
+    mpc_utils.reference_generator,
+    use_terrain_estimation,
+    N,
+    dt,
+    n_joints,
+    n_contact,
+    foot0=p_legs0,
+    q0=q0,
+    clearence_speed=clearance_speed,
+)
 
 def dynamics(model, mjx_model, contact_id, body_id):
     return partial(

mpx/config/config_h1_jump_forward.py

@@ -20,6 +20,7 @@
 step_height = base.step_height
 initial_height = base.initial_height
 robot_height = base.robot_height
+clearance_speed = base.clearance_speed
 
 p0 = base.p0
 quat0 = base.quat0
@@ -31,6 +32,8 @@
 n = base.n
 m = base.m
 grf_as_state = base.grf_as_state
+foot_slice = base.foot_slice
+leg_slice = base.leg_slice
 u_ref = base.u_ref
 W = base.W
 use_terrain_estimation = False
@@ -44,6 +47,17 @@
 hessian_approx = base.hessian_approx
 dynamics = base.dynamics
 MPCWrapper = base.MPCWrapper
+reference_generator = partial(
+    mpc_utils.reference_generator,
+    use_terrain_estimation,
+    N,
+    dt,
+    n_joints,
+    n_contact,
+    foot0=p_legs0,
+    q0=q0,
+    clearence_speed=clearance_speed,
+)
 
 reference = partial(
     mpc_utils.reference_humanoid_jump_forward,

mpx/config/config_h1_kinodynamic.py

@@ -5,6 +5,7 @@
 import jax.numpy as jnp
 
 import mpx.utils.models as mpc_dyn_model
+import mpx.utils.mpc_utils as mpc_utils
 import mpx.utils.mpc_wrapper as base_mpc_wrapper
 import mpx.utils.objectives as mpc_objectives
 
@@ -24,6 +25,7 @@
 step_height = 0.08
 initial_height = 0.9
 robot_height = 0.9
+clearance_speed = 0.2
 
 p0 = jnp.array([0.0, 0.0, 0.9])
 quat0 = jnp.array([1.0, 0.0, 0.0, 0.0])
@@ -47,6 +49,8 @@
 n = 13 + 2 * n_joints + 3 * n_contact
 m = n_joints + 3 * n_contact
 grf_as_state = False
+foot_slice = slice(13 + 2 * n_joints, 13 + 2 * n_joints + 3 * n_contact)
+leg_slice = foot_slice
 u_ref = jnp.zeros(m)
 
 Qp = jnp.diag(jnp.array([0.0, 0.0, 1e4]))
@@ -64,7 +68,7 @@
 Qleg = jnp.diag(jnp.tile(jnp.array([1e5, 1e5, 1e5]), n_contact))
 Qdq_cmd = jnp.diag(jnp.ones(n_joints)) * 1e-2
 Qgrf = jnp.diag(jnp.ones(3 * n_contact)) * 1e-3
-W = jax.scipy.linalg.block_diag(Qp, Qrot, Qq, Qdp, Qomega, Qdq, Qleg, Qdq_cmd, Qgrf)
+W = {"pos": Qp, "rot": Qrot, "q": Qq, "vel": Qdp, "omega": Qomega, "dq": Qdq, "contact": Qleg, "dq_cmd": Qdq_cmd, "grf": Qgrf}
 
 use_terrain_estimation = False
 initial_state = jnp.concatenate([p0, quat0, q0, jnp.zeros(6 + n_joints), p_legs0])
@@ -94,6 +98,17 @@
 
 cost = partial(mpc_objectives.h1_kinodynamic_obj, n_joints, n_contact, N)
 hessian_approx = None
+reference_generator = partial(
+    mpc_utils.reference_generator,
+    use_terrain_estimation,
+    N,
+    dt,
+    n_joints,
+    n_contact,
+    foot0=p_legs0,
+    q0=q0,
+    clearence_speed=clearance_speed,
+)
 
 
 def dynamics(model, mjx_model, contact_id, body_id):