IsaacLab Implementation?

[arjun-2612]

Good evening,

I am wondering if this approach can be ported to IsaacLab, or if you have an example of parallelized MPC in IsaacLab. Currently I have a working version in IsaacLab, but I find that with massive environments in parallel, it is really slow. I am wondering if your implementation with JAX is also given in IsaacLab. Please let me know.

Best,
Arjun.

[lorenzo96-cmd]

Hi Arjun,

Yes, you can definitely use MPX with IsaacLab. I’ve already tested this, and in the future I’ll include a simple example in the docs. In the meantime, you can use this example
as a reference. In particular, you can use the same function
to compute the torques for all environments in parallel.

A few considerations:

• Remember to set XLA_PYTHON_CLIENT_PREALLOCATE=false to avoid JAX preallocating 75% of the GPU memory.

• Use jax_dlpack.from_dlpack() to convert x0 and the inputs from PyTorch → JAX.

• Use torch_dlpack.from_dlpack() to convert the torques back from JAX → PyTorch.

• Keep an eye on the joint order and the reference frame for the velocities in your observation so they stay consistent with MPX.

Best,
Lorenzo

P.S. If you want to try it out, we can use this issue to share tips and help you integrate it.

[LockedFlysher]

I have tested this, it runs fast, 70,000 steps/s for SRBD.

[arjun-2612]

Oh I see. So I will convert my torch tensors to JAX, run the MPX function given in the example, and once I solve the GRFs, I can convert back to torch? Right now I have my centroidal MPC formulation fully using torch and I use the qpth QPFunction() to calculate GRFs. So if I have understood what you have said correctly, that would really help a lot! Let me try this and let you know.

[arjun-2612]

Also another question: does your MPC implementation give out position commands by any chance? Or is it only torques? I'd like to implement this on a robot that can't do torque control directly. Please let me know.

[lorenzo96-cmd]

If you want to use the SRBD model, it doesn’t compute the inverse kinematics so you won’t have the joint-space references (but this can be easily implemented in a batched way).
On the other hand, if you use the WB model, you can obtain q_des and dq_des from here.
Currently, these two outputs are not assigned, but you can simply add them and return them as is done for the torque.

[LockedFlysher]

And full centroidal dynamics at 3000 steps/s, with q_des/dq_des in optimized trajectory.

[arjun-2612]

Hey all, reopening this issue after looking into the IsaacLab side implementation and I have a few questions.

So I have the manager based environment for my RL environment in IsaacLab, and I have an action config that implements centroidal MPC and solves the QP. Now I find that the qpth optimizer I am using to solve the QP is very unoptimized, so when I run 4096 environments to train the RL, it takes forever to run even an epoch.

I have attached below my code segment which calculates the foot forces after building the dynamics, objective, and constraint matrices for MPC. I follow the MIT Cheetah Raibert + MPC paper implementation for my quadruped. The below code is part of the entire action config in IsaacLab.

The example you linked earlier, do I just import the file and run the solve part alone? Or should I change a bunch of things? Please let me know. Sorry for the trouble!

def apply_actions(self):
        # Uncomment all the lines commented out to run the MPC with Raibert. For RL, comment them back
        state_changed = self._leg_state != self._prev_leg_state
        if (self._counter % self._decimation == 0) or (torch.any(state_changed)):
            # Acquire all matrices necessary
            # self._compute_X()
            self._compute_Xd()
            # self._compute_centroidal_dynamics_matrices()
            # self._compute_qp_mats()
            # self._compute_objective_horizon_matrix()
            # G_ineq, h_ineq = self._compute_constraint_horizon_matrices()

            # # Convert to double data type
            # G_ineq = G_ineq.to(self.qp_dtype)
            # h_ineq = h_ineq.to(self.qp_dtype)
            # G = self.G.to(self.qp_dtype)
            # p = self.p.to(self.qp_dtype)

            # Aeq = torch.empty((self.num_envs, 0, self.nu*self.cfg.horizon_length), device=self.device, dtype=self.qp_dtype)
            # beq = torch.empty((self.num_envs, 0), device=self.device, dtype=self.qp_dtype)

            # f_w = QPFunction(check_Q_spd=False, eps=1e-5, maxIter=150)(G, p, G_ineq, h_ineq, Aeq, beq)  # (E, 12H)
            # f_w0 = f_w[:, :self.nu]  # (E,12) -> only need the first step (receding horizon MPC)
            # self.mpc_fw = f_w0.view(-1, self.nf, 3).to(torch.float32)  # (E,4,3)

        self.update_command()
        self._prev_leg_state = self._leg_state
        
        # # Jacobian for torques
        # Jw = self._compute_frame_jacobian()
        # Jw_T = Jw[:, :, 0:3, :].transpose(-1, -2) # [E, 4, J, 3]

        # # Get the desired joint positions from the command directly
        # q_desired = self.joint_pos_desired
        # tau_ff = self._compute_feedforward_torque(Jw[:, :, 0:3, :])

        # # Torques for joint positions from IK
        # q_meas = self._asset.data.joint_pos[:, self._joint_ids]
        # dq_meas = self._asset.data.joint_vel[:, self._joint_ids]
        # tau_swing = self.cfg.kp * (q_desired - q_meas) + self.cfg.kd * (-dq_meas) + tau_ff

        # # Torques from QP foot forces
        # tau_stance = -torch.einsum('elnd,eld->en', Jw_T, self.mpc_fw)
        
        # # Swing mask to decide which torques to take for which feet
        # mask = torch.cat([self.leg_state, self.leg_state, self.leg_state], dim=1)
        # self.tau = torch.where(
        #     mask,
        #     tau_swing,
        #     tau_stance,
        # )

        # # Set target to actuator compute() function
        # self._asset.set_joint_effort_target(self.tau, joint_ids=self._joint_ids)
        
        self._counter += 1

[arjun-2612]

Ok so I have incorporated what the example file for multi_env.py has in your repo with my own codebase. Below is the code for running the JAX based MPC from your repo.

"""
Pure JAX MPC for Husky-B using MPX's primal-dual iLQR optimizer.
No Mujoco dependency - all state comes from IsaacLab.
"""

import jax
import jax.numpy as jnp
from functools import partial

# Import MPX's optimizer
try:
    import mpx.primal_dual_ilqr.primal_dual_ilqr.optimizers as optimizers
except ImportError:
    raise ImportError(
        "MPX not found. Please install: pip install git+https://github.com/iit-DLSLab/mpx.git"
    )


# -- Configuration for Husky-B --

class HuskyBMPCConfig:
    """Configuration for Husky-B MPC."""
    
    # MPC parameters
    N = 12              # Horizon length
    dt = 0.02           # Timestep (50 Hz)
    n_contact = 4       # Number of feet
    n_x = 13            # State dimension (rpy, pos, omega, vel, -g)
    n_u = 12            # Control dimension (3 GRFs per foot)
    
    # Robot parameters
    mass = 6.11         # kg
    inertia = jnp.diag(jnp.array([0.05, 0.1, 0.1]))  # Body inertia (adjust based on URDF)
    inertia_inv = jnp.linalg.inv(inertia)
    
    # Cost weights (matching your RMPCActionCfg q_diag, r_diag)
    Q_diag = jnp.array([5.0, 40.0, 1.0, 5.0, 5.0, 50.0, 0.01, 0.1, 0.2, 10.0, 10.0, 1.5, 0.0])
    R_diag = jnp.tile(jnp.array([5e-4, 5e-4, 1e-3]), n_contact)
    
    # Constraints
    friction_coeff = 0.6
    f_min_ratio = 0.03
    f_max_ratio = 1.5


config = HuskyBMPCConfig()


# -- Dynamics model (centroidal SRBD) --

def huskyb_centroidal_dynamics(x, u, parameter):
    """
    Single-rigid-body centroidal dynamics for Husky-B.
    
    Args:
        x: (13,) = [roll, pitch, yaw, px, py, pz, wx, wy, wz, vx, vy, vz, -g]
        u: (12,) = [fx0, fy0, fz0, ..., fx3, fy3, fz3] (GRFs in world frame)
        parameter: (21,) auxiliary info:
            - [:12]: foot positions relative to COM (flatten of 4x3)
            - [12:21]: inertia inverse matrix (flatten of 3x3)
    
    Returns:
        x_next: (13,)
    """
    # Extract state
    roll, pitch, yaw = x[0], x[1], x[2]
    pos = x[3:6]
    omega = x[6:9]
    vel = x[9:12]
    g = x[12]  # -9.81
    
    # Reshape forces
    f = u.reshape((config.n_contact, 3))  # (4, 3)
    
    # Sum all forces
    f_total = jnp.sum(f, axis=0)  # (3,)
    
    # Linear acceleration: a = (sum_f) / m + [0, 0, g]
    acc = f_total / config.mass + jnp.array([0.0, 0.0, g])
    
    # Angular acceleration
    # parameter[:12] contains foot positions relative to COM (flatten of 4x3)
    # parameter[12:21] contains inertia inverse (flatten of 3x3)
    r = parameter[:12].reshape((config.n_contact, 3))  # (4, 3)
    inertia_inv = parameter[12:21].reshape((3, 3))  # (3, 3)
    
    # Torque = sum(r x f)
    torque = jnp.sum(jnp.cross(r, f), axis=0)  # (3,)
    
    # alpha = I^-1 @ torque (using dynamically computed inertia inverse)
    alpha = inertia_inv @ torque
    
    # Simple Euler integration (MPX uses ZOH for A,B; we match that spirit)
    # For proper match, would use your _compute_centroidal_dynamics_matrices logic
    # Here: simplified for speed
    
    # RPY rate (approximate, should use proper rotation matrix)
    # For small angles: rpy_dot ≈ omega in body frame rotated to world
    # Simplified: rpy_dot ≈ omega (valid for small pitch/roll)
    rpy_next = jnp.array([roll, pitch, yaw]) + omega * config.dt
    pos_next = pos + vel * config.dt
    omega_next = omega + alpha * config.dt
    vel_next = vel + acc * config.dt
    g_next = g  # constant
    
    x_next = jnp.concatenate([rpy_next, pos_next, omega_next, vel_next, jnp.array([g_next])])
    return x_next


# -- Cost function --

def huskyb_cost(W, reference, x, u, t):
    """
    Quadratic tracking cost with friction cone penalties.
    
    Args:
        W: (not used, kept for MPX API compatibility)
        reference: (N+1, n_x) desired state trajectory
        x: (n_x,) current state
        u: (n_u,) current control
        t: timestep index
    
    Returns:
        cost: scalar
    """
    x_ref = reference[t]
    x_err = x - x_ref
    
    # State tracking cost
    Q = jnp.diag(config.Q_diag)
    cost_x = 0.5 * x_err @ Q @ x_err
    
    # Control cost
    R = jnp.diag(config.R_diag)
    cost_u = 0.5 * u @ R @ u
    
    # Friction cone soft penalty (MPX often uses barrier or penalty)
    mu = config.friction_coeff
    for i in range(config.n_contact):
        fx, fy, fz = u[3*i], u[3*i+1], u[3*i+2]
        # Cone: mu*fz >= sqrt(fx^2 + fy^2)
        lateral = jnp.sqrt(fx**2 + fy**2 + 1e-6)
        cone_margin = mu * fz - lateral
        # Penalty if violated
        cost_u += jnp.where(cone_margin < 0, -1e3 * cone_margin, 0.0)
        
        # Force bounds (soft)
        f_min = config.f_min_ratio * config.mass * 9.81 / config.n_contact
        f_max = config.f_max_ratio * config.mass * 9.81
        cost_u += jnp.where(fz < f_min, 1e2 * (f_min - fz)**2, 0.0)
        cost_u += jnp.where(fz > f_max, 1e2 * (fz - f_max)**2, 0.0)
    
    return cost_x + cost_u


# -- Hessian approximation (Gauss-Newton) --

def huskyb_hessian_gn(x, u, t):
    """
    Gauss-Newton Hessian approximation for the cost.
    
    Returns:
        Hxx: (n_x, n_x)
        Huu: (n_u, n_u)
        Hxu: (n_x, n_u)
    """
    Hxx = jnp.diag(config.Q_diag)
    Huu = jnp.diag(config.R_diag)
    Hxu = jnp.zeros((config.n_x, config.n_u))
    return Hxx, Huu, Hxu


# -- Build MPC solver --

# Partial application of cost/dynamics/hessian
_cost_fn = partial(huskyb_cost, None)  # W not used
_dynamics_fn = huskyb_centroidal_dynamics
_hessian_fn = huskyb_hessian_gn

# Build MPC work function
_mpc_work = partial(optimizers.mpc, _cost_fn, _dynamics_fn, _hessian_fn, False)

# JIT-compile batched version
def _solve_one(reference, parameter, W, x0, X0, U0, V0):
    """Single-env MPC solve."""
    X, U, V = _mpc_work(reference, parameter, W, x0, X0, U0, V0)
    return X, U, V

solve_batched = jax.jit(jax.vmap(_solve_one, in_axes=(0, 0, 0, 0, 0, 0, 0)))


# -- Public API --

def init_warm_start(num_envs):
    """
    Initialize warm-start trajectories for MPC.
    
    Args:
        num_envs: number of parallel environments
    
    Returns:
        X0: (E, N+1, n_x) initial state trajectory
        U0: (E, N, n_u) initial control trajectory
        V0: (E, N+1, n_x) initial dual variables
    """
    X0 = jnp.zeros((num_envs, config.N + 1, config.n_x))
    U0 = jnp.zeros((num_envs, config.N, config.n_u))
    V0 = jnp.zeros((num_envs, config.N + 1, config.n_x))
    return X0, U0, V0


def mpc_step(x0, reference, parameter, X0, U0, V0):
    """
    Run batched MPC update.
    
    Args:
        x0: (E, n_x) current state
        reference: (E, N+1, n_x) desired trajectory
        parameter: (E, N, n_param) auxiliary info (foot positions relative to COM per timestep)
        X0, U0, V0: (E, N+1, n_x), (E, N, n_u), (E, N+1, n_x) warm starts
    
    Returns:
        X: (E, N+1, n_x) optimized state trajectory
        U: (E, N, n_u) optimized control trajectory
        V: (E, N+1, n_x) dual variables
        u0: (E, n_u) first control (foot forces for this step)
    """
    num_envs = x0.shape[0]
    W = jnp.tile(jnp.eye(config.n_x + config.n_u), (num_envs, 1, 1))  # dummy weight matrix
    
    X, U, V = solve_batched(reference, parameter, W, x0, X0, U0, V0)
    
    u0 = U[:, 0, :]  # First control step (E, 12) - foot forces
    return X, U, V, u0


def shift_trajectories(X, U, V):
    """
    Warm-start for next MPC iteration: shift trajectories forward in time.
    
    Args:
        X, U, V: current solution trajectories
    
    Returns:
        X0, U0, V0: shifted trajectories for next iteration
    """
    # Shift and repeat last element
    X0 = jnp.concatenate([X[:, 1:, :], X[:, -1:, :]], axis=1)
    U0 = jnp.concatenate([U[:, 1:, :], U[:, -1:, :]], axis=1)
    V0 = jnp.concatenate([V[:, 1:, :], V[:, -1:, :]], axis=1)
    return X0, U0, V0

and here is how I am using it in my action config in IsaacLab

def apply_actions(self):
        state_changed = self._leg_state != self._prev_leg_state
        if (self._counter % self._decimation == 0) or (torch.any(state_changed)):
            # Use MPX JAX MPC
            self._compute_X()
            self._compute_Xd()

            # Build state x0 (E, 13): [rpy, pos, omega, vel, -g]
            x0_torch = self.xt  # (E, 13)

            # Build reference trajectory (E, N+1, 13)
            reference_torch = self.xd  # (E, N, 13) -> need (E, N+1, 13)
            # Append last desired state to match N+1
            reference_torch = torch.cat(
                [reference_torch, reference_torch[:, -1:, :]], dim=1
            )  # (E, N+1, 13)

            # Compute composite inertia inverse (E, 3, 3)
            inertia_inv_torch = self._compute_composite_inertia_world()  # (E, 3, 3)
            # Flatten to (E, 9) for passing through parameter
            inertia_inv_flat = inertia_inv_torch.flatten(start_dim=1)  # (E, 9)

            # Build parameter: foot positions relative to COM + inertia inverse
            # For simplicity, repeat current foot positions for all timesteps
            com_pos_w = self._asset.data.root_com_pos_w.clone()  # (E, 3)
            foot_pos_w = self.foot_pos_w.clone()  # (E, 4, 3)
            foot_rel = (foot_pos_w - com_pos_w.unsqueeze(1)).flatten(start_dim=1)  # (E, 12)
            # Concatenate with inertia: (E, 12 + 9) = (E, 21) per timestep
            param_single = torch.cat([foot_rel, inertia_inv_flat], dim=1)  # (E, 21)
            # Tile for horizon: (E, N, 21)
            parameter_torch = param_single.unsqueeze(1).repeat(1, huskyb_mpc_jax.config.N, 1)

            # Convert to numpy for JAX
            x0_np = x0_torch.detach().cpu().numpy()
            reference_np = reference_torch.detach().cpu().numpy()
            parameter_np = parameter_torch.detach().cpu().numpy()

            # Call MPX MPC
            X, U, V, u0_np = huskyb_mpc_jax.mpc_step(
                x0_np, reference_np, parameter_np,
                self._mpx_X0, self._mpx_U0, self._mpx_V0
            )

            # Shift trajectories for next warm-start
            self._mpx_X0, self._mpx_U0, self._mpx_V0 = huskyb_mpc_jax.shift_trajectories(X, U, V)

            # Convert back to torch: u0_np is (E, 12) foot forces
            self.mpc_fw = torch.from_numpy(u0_np).to(self.device).view(-1, self.nf, 3)  # (E, 4, 3)

Have I done this right? I haven't yet tested this out, but will do so later this week and let you know!

[lorenzo96-cmd]

Hi @arjun-2612,
I noticed a couple of odd things in the code snippet you shared.

• You have-gin your state, is that gravity? If so, it should be a parameter, not an optimization variable.

• I believe the integration of the Euler angles is incorrect.

• The Hessian approximation seems to be dropping the barrier term completely.

If you want to use the SRBD model, you can use the one in mpx and just change the mass and inertia in the config.

I can come by later next week with more suggestions.

@LockedFlysher , could you share a snippet of your code? If you have a clean working example in IsaacLab, we could include it in the examples list. If you’re interested, feel free to open a PR.

[arjun-2612]

Hey @lorenzo96-cmd

Yeah I include gravity as a 13th state variable. The MIT cheetah paper does this to compensate for gravity with the foot forces.

Oh, I see. I will check with your code and see how you do it then.

Oh yes, I will need to complete that piece of the implementation.

I am computing the inertia of each link in the robot in IsaacLab, and aggregating them together, hence I pass them into the code as a parameter.

Oh yeah if there is already some code written that would be great to look at. Please let me know!

[lorenzo96-cmd]

I’m familiar with the paper you’re referring to and also in their case, g is a constant not an optimisation variable. You’re not seeking the optimal gravity.

[arjun-2612]

Oh I see. Ok I have made the changes to reflect what your API needs. Here is my updated code.

MPX code file

"""
Pure JAX MPC for Husky-B using MPX's primal-dual iLQR optimizer.
Uses exact implementations copied from MPX repo for compatibility.
Uses dlpack for zero-copy GPU tensor conversion between PyTorch and JAX.
"""

import os
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'

import jax
jax.config.update('jax_platforms', 'cuda')

import jax.scipy as jsp
import jax.numpy as jnp
from functools import partial
from mujoco.mjx._src import math as mjx_math

# dlpack conversion utilities
from jax.dlpack import from_dlpack as jax_from_dlpack
from jax.dlpack import to_dlpack as jax_to_dlpack

# Import MPX's optimizer
try:
    import mpx.primal_dual_ilqr.primal_dual_ilqr.optimizers as optimizers
except ImportError:
    raise ImportError(
        "MPX not found. Please install: pip install git+https://github.com/iit-DLSLab/mpx.git"
    )


# -- Configuration for Husky-B --

class HuskyBMPCConfig:
    """Configuration for Husky-B MPC."""
    
    # MPC parameters
    N = 12              # Horizon length
    dt = 0.02           # Timestep (50 Hz)
    n_contact = 4       # Number of feet
    n_x = 13            # State dimension (pos, quat, vel, omega)
    n_u = 12            # Control dimension (3 GRFs per foot)
    
    # Robot parameters
    mass = 6.11         # kg
    _inertia_diag = jnp.array([0.05, 0.1, 0.1])  # Body inertia diagonal
    gravity = 9.81      # m/s^2
    
    # Cost weights (matching your RMPCActionCfg q_diag, r_diag)
    # Note: Qrot is 3D even though quat residual is 4D - this is how MPX does it
    Q_diag = jnp.array([
        5.0, 5.0, 50.0,      # position (3)
        5.0, 40.0, 1.0,      # rotation (3) - NOT 4!
        10.0, 10.0, 1.5,     # linear velocity (3)
        0.01, 0.1, 0.2       # angular velocity (3)
    ])
    R_diag = jnp.tile(jnp.array([5e-4, 5e-4, 1e-3]), n_contact)
    
    friction_coeff = 0.7
    
    @classmethod
    def get_inertia(cls):
        return jnp.diag(cls._inertia_diag)
    
    @classmethod
    def get_inertia_inv(cls):
        return jnp.linalg.inv(jnp.diag(cls._inertia_diag))


config = HuskyBMPCConfig()


# -- EXACT COPIES FROM MPX REPO --

def penalty(constraint, alpha=0.1, sigma=5):
    """Log barrier penalty (exact from MPX)."""
    def safe_log(x):
        x = jnp.clip(x, 1e-10, 1e6)
        return jnp.log(x)
    quadratic_barrier = alpha / 2 * (jnp.square((constraint - 2 * sigma) / sigma) - jnp.ones_like(constraint))
    log_barrier = -alpha * safe_log(constraint)
    return jnp.clip(jnp.where(constraint > sigma, log_barrier, quadratic_barrier + log_barrier), 0, 1e8)


def quadruped_srbd_dynamics(mass, inertia, inertia_inv, dt, x, u, t, parameter):
    """Exact from MPX."""
    p = x[:3]
    quat = x[3:7]
    dp = x[7:10]
    omega = x[10:13]
    grf = u

    contact = parameter[t, :4]
    p_legs = parameter[t, 4:].reshape((4, 3))

    dp_next = dp + (jnp.array([0, 0, -9.81]) + (1 / mass) * (
        grf[:3] * contact[0] + grf[3:6] * contact[1] + 
        grf[6:9] * contact[2] + grf[9:12] * contact[3]
    )) * dt

    p0 = p_legs[0]
    p1 = p_legs[1]
    p2 = p_legs[2]
    p3 = p_legs[3]

    b_R_w = mjx_math.quat_to_mat(quat)
    base_angular_wrench = (
        jnp.cross(p0 - p, grf[:3]) * contact[0] + 
        jnp.cross(p1 - p, grf[3:6]) * contact[1] + 
        jnp.cross(p2 - p, grf[6:9]) * contact[2] + 
        jnp.cross(p3 - p, grf[9:12]) * contact[3]
    )
    omega_next = omega + inertia_inv @ (b_R_w.T @ base_angular_wrench - jnp.cross(omega, inertia @ omega)) * dt

    p_new = p + dp_next * dt
    quat_new = mjx_math.quat_integrate(quat, omega_next, dt)

    x_next = jnp.concatenate([p_new, quat_new, dp_next, omega_next])
    return x_next


def quadruped_srbd_obj(n_contact, N, W, reference, x, u, t):
    """Exact from MPX."""
    p = x[:3]
    quat = x[3:7]
    dp = x[7:10]
    omega = x[10:13]
    grf = u

    p_ref = reference[t, :3]
    quat_ref = reference[t, 3:7]
    dp_ref = reference[t, 7:10]
    omega_ref = reference[t, 10:13]

    mu = 0.7
    friction_cone = mu * grf[2::3] - jnp.sqrt(
        jnp.square(grf[1::3]) + jnp.square(grf[::3]) + jnp.ones(n_contact) * 1e-2
    )
    friction_cone = jnp.clip(penalty(friction_cone), 1e-6, 1e6)
    
    stage_cost = (
        (p - p_ref).T @ W[:3, :3] @ (p - p_ref) + 
        mjx_math.quat_sub(quat, quat_ref).T @ W[3:6, 3:6] @ mjx_math.quat_sub(quat, quat_ref) +
        (dp - dp_ref).T @ W[6:9, 6:9] @ (dp - dp_ref) + 
        (omega - omega_ref).T @ W[9:12, 9:12] @ (omega - omega_ref) +
        grf.T @ W[12:12 + 3 * n_contact, 12:12 + 3 * n_contact] @ grf + 
        jnp.sum(friction_cone)
    )
    term_cost = (
        (p - p_ref).T @ W[:3, :3] @ (p - p_ref) + 
        mjx_math.quat_sub(quat, quat_ref).T @ W[3:6, 3:6] @ mjx_math.quat_sub(quat, quat_ref) +
        (dp - dp_ref).T @ W[6:9, 6:9] @ (dp - dp_ref) + 
        (omega - omega_ref).T @ W[9:12, 9:12] @ (omega - omega_ref)
    )

    return jnp.where(t == N, 0.5 * term_cost, 0.5 * stage_cost)


def quadruped_srbd_hessian_gn(n_contact, W, reference, x, u, t, parameter):
    """Gauss-Newton Hessian matching MPX exactly. Parameter is passed from closure via partial."""
    
    def residual(x_in, u_in):
        """Residual vector for Gauss-Newton (25D with full 4D quat)."""
        p = x_in[:3]
        quat = x_in[3:7]
        dp = x_in[7:10]
        omega = x_in[10:13]
        grf = u_in

        p_ref = reference[t, :3]
        quat_ref = reference[t, 3:7]
        dp_ref = reference[t, 7:10]
        omega_ref = reference[t, 10:13]

        p_res = (p - p_ref).T
        quat_res = mjx_math.quat_sub(quat, quat_ref).T  # Full 4D
        dp_res = (dp - dp_ref).T
        omega_res = (omega - omega_ref).T
        grf_res = grf.T

        return jnp.concatenate([p_res, quat_res, dp_res, omega_res, grf_res])
    
    def friction_constraint(u_in):
        """Friction cone constraint."""
        grf = u_in
        mu = 0.7
        friction_cone = mu * grf[2::3] - jnp.sqrt(
            jnp.square(grf[1::3]) + jnp.square(grf[::3]) + jnp.ones(n_contact) * 1e-2
        )
        return friction_cone
    
    # Create jacobian functions (not call them yet)
    J_x_fn = jax.jacobian(residual, 0)
    J_u_fn = jax.jacobian(residual, 1)
    
    contact = parameter[t, :4]  # Contact comes from parameter, not reference
    hessian_penalty = jax.grad(jax.grad(penalty))
    J_friction_cone_fn = jax.jacobian(friction_constraint)
    
    # Now call the jacobian functions with (x, u)
    H_penalty = jnp.diag(
        jnp.clip(jax.vmap(hessian_penalty)(friction_constraint(u)), 1e-6, 1e6) * contact
    )
    H_constraint = J_friction_cone_fn(u).T @ H_penalty @ J_friction_cone_fn(u)
    
    return J_x_fn(x, u).T @ W @ J_x_fn(x, u), J_u_fn(x, u).T @ W @ J_u_fn(x, u) + H_constraint, J_x_fn(x, u).T @ W @ J_u_fn(x, u)


# -- Build MPC solver --

def _solve_one(reference, parameter, W, x0, X0, U0, V0):
    """Single-env MPC solve using partials matching MPX wrapper.
    
    Args:
        reference: (N+1, 13) state trajectory only
        parameter: (N, 16) contact_schedule + foot_positions per timestep
        W: (24, 24) weight matrix
        x0, X0, U0, V0: initial state and warm starts
    """
    # Create partials matching MPX's mpc_wrapper_srbd.py
    mass = config.mass
    inertia = config.get_inertia()
    inertia_inv = config.get_inertia_inv()
    dt = config.dt
    n_contact = config.n_contact
    N = config.N
    
    dynamics_fn = partial(quadruped_srbd_dynamics, mass, inertia, inertia_inv, dt)
    cost_fn = partial(quadruped_srbd_obj, n_contact, N)
    hess_fn = partial(quadruped_srbd_hessian_gn, n_contact, parameter=parameter)  # Bind parameter from closure
    
    X, U, V = optimizers.mpc(
        cost_fn,
        dynamics_fn,
        hess_fn,
        False,
        reference,
        parameter,
        W,
        x0,
        X0,
        U0,
        V0,
    )
    
    return X, U, V


# Batched solver
solve_batched = jax.jit(
    jax.vmap(_solve_one, in_axes=(0, 0, None, 0, 0, 0, 0))
)

def init_warm_start(num_envs):
    """
    Initialize warm-start trajectories for MPC on GPU.
    
    Args:
        num_envs: number of parallel environments
    
    Returns:
        X0: (E, N+1, n_x) initial state trajectory (JAX array on GPU)
        U0: (E, N, n_u) initial control trajectory (JAX array on GPU)
        V0: (E, N+1, n_x) initial dual variables (JAX array on GPU)
    """
    X0 = jnp.zeros((num_envs, config.N + 1, config.n_x))
    U0 = jnp.zeros((num_envs, config.N, config.n_u))
    V0 = jnp.zeros((num_envs, config.N + 1, config.n_x))
    return X0, U0, V0


def mpc_step(x0, reference, parameter, X0, U0, V0):
    """
    Run batched MPC update with zero-copy GPU tensor conversion.
    
    Args:
        x0: (E, 13) PyTorch tensor - current state [pos, quat, vel, omega]
        reference: (E, N+1, 13) PyTorch tensor - state trajectory only
        parameter: (E, N, 16) PyTorch tensor - contact + foot positions per timestep
        X0, U0, V0: JAX arrays on GPU - warm starts
    
    Returns:
        X, U, V: optimized trajectories (JAX arrays)
        u0_torch: (E, 12) first control as PyTorch tensor
    """
    from torch.utils.dlpack import from_dlpack as torch_from_dlpack
    
    # Debug input sanity
    print(f"[MPC ENTRY] x0: shape={x0.shape}, nan={x0.isnan().any()}, max={x0.abs().max()}")
    print(f"[MPC ENTRY] reference: shape={reference.shape}, nan={reference.isnan().any()}, max={reference.abs().max()}")
    print(f"[MPC ENTRY] parameter: shape={parameter.shape}, nan={parameter.isnan().any()}, max={parameter.abs().max()}")
    
    # Convert to JAX (zero-copy) - automatically on GPU via dlpack
    x0_jax = jax_from_dlpack(x0.contiguous())
    reference_jax = jax_from_dlpack(reference.contiguous())
    parameter_jax = jax_from_dlpack(parameter.contiguous())
    
    num_envs = x0_jax.shape[0]
    
    # Ensure we're on GPU device
    gpu_device = jax.devices('gpu')[0] if jax.devices('gpu') else jax.devices('cpu')[0]
    
    # Build weight matrix (block-diagonal Q, R) matching MPX
    # Residual is 25D: p(3) + quat(4) + dp(3) + omega(3) + grf(12)
    # But W is 24D: p(3) + quat(3) + dp(3) + omega(3) + grf(12)
    # The 4D quat residual uses a 3D weight matrix (this is MPX's design)
    Qp = jnp.diag(config.Q_diag[:3])
    Qrot = jnp.diag(config.Q_diag[3:6])
    Qdp = jnp.diag(config.Q_diag[6:9])
    Qomega = jnp.diag(config.Q_diag[9:12])
    R = jnp.diag(config.R_diag)
    
    W = jsp.linalg.block_diag(Qp, Qrot, Qdp, Qomega, R)
    
    # Place W on GPU
    with jax.default_device(gpu_device):
        W = jax.device_put(W, gpu_device)
        X0_gpu = jax.device_put(X0, gpu_device)
        U0_gpu = jax.device_put(U0, gpu_device)
        V0_gpu = jax.device_put(V0, gpu_device)
        
        # Call batched solver - W is NOT vmapped (in_axes=None), so same for all envs
        X, U, V = solve_batched(reference_jax, parameter_jax, W, x0_jax, X0_gpu, U0_gpu, V0_gpu)
    
    print(f"[MPC EXIT] U: shape={U.shape}, nan={jnp.isnan(U).any()}, max={jnp.abs(U).max()}")
    print(f"[MPC EXIT] X: shape={X.shape}, nan={jnp.isnan(X).any()}, max={jnp.abs(X).max()}")
    
    # Convert first control back to PyTorch
    u0_jax = U[:, 0, :]
    u0_torch = torch_from_dlpack(u0_jax)
    
    return X, U, V, u0_torch


def shift_trajectories(X, U, V):
    """
    Warm-start for next MPC iteration: shift trajectories forward in time.
    
    Args:
        X, U, V: current solution trajectories
    
    Returns:
        X0, U0, V0: shifted trajectories for next iteration
    """
    # Shift and repeat last element
    X0 = jnp.concatenate([X[:, 1:, :], X[:, -1:, :]], axis=1)
    U0 = jnp.concatenate([U[:, 1:, :], U[:, -1:, :]], axis=1)
    V0 = jnp.concatenate([V[:, 1:, :], V[:, -1:, :]], axis=1)
    return X0, U0, V0

The action in IsaacLab

def apply_actions(self):
        state_changed = self._leg_state != self._prev_leg_state
        if (self._counter % self._decimation == 0) or (torch.any(state_changed)):
            # New MPX implementation
            self._compute_X()
            self._compute_Xd()

            # MPX expects: (E, N, n_contact, 3)
            com_pos_w = self._asset.data.root_com_pos_w.clone()  # (E, 3)
            foot_pos_w = self.foot_targets_w.clone()  # (E, 4, 3)
            foot_rel_com = foot_pos_w - com_pos_w.unsqueeze(1)  # (E, 4, 3)
            
            # Build parameter tensor: (E, N, 4+12) = [contact_schedule(4), foot_positions(12)]
            # For now, assume all contacts active
            contact_schedule = torch.ones((self.num_envs, huskyb_mpc_jax.config.N, 4), device=self.device)
            foot_pos_flat = foot_rel_com.unsqueeze(1).expand(-1, huskyb_mpc_jax.config.N, -1, -1)  # (E, N, 4, 3)
            foot_pos_flat = foot_pos_flat.reshape(self.num_envs, huskyb_mpc_jax.config.N, -1)  # (E, N, 12)
            parameter_torch = torch.cat([contact_schedule, foot_pos_flat], dim=-1)  # (E, N, 16)
            
            # Build reference tensor: (E, N+1, 13) = [state_trajectory only]
            # Contact and foot positions go in parameter tensor, NOT reference
            reference_torch = self.xd  # (E, N+1, 13)
            
            # Debug: Check tensor shapes and values before MPC
            print(f"[DEBUG MPC Input]")
            print(f"  xt shape: {self.xt.shape}, isnan: {torch.isnan(self.xt).any()}, vals: {self.xt[0]}")
            print(f"  reference shape: {reference_torch.shape}, isnan: {torch.isnan(reference_torch).any()}")
            print(f"  reference[0,0,:13]: {reference_torch[0, 0, :13]}")
            print(f"  reference[0,-1,:13]: {reference_torch[0, -1, :13]}")
            print(f"  parameter shape: {parameter_torch.shape}, isnan: {torch.isnan(parameter_torch).any()}")
            
            X, U, V, u0_torch = huskyb_mpc_jax.mpc_step(
                self.xt, reference_torch, parameter_torch,
                self._mpx_X0, self._mpx_U0, self._mpx_V0
            )
            
            print(f"[DEBUG MPC Output]")
            print(f"  U shape: {U.shape}, isnan: {jnp.isnan(U).any()}")
            print(f"  u0_torch shape: {u0_torch.shape}, isnan: {torch.isnan(u0_torch).any()}, vals: {u0_torch[0]}")
            
            self._mpx_X0, self._mpx_U0, self._mpx_V0 = huskyb_mpc_jax.shift_trajectories(X, U, V)
            self.mpc_fw = u0_torch.reshape(self.num_envs, 4, 3)

This code here runs fast but gives NaNs everytime for foot forces. Here is a log of what I see.

[DEBUG MPC Input]
  xt shape: torch.Size([128, 13]), isnan: False, vals: tensor([ 1.4271e+01, -1.2849e+01,  4.8000e-01,  9.5497e-01,  1.0430e-02,
        -2.6833e-02,  2.9530e-01,  1.6293e-01, -1.4361e-01,  0.0000e+00,
        -3.5810e-02,  5.0996e-02, -1.4339e-01], device='cuda:0')
  reference shape: torch.Size([128, 13, 13]), isnan: False
  reference[0,0,:13]: tensor([ 14.2707, -12.8485,   0.4700,   0.9554,   0.0000,   0.0000,   0.2954,
          0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000],
       device='cuda:0')
  reference[0,-1,:13]: tensor([ 14.2707, -12.8485,   0.4700,   0.9554,   0.0000,   0.0000,   0.2954,
          0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000],
       device='cuda:0')
  parameter shape: torch.Size([128, 12, 16]), isnan: False
[MPC ENTRY] x0: shape=torch.Size([128, 13]), nan=False, max=14.27074146270752
[MPC ENTRY] reference: shape=torch.Size([128, 13, 13]), nan=False, max=14.27074146270752
[MPC ENTRY] parameter: shape=torch.Size([128, 12, 16]), nan=False, max=1.0
[MPC EXIT] U: shape=(128, 12, 12), nan=True, max=nan
[MPC EXIT] X: shape=(128, 13, 13), nan=True, max=nan
[DEBUG MPC Output]
  U shape: (128, 12, 12), isnan: True
  u0_torch shape: torch.Size([128, 12]), isnan: True, vals: tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
       device='cuda:0')

Not really sure what's going on because I am not very versed in JAX API, so hoping you can point out my issue.