working example obstacle avoidance

Author: lorenzo96-cmd

README.md

@@ -84,36 +84,6 @@ python mpx/examples/mjx_quad.py
 
 Use the keyboard's arrows to control the robot.
 
-## Diffusion-MPPI-guided FDDP
-
-MPX includes a receding-horizon manipulation solver that anneals a diagonal
-Gaussian in control-knot space, then refines its mean with FDDP. The local
-solver adds a temporary Gaussian precision to `l_u` and `l_uu`, decays that
-guidance each iteration, and reports physical task cost separately from prior
-cost. Original FDDP and primal-dual modes remain available.
-
-```bash
-# Interactive Push-T with live samples, mean, FDDP prediction, and history
-pixi run python -m mpx.examples.push_t_guided_mpc --viewer --verbose
-
-# Deterministic Push-T benchmark and plots
-pixi run python -m mpx.examples.push_t_guided_mpc \
-  --headless --seeds 0 1 2 3 4 --plot
-
-# Interactive full-MJX AgileX Piper box push
-pixi run python -m mpx.examples.agilex_box_push_guided_mpc --viewer --verbose
-
-# Nominal and mismatched-plant Piper benchmarks
-pixi run python -m mpx.examples.agilex_box_push_guided_mpc \
-  --headless --seeds 0 1 2 3 4 --condition nominal --plot
-pixi run python -m mpx.examples.agilex_box_push_guided_mpc \
-  --headless --condition heavy
-```
-
-Both examples accept `--mode fddp`, `--mode mppi`, `--mode mppi_fddp`, or
-`--mode guided`. See [`docs/diffusion_mppi_fddp.md`](docs/diffusion_mppi_fddp.md)
-for equations, ablations, video commands, diagnostics, model attribution, and
-measured limitations.
 
 > **Note:**  
 The first time running the script it can take more than a minute to JIT the solver

mpx/config/config_piper.py

@@ -27,9 +27,6 @@
 merit_penalty = jnp.array(1e-6,dtype=jnp.float32)
 
 
-q0 = jnp.array([0, 0, -0., -0., 0, 0])
-q0_init = q0
-
 
 n_joints = 6
 n_contact = len(contact_frame)
@@ -43,12 +40,12 @@
 qacc_slice = slice(0, nv)
 tau_slice = slice(nv, nv + n_joints)
 
-Qq = jnp.diag(jnp.ones(n_joints)) * 0.01
+Qq = jnp.diag(jnp.ones(n_joints)) * 0.1
 Qdq = jnp.diag(jnp.ones(n_joints)) * 0.2
-Qacc = jnp.diag(jnp.ones(nv)) * 0.001
+Qacc = jnp.diag(jnp.ones(nv)) * 0.01
 Qtau = jnp.diag(jnp.ones(n_joints)) * 0.0001
-Qee = jnp.diag(jnp.ones(3)) * 100.0
-Qee_final = jnp.diag(jnp.ones(3)) * 500.0
+Qee = jnp.diag(jnp.ones(3)) * 200.0
+Qee_final = jnp.diag(jnp.ones(3)) * 1000.0
 W = {
     "Qq": Qq,
     "Qdq": Qdq,
@@ -58,8 +55,9 @@
     "Qee_final": Qee_final,
 }
 
-initial_state = jnp.concatenate([q0, jnp.zeros(nv)])
 
+obstacle_center = jnp.array([0.3, 0.0, 0.6])
+obstacle_radius = 0.2
 max_torque = jnp.array([32.0, 32.0, 32.0, 8.0, 8.0, 8.0])
 min_torque = -max_torque
 _MODEL = mujoco.MjModel.from_xml_path(model_path)
@@ -68,6 +66,9 @@
     raise ValueError("All controlled Piper joints must define position limits")
 joint_position_min = jnp.asarray(_MODEL.jnt_range[:n_joints, 0])
 joint_position_max = jnp.asarray(_MODEL.jnt_range[:n_joints, 1])
+q0 = joint_position_min + 0.5 * (joint_position_max - joint_position_min)
+q0_init = q0
+initial_state = jnp.concatenate([q0, jnp.zeros(nv)])
 barrier_parameter = jnp.array(1.0, dtype=jnp.float32)
 multiplier_regularization = jnp.array(1e-8, dtype=jnp.float32)
 _INITIAL_DATA = mujoco.MjData(_MODEL)
@@ -135,13 +136,27 @@ def inequality(x, u, t, parameter):
     """Joint position bounds, feasible when values are <= 0."""
     del u, t, parameter
     qpos, _ = _state_parts(x)
+    data = mjx.make_data(_MJX_MODEL)
+    data = data.replace(qpos=qpos)
+    data = smooth.kinematics(_MJX_MODEL, data)
+    end_effector_position = data.geom_xpos[_CONTACT_ID[0]]
+    ## obstacle avoidance constraints - sphere centered at (0.3, 0.0, 0.4) with radius 0.2
+    distance_to_obstacle = jnp.linalg.norm(end_effector_position - obstacle_center)
+    obstacle_constraint = jnp.array([obstacle_radius - distance_to_obstacle])
     return jnp.concatenate(
         [
             joint_position_min - qpos,
             qpos - joint_position_max,
+            obstacle_constraint
         ]
     )
 
+def pseudo_huber(error, weight, delta):
+    squared_error = error.T @ weight @ error
+    return delta**2 * (
+        jnp.sqrt(1.0 + squared_error / delta**2) - 1.0
+    )
+
 def cost(W, reference, x, u, t):
     qpos, qvel = _state_parts(x)
     q = qpos
@@ -159,14 +174,15 @@ def cost(W, reference, x, u, t):
 
     end_effector_position = data.geom_xpos[_CONTACT_ID[0]]
     end_effector_error = end_effector_position - p_ref
+    ee_tracking_cost = pseudo_huber(end_effector_error, W["Qee"], delta=0.1)
     joint_error = q - q_ref
 
     stage_cost = (
           joint_error.T @ W["Qq"] @ joint_error
         + (dq - dq_ref).T @ W["Qdq"] @ (dq - dq_ref)
         + acc.T @ W["Qacc"] @ acc
         + tau.T @ W["Qtau"] @ tau
-        + end_effector_error.T @ W["Qee"] @ end_effector_error
+        + ee_tracking_cost
     )
     terminal_cost = (
         joint_error.T @ W["Qq"] @ joint_error

mpx/examples/agile_x.py

@@ -9,6 +9,9 @@
 
 
 import jax
+
+jax.config.update("jax_enable_x64", True)
+
 import jax.numpy as jnp
 import mujoco
 import mujoco.viewer
@@ -36,18 +39,29 @@ def solve_mpc(mpc_data, qpos, qvel, command):
 
     return solve_mpc
 
-def sample_reachable_position(rng, model, data, end_effector_id, qpos):
-    joint_range = model.jnt_range[: config.n_joints]
-    sample_qpos = rng.uniform(
-        low=joint_range[:, 0],
-        high=joint_range[:, 1],
-    )
-    data.qpos[:] = sample_qpos
-    data.qvel[:] = 0.0
-    mujoco.mj_forward(model, data)
-    if data.geom_xpos[end_effector_id][2] < 0.4:
-        return sample_reachable_position(rng, model, data, end_effector_id, qpos)
-    return data.geom_xpos[end_effector_id].copy()
+# def sample_reachable_position(rng, model, data, end_effector_id, qpos):
+#     joint_range = model.jnt_range[: config.n_joints]
+#     sample_qpos = rng.uniform(
+#         low=joint_range[:, 0],
+#         high=joint_range[:, 1],
+#     )
+#     data.qpos[:] = sample_qpos
+#     data.qvel[:] = 0.0
+#     mujoco.mj_forward(model, data)
+#     if data.geom_xpos[end_effector_id][2] < 0.4:
+#         return sample_reachable_position(rng, model, data, end_effector_id, qpos)
+#     return data.geom_xpos[end_effector_id].copy()
+def sample_position(indx,obstacle_center,obstacle_radius):
+    if indx == 0:
+        return np.array([0,0,obstacle_radius+0.01]) + obstacle_center
+    elif indx == 1:
+        return np.array([0.0,obstacle_radius+0.01,0.0]) + obstacle_center
+    elif indx == 2:
+        return np.array([0.0,0.0,-obstacle_radius-0.01]) + obstacle_center
+    elif indx == 3:
+        return np.array([0.0,-obstacle_radius-0.01,0.0]) + obstacle_center  
+    else:
+        raise ValueError("Invalid index for sample_position")
 
 def main():
     model = mujoco.MjModel.from_xml_path(dir_path + "/../data/piper_l/scene_flat.xml")
@@ -65,9 +79,8 @@ def main():
     end_effector_id = mujoco.mj_name2id(
         model, mujoco.mjtObj.mjOBJ_GEOM, config.contact_frame[0]
     )
-    desired_position = sample_reachable_position(
-        rng, model, target_data, end_effector_id, config.q0
-    )
+    des_pos_idx = 0
+    desired_position = sample_position(des_pos_idx, config.obstacle_center, config.obstacle_radius)
 
     data.qpos = np.asarray(config.q0)
     mujoco.mj_forward(model, data)
@@ -86,7 +99,7 @@ def main():
     q = jnp.zeros(config.n_joints)
     dq = jnp.zeros(config.n_joints)
     def step_controller():
-        nonlocal counter, tau, mpc_data, desired_position, target_start_counter, q, dq
+        nonlocal counter, tau, mpc_data, desired_position, target_start_counter, q, dq, des_pos_idx
 
         qpos = data.qpos.copy()
         qvel = data.qvel.copy()
@@ -107,9 +120,8 @@ def step_controller():
         counter += 1
         timed_out = counter - target_start_counter >= max_steps_per_target
         if timed_out:
-            desired_position = sample_reachable_position(
-                rng, model, target_data, end_effector_id, data.qpos
-            )
+            des_pos_idx = (des_pos_idx + 1) % 4
+            desired_position = sample_position(des_pos_idx, config.obstacle_center, config.obstacle_radius)
             target_start_counter = counter
             mpc_data = reset_mpc(
                 mpc.make_data(), data.qpos.copy(), data.qvel.copy()
@@ -120,13 +132,19 @@ def step_controller():
         model,
         data
     ) as viewer:
-        viewer.sync()
+        
         desired_position_geom_id = sim_utils.render_sphere(
                 viewer,
                 desired_position,
                 diameter=0.05,
                 color=np.array([0.0, 0.0, 1.0, 0.5]))
-
+        obstacle_geom_id = sim_utils.render_sphere(
+                viewer,
+                config.obstacle_center,
+                diameter=config.obstacle_radius*2,
+                color=np.array([1.0, 0.0, 0.0, 0.5])
+        )
+        viewer.sync()
         while viewer.is_running():
             tic = timer()
             step_controller()