modified spot arm script to plot the end effector reference

Author: lorenzo96-cmd

mpx/examples/mjx_spot_arm.py

@@ -17,7 +17,7 @@
 import numpy as np
 
 import mpx.config.config_spot_arm as config
-import mpx.utils.mpc_wrapper as mpc_wrapper
+# import mpx.utils.mpc_wrapper as mpc_wrapper
 import mpx.utils.sim as sim_utils
 
 jax.config.update("jax_compilation_cache_dir", "./jax_cache")
@@ -32,7 +32,7 @@ def solve_mpc(mpc_data, qpos, qvel, foot, command, contact):
             mpc.initial_state
             .at[mpc.qpos_slice].set(qpos)
             .at[mpc.qvel_slice].set(qvel)
-            .at[mpc.foot_slice].set(foot)
+            .at[config.leg_slice].set(foot)
         )
         return mpc.run(mpc_data, x0, command, contact)
 
@@ -49,7 +49,7 @@ def main(headless=False, steps=500, scene="flat"):
     model.opt.timestep = 1 / sim_frequency
 
     contact_ids = sim_utils.geom_ids(model, config.contact_frame)
-    mpc = mpc_wrapper.MPCWrapper(config, limited_memory=True)
+    mpc = config.MpcWrapper(config, limited_memory=True)
     command_handle = sim_utils.KeyboardVelocityCommand()
     solve_mpc = _build_solve_fn(mpc)
     reset_mpc = jax.jit(mpc.reset)
@@ -62,7 +62,7 @@ def main(headless=False, steps=500, scene="flat"):
     mpc_data = reset_mpc(mpc.make_data(), data.qpos.copy(), data.qvel.copy(), foot)
 
     warm_command = jnp.asarray(command_handle.mpc_input(config.robot_height))
-    warm_contact = jnp.asarray(sim_utils.estimate_contacts(data, contact_ids))
+    warm_contact = jnp.asarray(sim_utils.estimate_contacts(data, contact_ids[:4]))
     mpc_data, tau = solve_mpc(
         mpc_data,
         data.qpos.copy(),
@@ -73,7 +73,8 @@ def main(headless=False, steps=500, scene="flat"):
     )
     tau.block_until_ready()
     mpc_data = reset_mpc(mpc_data, data.qpos.copy(), data.qvel.copy(), foot)
-
+    arm_ref_fun = jax.jit(config.extra_ref_fun)
+    arm_ref_data = {"foot": jnp.zeros((config.N + 1, 3))}
     period = int(sim_frequency / config.mpc_frequency)
     print(f"Controller period: {period} steps at {sim_frequency} Hz simulation frequency.")
     counter = 0
@@ -90,7 +91,7 @@ def step_controller():
             foot = jnp.asarray(sim_utils.geom_positions(data, contact_ids))
 
             command = jnp.asarray(command_handle.mpc_input(config.robot_height))
-            contact = jnp.asarray(sim_utils.estimate_contacts(data, contact_ids))
+            contact = jnp.asarray(sim_utils.estimate_contacts(data, contact_ids[:4]))
             print(f"Contact: {contact}")
             print(foot)
             print(f"Command: {command}")
@@ -111,7 +112,6 @@ def step_controller():
             # The shifted warm start is the next joint target used by the PD stabilizer.
             q_ref = mpc_data.X0[0, 7 : 7 + config.n_joints]
             print(f"MPC time: {1e3 * (stop - start):.2f} ms")
-
         data.ctrl = np.asarray(tau)
         mujoco.mj_step(model, data)
         counter += 1
@@ -120,19 +120,25 @@ def step_controller():
         for _ in range(steps):
             step_controller()
         return
-
+    arm_ref_sphere = []
+    arm_pred_sphere = []
     with mujoco.viewer.launch_passive(
         model,
         data,
         key_callback=command_handle.key_callback,
     ) as viewer:
         viewer.sync()
+        
         while viewer.is_running():
             overlay_text = command_handle.consume_overlay_text()
             tic = timer()
             if overlay_text is not None:
                 viewer.set_texts((None, None, *overlay_text))
             step_controller()
+            arm_ref_data["foot"] = jnp.zeros((config.N + 1, 1))
+            arm_ref_data_new = arm_ref_fun(arm_ref_data,counter * model.opt.timestep)
+            arm_ref_sphere = sim_utils.render_sphere_trajectory(viewer, arm_ref_data_new["foot"][:,1:],np.ones(config.N+1),0.03,np.array([0.0,1.0,0.0,1.0]),geom_ids = arm_ref_sphere)
+            arm_pred_sphere = sim_utils.render_sphere_trajectory(viewer, mpc_data.X0[:,13+2*config.n_joints + 12 :13+2*config.n_joints+15],np.ones(config.N+1),0.03,geom_ids = arm_pred_sphere)
             toc = timer()
             if toc - tic < model.opt.timestep:
                 sleep_time = model.opt.timestep - (toc - tic)