Slow compile for new env

Hi,

I'm trying to add a new environment. The model is a manipulator with 6 joints. I always run into slow compilation issue, and I have tried everything to reduce the complexity of the model.

```
E1109 19:24:47.543532 2695810 slow_operation_alarm.cc:73]
********************************
[Compiling module jit_generate_eval_unroll for GPU] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
********************************
E1109 19:28:30.719801 2695453 slow_operation_alarm.cc:140] The operation took 5m43.176337956s

********************************
[Compiling module jit_generate_eval_unroll for GPU] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
********************************
```

Could you take a look at [this PR](https://github.com/thuaduc/mujoco_playground/pull/1) and helps me find the issue? I really appreciate any helps.

Best,
Duc