Merge pull request #16 from iit-DLSLab/develop

MPX v0.1

Author: lorenzo96-cmd

.gitmodules

@@ -1,3 +1,3 @@
-[submodule "primal_dual_ilqr"]
-	path = mpx/primal_dual_ilqr
-	url = https://github.com/iit-DLSLab/primal_dual_ilqr.git
+[submodule "jax_ocp_solvers"]
+	path = mpx/jax_ocp_solvers
+	url = https://github.com/iit-DLSLab/jax_ocp_solvers.git

mpx/config/config_acrobot_swingup.py

@@ -0,0 +1,98 @@
+import os
+
+import jax
+import jax.numpy as jnp
+
+task_name = "acrobot_swingup"
+benchmark_mode = "direct"
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+scene_path = os.path.abspath(os.path.join(dir_path, "..")) + "/data/acrobot/scene.xml"
+
+N = 200
+n = 4
+m = 1
+dt = 0.01
+
+m1 = 1.0
+m2 = 1.0
+l1 = 0.5
+l2 = 0.5
+lc1 = 0.5 * l1
+lc2 = 0.5 * l2
+g = 9.81
+I1 = m1 * l1 * l1 / 12.0
+I2 = m2 * l2 * l2 / 12.0
+
+parameter = jnp.zeros(N + 1)
+reference = jnp.zeros(N + 1)
+
+
+def dynamics(x, u, t, parameter):
+    del t, parameter
+    theta1_dot = x[0]
+    theta2_dot = x[1]
+    theta1 = x[2]
+    theta2 = x[3]
+
+    d11 = (
+        I1
+        + I2
+        + m1 * lc1 * lc1
+        + m2 * (l1 * l1 + lc2 * lc2 + 2.0 * l1 * lc2 * jnp.cos(theta2))
+    )
+    d12 = I2 + m2 * (lc2 * lc2 + l1 * lc2 * jnp.cos(theta2))
+    d21 = d12
+    d22 = I2 + m2 * lc2 * lc2
+
+    c11 = -2.0 * m2 * l1 * lc2 * jnp.sin(theta2) * theta2_dot
+    c12 = -m2 * l1 * lc2 * jnp.sin(theta2) * theta2_dot
+    c21 = m2 * l1 * lc2 * jnp.sin(theta2) * theta1_dot
+    c22 = 0.0
+
+    g1 = m1 * g * lc1 * jnp.sin(theta1) + m2 * g * (l1 * jnp.sin(theta1) + lc2 * jnp.sin(theta1 + theta2))
+    g2 = m2 * lc2 * g * jnp.sin(theta1 + theta2)
+
+    D = jnp.array([[d11, d12], [d21, d22]])
+    C = jnp.array([[c11, c12], [c21, c22]])
+    G = jnp.array([g1, g2])
+
+    theta_dot = jnp.array([theta1_dot, theta2_dot])
+    theta_dot_new = theta_dot + dt * jnp.linalg.inv(D) @ (jnp.array([0.0, u[0]]) - C @ theta_dot - G)
+    theta_new = jnp.array([theta1, theta2]) + dt * theta_dot_new
+    return jnp.concatenate([theta_dot_new, theta_new])
+
+
+x0 = jnp.array([0.0, 0.0, 0.0, 0.0])
+x_ref = jnp.array([0.0, 0.0, jnp.pi, 0.0])
+u_ref = jnp.array([0.0])
+
+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))
+
+
+def cost(W, reference, x, u, t):
+    del W, reference
+    stage_cost = (x - x_ref).T @ Q @ (x - x_ref) + (u - u_ref).T @ R @ (u - u_ref)
+    term_cost = (x - x_ref).T @ Q_f @ (x - x_ref)
+    return jnp.where(t == N, 0.5 * term_cost, 0.5 * stage_cost)
+
+
+hessian_approx = None
+
+
+initial_X0 = jnp.tile(x0, (N + 1, 1))
+initial_U0 = jnp.tile(u_ref, (N, 1))
+initial_V0 = jnp.zeros((N + 1, n))
+
+solver_mode = "fddp"
+
+
+def state_to_qpos(x):
+    return x[2:]
+
+
+def state_to_qvel(x):
+    return x[:2]

mpx/config/config_aliengo.py

@@ -63,10 +63,27 @@
 
 use_terrain_estimation = True  # Flag to use terrain estimation
 
-cost = partial(mpc_objectives.quadruped_wb_obj,True)
-hessian_approx = partial(mpc_objectives.quadruped_wb_hessian_gn,True)
-dynamics = mpc_dyn_model.quadruped_wb_dynamics
+_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,
+    )
 # dynamics = mpc_dyn_model.quadruped_wb_dynamics_learned_contact_model
 # dynamics = mpc_dyn_model.quadruped_wb_dynamics_explicit_contact
 max_torque = 35
 min_torque = -35
+
+solver_mode = "primal_dual"  # Solver mode for the optimization problem

mpx/config/config_aliengo_trot_two_step.py

@@ -0,0 +1,66 @@
+from functools import partial
+
+import mpx.config.config_aliengo as base
+import mpx.utils.mpc_utils as mpc_utils
+import mpx.utils.objectives as mpc_objectives
+
+task_name = "aliengo_trot_two_step"
+
+model_path = base.model_path
+contact_frame = base.contact_frame
+body_name = base.body_name
+
+dt = 0.02
+N = 60
+mpc_frequency = base.mpc_frequency
+
+timer_t = base.timer_t
+duty_factor = 0.5
+step_freq = base.step_freq
+step_height = 0.08
+initial_height = base.initial_height
+robot_height = base.robot_height
+
+p0 = base.p0
+quat0 = base.quat0
+q0 = base.q0
+p_legs0 = base.p_legs0
+
+n_joints = base.n_joints
+n_contact = base.n_contact
+n = base.n
+m = base.m
+grf_as_state = base.grf_as_state
+u_ref = base.u_ref
+
+Qp = base.Qp
+Qrot = base.Qrot
+Qq = base.Qq
+Qdp = base.Qdp
+Qomega = base.Qomega
+Qdq = base.Qdq
+Qtau = base.Qtau
+Q_grf = base.Q_grf
+Qleg = base.Qleg
+W = base.W
+
+use_terrain_estimation = False
+initial_state = base.initial_state
+
+cost = partial(mpc_objectives.quadruped_wb_obj, True, n_joints, n_contact, N)
+hessian_approx = base.hessian_approx
+dynamics = base.dynamics
+
+reference = partial(
+    mpc_utils.reference_quadruped_trot_two_step,
+    base_height=robot_height,
+    total_forward=0.45,
+    step_length=0.16,
+    step_height=step_height,
+    settle_time=0.10,
+    phase_time=0.16,
+)
+
+solver_mode = "fddp"
+max_torque = base.max_torque
+min_torque = base.min_torque

mpx/config/config_barrel_roll.py

@@ -82,10 +82,26 @@
 
 use_terrain_estimation = True  # Flag to use terrain estimation
 
-cost = partial(mpc_objectives.quadruped_wb_obj,False)
-hessian_approx = partial(mpc_objectives.quadruped_wb_hessian_gn,False)
-dynamics = mpc_dyn_model.quadruped_wb_dynamics
+_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, False, 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,
+    )
 reference = mpc_utils.reference_barell_roll
+solver_mode = "fddp"  # Solver mode for the optimization problem
 # dynamics = mpc_dyn_model.quadruped_wb_dynamics_learned_contact_model
 # dynamics = mpc_dyn_model.quadruped_wb_dynamics_explicit_contact
 max_torque = 40

mpx/config/config_go2.py

@@ -64,9 +64,24 @@
 
 use_terrain_estimation = True  # Flag to use terrain estimation
 
-cost = partial(mpc_objectives.quadruped_wb_obj,True)
-hessian_approx = partial(mpc_objectives.quadruped_wb_hessian_gn,True)
-dynamics = mpc_dyn_model.quadruped_wb_dynamics
+_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,
+    )
 # dynamics = mpc_dyn_model.quadruped_wb_dynamics_learned_contact_model
 # dynamics = mpc_dyn_model.quadruped_wb_dynamics_explicit_contact
 max_torque = 25

mpx/config/config_h1.py

@@ -4,6 +4,7 @@
 import mpx.utils.objectives as mpc_objectives
 import os 
 import sys 
+from functools import partial
 dir_path = os.path.dirname(os.path.realpath(__file__))
 model_path = os.path.abspath(os.path.join(dir_path, '..')) + '/data/unitree_h1/mjx_h1_walk_real_feet.xml'  # Path to the MuJoCo model XML file
 # Joint names and related configuration
@@ -75,9 +76,26 @@
 
 use_terrain_estimation = False  # Flag to use terrain estimation
 
-cost = mpc_objectives.h1_wb_obj
-hessian_approx = mpc_objectives.h1_wb_hessian_gn
-dynamics = mpc_dyn_model.h1_wb_dynamics
+_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.h1_wb_obj, n_joints, n_contact, N)
+hessian_approx = None
+
+def dynamics(model, mjx_model, contact_id, body_id):
+    return partial(
+        mpc_dyn_model.h1_wb_dynamics,
+        model,
+        mjx_model,
+        contact_id,
+        body_id,
+        n_joints,
+        dt,
+    )
+
+solver_mode = "primal_dual"  # Solver mode for the optimization problem
 
 max_torque = 1000
-min_torque = -1000  
+min_torque = -1000  

mpx/config/config_h1_jump_forward.py

@@ -0,0 +1,60 @@
+from functools import partial
+
+import mpx.config.config_h1_kinodynamic as base
+import mpx.utils.mpc_utils as mpc_utils
+import mpx.utils.objectives as mpc_objectives
+
+task_name = "h1_jump_forward"
+
+model_path = base.model_path
+contact_frame = base.contact_frame
+body_name = base.body_name
+
+dt = 0.02
+N = 50
+mpc_frequency = base.mpc_frequency
+
+timer_t = base.timer_t
+duty_factor = 1.0
+step_freq = base.step_freq
+step_height = base.step_height
+initial_height = base.initial_height
+robot_height = base.robot_height
+
+p0 = base.p0
+quat0 = base.quat0
+q0 = base.q0
+p_legs0 = base.p_legs0
+
+n_joints = base.n_joints
+n_contact = base.n_contact
+n = base.n
+m = base.m
+grf_as_state = base.grf_as_state
+u_ref = base.u_ref
+W = base.W
+use_terrain_estimation = False
+initial_state = base.initial_state
+
+joint_kp = base.joint_kp
+joint_kd = base.joint_kd
+torque_limits = base.torque_limits
+
+cost = partial(mpc_objectives.h1_kinodynamic_obj, n_joints, n_contact, N)
+hessian_approx = base.hessian_approx
+dynamics = base.dynamics
+MPCWrapper = base.MPCWrapper
+
+reference = partial(
+    mpc_utils.reference_humanoid_jump_forward,
+    base_height=robot_height,
+    crouch_height=0.82,
+    apex_height=1.02,
+    jump_distance=0.35,
+    foot_shift=0.18,
+    foot_lift=0.10,
+)
+
+solver_mode = "fddp"
+max_torque = base.max_torque
+min_torque = base.min_torque