Using proxddp with MuJoCo's dynamics

[tomekatat]

Hi everyone,

I am trying to use proxDDP with MuJoCo's dynamics. I defined a discrete dynamics class for aligator in the following way:

import aligator
import numpy as np
from aligator import constraints, manifolds, dynamics
import mujoco as mj

class DiscDynamicModel(dynamics.ExplicitDynamicsModel):
    shared_model = None
    shared_data = None
    shared_timestep = None

    def __init__(self):
        self.mj_model: mj.MjModel = DiscDynamicModel.shared_model
        self.mj_data: mj.MjData = DiscDynamicModel.shared_data
        self.nx = 2 * self.mj_model.nv
        self.NU = self.mj_model.nu
        self.nq = self.mj_model.nq
        self.nv = self.mj_model.nv

        space = manifolds.VectorSpace(self.nx)
        self.dt = DiscDynamicModel.shared_timestep

        super().__init__(space, self.NU)

    def __getinitargs__(self):
        return ()

    def __deepcopy__(self, memo):
        # Return self or a lightweight new object (without deep copying model/data)
        cls = self.__class__
        new_obj = cls.__new__(cls)
        memo[id(self)] = new_obj

        # Only copy the minimal state
        new_obj.mj_model = self.mj_model
        new_obj.mj_data = self.mj_data
        new_obj.dt = self.dt
        new_obj.nx = self.nx
        new_obj.NU = self.NU
        new_obj.nq = self.nq
        new_obj.nv = self.nv

        dynamics.ExplicitDynamicsModel.__init__(new_obj, self.space, self.NU)
        return new_obj

    def forward(self, x, u, data: aligator.dynamics.ExplicitDynamicsData):
        self.mj_data.qpos = x[:self.nv].copy()
        self.mj_data.qvel = x[self.nv:].copy()
        self.mj_data.ctrl = u.copy()

        mj.mj_step(self.mj_model, self.mj_data)

        x_next = np.hstack((self.mj_data.qpos, self.mj_data.qvel))
        data.xnext[:] = x_next.copy()

    def dForward(self, x, u, data: aligator.dynamics.ExplicitDynamicsData):
        self.mj_data.qpos = x[:self.nv].copy()
        self.mj_data.qvel = x[self.nv:].copy()
        self.mj_data.ctrl = u.copy()

        A = np.zeros((self.nx, self.nx))
        B = np.zeros((self.nx, self.NU))

        mj.mj_step(self.mj_model, self.mj_data)
        mj.mjd_transitionFD(self.mj_model, self.mj_data, 1e-9, False, A, B, None, None)

        data.Jx[:,:] = A.copy()
        if self.NU == 1:
            data.Ju[:] = B[:,0].copy()
        else:
            data.Ju[:,:] = B.copy()

This way, I have run successfully proxddp, however it was really slow. I redefined the deepcopy function not to copy the MuJoCo model and data at every call. This helped, however it is still not enough for my application. My question is, that is there any way to make the solver faster or to insert MuJoCo dynamics into proxDDP in a more effective way?

Thanks for any help!