Python bindings - huge performance overhead

[davidhozic]

Intro

Hi!

I am a graduate student, I use MuJoCo for a reinforcement learning.

My setup

• Mujoco version (PIP): 3.2.3
• API: Python bindings
• Architecture: amd64
• OS: Kubuntu 24.04

My question

Maybe I'm using Mujoco bindings incorrectly, but I have a problem with huge overheads when iterating the MjData.contact array. My RL environment is setup to have a timestep of 20 ms, which I internally divide to 2 ms sub-steps - I call mj_step 10 times in a loop. After each mj_step call, I iterate the MjData.contact array and check the IDs of the contact geoms. I am able to do about 720 RL steps in one real-world second when doing this, however the performance doubles to 1400 RL steps per real-world second if I comment out the collision reading.

Is there some faster way I can access the collision data, or the data struct in general?

Minimal model and/or code that explain my question

def _updateCollisionNormals(self) -> list[float]:
        collisions = self.collisions
        for contact in self.data.contact:
            for rod_i, (min_i, max_i) in enumerate(self.collision_map):
                pidx = contact.geom[1]
                if min_i <= pidx <= max_i:
                    break  # TODO

720 steps in one second:

def stepSimulation(self):
        self.clearCollisionNormals()

        # Step simulation
        for i in range(self.internal_step_factor):
            mj.mj_step(MODEL, self.data)
            self._updateCollisionNormals()

...

1400 steps in one second:

def stepSimulation(self):
        self.clearCollisionNormals()

        # Step simulation
        for i in range(self.internal_step_factor):
            mj.mj_step(MODEL, self.data)
            # self._updateCollisionNormals()

...

Confirmations

• I searched the latest documentation thoroughly before posting.
• I searched previous Issues and Discussions, I am certain this has not been raised before.