Nan issues when trying to train robot with extendable boom arms

[S-A-M-J]

Hi, for my master thesis i'm trying to train a robot with 4 extendable arms to traverse a random cave like environment.
The robot is based on the basic/reachbot.xml with each arm being an basic/reachbotArm.xml. The cave is represented by the cave_boxes.xml (about 4000 static boxes in an xml) and everything is inserted into the scene.xml. All files can be found in the repo here as i did not manage to attach them here somehow.

I am now trying to train the robot to move through the cave using brax (in turn using MJX) based on the mujoco playground examples with the code from cave_exploration.py and run_cave_exploration.py.
However it seems that whatever i do unless i scale the actions of the robot or its PD values so far down that it basically cant move any kind of training will always eventually lead to nan values appearing, but usually at different times.

For training i am using the run_cave_exploration.py file from the repo under tasks/cave_exploration

Things i have tried:

• setting the matmul precision to high
• debugging the nan but it only shows me that the issue happens within the training process.
• running only on cpu but it takes too long before an error appears.
• changing the integrator from euler to implicit fast.
• reducing the timestep up to 0.0005.
• reducing the length of the slide joints from 1.7 to 0.5
• reducing kp, kd value -> only works if so low that robot can't move
• reducing action scale -> only works if so low that robot can't move

Some of these seem to somewhat help to delay the appearance of the nan value but none has been able to make the simulation stable or even allow it to run for 10Mio steps. The problem however seems to be related to the physics being unstable since if the robot can barely move the training will at least run through the steps. Therefore im wondering if what i'm trying is just not possible due to the geometry or something else or anybody has any other ideas what i could try to stabilize the simulation or where the nan values are coming from.

All of the code can also be found here

I would really appreciate any help since i am on the brink of desperation.

The hardware i am using for training is an A100.

[S-A-M-J]

The answer was that due to solref and solimp settings geoms were clipping through each other and this was most likely resulting in infinite forces during training leading to nan values.