- add extra_qref_fn in reference generator

Author: Schulze18

mpx/config/config_spot_arm.py

@@ -111,4 +111,30 @@ def dynamics(model, mjx_model, contact_id, body_id):
 # Torque bounds used by the MPC cost / clipping.
 max_torque = 500
 min_torque = -500
-solver_mode = "primal_dual"  # Solver mode for the optimization problem
+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/utils/mpc_utils.py

@@ -35,8 +35,8 @@ def terrain_orientation(liftoff_pos,Ryaw):
 
     return jnp.roll(quat,1)
 
-@partial(jax.jit, static_argnums=(0,1,2,3,4,5))
-def reference_generator(use_terrain_estimator,N,dt,n_joints,n_contact,mass,foot0,q0,t_timer, x, foot, input, duty_factor, step_freq,step_height,liftoff,contact,clearence_speed):
+@partial(jax.jit, static_argnums=(0,1,2,3,4,5), static_argnames=("extra_qref_fn",))
+def reference_generator(use_terrain_estimator,N,dt,n_joints,n_contact,mass,foot0,q0,t_timer, x, foot, input, duty_factor, step_freq,step_height,liftoff,contact,clearence_speed, current_time, extra_qref_fn=None, extra_qref_data=None):
     p = x[:3]
     quat = x[3:7]
     # q = x[7:7+n_joints]
@@ -141,6 +141,25 @@ def cubic_splineZ(current_foot, foothold, step_height,val):
     liftoff = liftoff.at[1::3].set(liftoff_y)
     liftoff = liftoff.at[2::3].set(liftoff_z)
 
+    # ## Reference for the arm
+    # nq_arm = 7
+    # arm_amp_ref = jnp.array([2.0, 0.5, 0.4, 0.6, 0.7, 0.8, 0.5])
+    # arm_freq_ref = jnp.array([0.2, 1.0, 1.2, 0.8, 0.9, 1.0, 0.5])
+    # larm_index = 0
+    # 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[larm_index:larm_index+nq_arm]
+
+    #     q_ref = q_ref.at[t,larm_index:larm_index+nq_arm].set(arm_pos)
+    #     return (q_ref)
+    # init_carry = q_ref
+    # q_ref = jax.lax.fori_loop(0, N+1, arm_fn, init_carry)
+
+    if extra_qref_fn is not None:
+        q_ref = extra_qref_fn(q_ref, current_time, extra_qref_data)
+
     return jnp.concatenate([p_ref, quat_ref, q_ref, dp_ref, omega_ref, foot_ref, contact_sequence,grf_ref], axis=1),jnp.concatenate([contact_sequence], axis=1), liftoff
 
 @partial(jax.jit, static_argnums=(0,1,2,3))

mpx/utils/mpc_wrapper.py

@@ -15,6 +15,7 @@ class MPCData(PyTreeNode):
     """Carry state for the pure functional MPC API."""
 
     dt: float
+    time: jnp.ndarray
     duty_factor: float
     step_freq: float
     step_height: float
@@ -24,6 +25,7 @@ class MPCData(PyTreeNode):
     U0: jnp.ndarray
     V0: jnp.ndarray
     W: jnp.ndarray
+    extra_qref_data: object = None
 
 
 mpx_data = MPCData
@@ -165,19 +167,27 @@ def __init__(self, config, limited_memory=False):
 
         reference_generator = getattr(config, "reference_generator", mpc_utils.reference_generator)
         clearance_speed = getattr(config, "clearance_speed", getattr(config, "clearence_speed", 0.2))
-        self._ref_gen = jax.jit(
-            partial(
-                reference_generator,
-                config.use_terrain_estimation,
-                config.N,
-                config.dt,
-                config.n_joints,
-                config.n_contact,
-                robot_mass,
-                foot0=config.p_legs0,
-                q0=config.q0,
-                clearence_speed=clearance_speed,
-            )
+
+        ref_gen_kwargs = {
+            "foot0": config.p_legs0,
+            "q0": config.q0,
+            "clearence_speed": clearance_speed,
+        }
+        if hasattr(config, "extra_qref_fn"):
+            ref_gen_kwargs["extra_qref_fn"] = config.extra_qref_fn
+
+        # `reference_generator` is already jitted in mpc_utils. Keep a single
+        # binding layer here to avoid argument-position ambiguities when an
+        # extra q-ref callback is provided from config.
+        self._ref_gen = partial(
+            reference_generator,
+            config.use_terrain_estimation,
+            config.N,
+            config.dt,
+            config.n_joints,
+            config.n_contact,
+            robot_mass,
+            **ref_gen_kwargs,
         )
         self._timer_run = jax.jit(mpc_utils.timer_run)
         self._update_warm_start = partial(
@@ -193,6 +203,7 @@ def make_data(self):
 
         return MPCData(
             dt=self.config.dt,
+            time=jnp.asarray(0.0, dtype=jnp.float32),
             duty_factor=self.config.duty_factor,
             step_freq=self.config.step_freq,
             step_height=self.config.step_height,
@@ -202,6 +213,7 @@ def make_data(self):
             U0=self.initial_U0,
             V0=self.initial_V0,
             W=self.config.W,
+            extra_qref_data=getattr(self.config, "extra_qref_data", None),
         )
 
     def control_output(self, x0, X, U, reference, parameter):
@@ -226,6 +238,8 @@ def _run_impl(self, data, x0, input, contact):
             input=input,
             liftoff=data.liftoff,
             contact=contact,
+            current_time=data.time,
+            extra_qref_data=data.extra_qref_data,
         )
 
         # Reference generation and solver execution stay on the pure JAX path.
@@ -255,7 +269,9 @@ def _run_impl(self, data, x0, input, contact):
             V,
         )
 
+        new_time = data.time + jnp.asarray(1 / self.mpc_frequency, dtype=data.time.dtype)
         data = data.replace(
+            time=new_time,
             X0=X0,
             U0=U0,
             V0=V0,
@@ -286,6 +302,7 @@ def reset(self, data, qpos, qvel, foot):
             U0=self.initial_U0,
             X0=jnp.tile(initial_state, (self.config.N + 1, 1)),
             V0=self.initial_V0,
+            time=jnp.asarray(0.0, dtype=jnp.float32),
             contact_time=self.config.timer_t,
             liftoff=jnp.ravel(foot),
         )