Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Running out of GPU Memory #2236

Open
2 tasks done
KamatMayur opened this issue Nov 20, 2024 · 2 comments
Open
2 tasks done

Running out of GPU Memory #2236

KamatMayur opened this issue Nov 20, 2024 · 2 comments
Labels
MJX Using JAX to run on GPU

Comments

@KamatMayur
Copy link

KamatMayur commented Nov 20, 2024 •
edited
Loading

Intro

Hi!

I am a student, I use MuJoCo for my research on RL.

My setup

mujoco version: 3.2.5
python api
64 bit
Ubuntu 24.04.1 LTS
RTX 2060 super, 8GB @ 2010 MHz

What's happening? What did you expect?

Running the same code for MuJoCo MJX Humanoid environment from the collab tutorial, but with my own humanoid model, gives me the following error

/home/mayur-kamat/anaconda3/envs/rl/lib/python3.12/site-packages/jax/_src/interpreters/xla.py:133: RuntimeWarning: overflow encountered in cast
  return np.asarray(x, dtypes.canonicalize_dtype(x.dtype))
2024-11-20 20:37:20.877478: W external/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc:3020] Can't reduce memory use below 2.12GiB (2278089352 bytes) by rematerialization; only reduced to 5.65GiB (6063398400 bytes), down from 5.70GiB (6125699228 bytes) originally
E1120 20:37:33.901627  165263 hlo_lexer.cc:443] Failed to parse int literal: 894515288310727292233
/home/mayur-kamat/anaconda3/envs/rl/lib/python3.12/site-packages/jax/_src/interpreters/xla.py:133: RuntimeWarning: overflow encountered in cast
  return np.asarray(x, dtypes.canonicalize_dtype(x.dtype))
/home/mayur-kamat/anaconda3/envs/rl/lib/python3.12/site-packages/jax/_src/interpreters/xla.py:133: RuntimeWarning: overflow encountered in cast
  return np.asarray(x, dtypes.canonicalize_dtype(x.dtype))
/home/mayur-kamat/anaconda3/envs/rl/lib/python3.12/site-packages/jax/_src/interpreters/xla.py:133: RuntimeWarning: overflow encountered in cast
  return np.asarray(x, dtypes.canonicalize_dtype(x.dtype))
2024-11-20 20:40:31.260133: W external/xla/xla/tsl/framework/bfc_allocator.cc:306] Allocator (GPU_0_bfc) ran out of memory trying to allocate 1015.14MiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.
XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 1064452096 bytes.

This happens on my Local Machine as well as the Collab T4 gpu. I would like to know what the issue is and how it can be resolved. Halving the environment count didn't solve it either. Besides I manually measure the size of the batches environments which was barely 300 MBs of data.

Steps for reproduction

The humanoid model I use is given below. Please use this model in the humanoid environment given in the colab tutorial to replicate the results.

Minimal model for reproduction

minimal XML 
<mujoco>
    <option gravity="0 0 -9.8" />
    <default class="main">
        <joint limited="true" solimplimit="0 0.99 0.01" stiffness="0.1" armature=".01" damping="1"/>
        <geom friction="0.8" solref="0.015 1" solimp="0.95 0.99 0.003"/>
        <motor ctrllimited="true" ctrlrange="-1 1"/>
        <default class="humanoid">
          <geom type="capsule" material="self"/>
          <default class="stiff_low">
            <joint stiffness=".5" damping="4"/>
          </default>
          <default class="stiff_medium">
            <joint stiffness="10" damping="5"/>
          </default>
          <default class="stiff_high">
            <joint stiffness="30" damping="10"/>
          </default>
          <default class="touch">
            <site group="3" rgba="0 0 1 .5"/>
          </default>
        </default>
    </default>
    <worldbody>
        <geom name="floor" type="plane" pos="0 0 0.01" size="10 10 0.1" material="floor_material" />
        <body name="pelvis" pos="1.2400539848490296e-16 -0.007071309722959995 0.9311913251876831"
            quat="-0.7069933123978012 0.012667131671581075 0.7069933123978012 -0.012667131671581075">
            <joint type="free" name="joint_pelvis" />
            <geom type="capsule" size="0.055"
                fromto="0.00 -0.006 0.05 0.00 -0.006 -0.05"
                density="1000" rgba="0.8 0.6 0.4 1" />
            
            <body name="spine_01"
                pos="0.06749367713928223 -0.00638774037361145 -5.3054370284229355e-15"
                quat="0.9994206427890385 2.3362404552790525e-18 1.0090234167493236e-17 -0.034034963892227064">
                <joint type="hinge" name="joint_spine_01_x" axis="1 0 0" range="-5 5"
                    class="stiff_high" />
                <joint type="hinge" name="joint_spine_01_y" axis="0 1 0" range="-3 3"
                    class="stiff_high" />
                <joint type="hinge" name="joint_spine_01_z" axis="0 0 1" range="-5 20"
                    class="stiff_high" />
                <geom type="capsule" size="0.03"
                    fromto="0.02 0.0 -0.06 0.02 0.0 0.06"
                    density="1000" rgba="0.8 0.6 0.4 1" />
                <body name="spine_02"
                    pos="0.07251763343811035 -0.00034980475902557373 -1.5774199709026743e-14"
                    quat="0.9997177507990976 5.392413112454937e-23 -2.2919371771182243e-22 -0.023757498546636644">
                    <joint type="hinge" name="joint_spine_02_x" axis="1 0 0" range="-10 10"
                        class="stiff_high" />
                    <joint type="hinge" name="joint_spine_02_y" axis="0 1 0" range="-6 6"
                        class="stiff_high" />
                    <joint type="hinge" name="joint_spine_02_z" axis="0 0 1" range="-10 20"
                        class="stiff_high" />
                    <geom type="capsule" size="0.03"
                        fromto="0.015 0.0 -0.07 0.015 0.0 0.07"
                        density="1000" rgba="0.8 0.6 0.4 1" />
                    <body name="spine_03" pos="0.06010448932647705 0.0 2.104476904969243e-14"
                        quat="0.9998893492393578 6.442381829385969e-23 6.845677000085893e-22 -0.014875795027281036">
                        <joint type="hinge" name="joint_spine_03_x" axis="1 0 0" range="-15 15"
                            class="stiff_high" />
                        <joint type="hinge" name="joint_spine_03_y" axis="0 1 0" range="-15 15"
                            class="stiff_high" />
                        <joint type="hinge" name="joint_spine_03_z" axis="0 0 1" range="-15 20"
                            class="stiff_high" />
                        <geom type="capsule" size="0.03"
                            fromto="0.025 0.0 -0.07 0.025 0.0 0.07"
                            density="1000" rgba="0.8 0.6 0.4 1" />
                        <body name="spine_04"
                            pos="0.0810236930847168 1.4901161193847656e-08 -2.8308088430859316e-14"
                            quat="0.9999134362924395 6.762786830704167e-23 1.323603538731091e-22 0.01315750439276834">
                            <joint type="hinge" name="joint_spine_04_x" axis="1 0 0" range="-15 15"
                                class="stiff_high" />
                            <joint type="hinge" name="joint_spine_04_y" axis="0 1 0" range="-15 15"
                                class="stiff_high" />
                            <joint type="hinge" name="joint_spine_04_z" axis="0 0 1" range="-10 20"
                                class="stiff_high" />
                            <geom type="capsule" size="0.035"
                                fromto="0.02 0.0 -0.08 0.02 0.0 0.08"
                                density="1000" rgba="0.8 0.6 0.4 1" />
                            <body name="spine_05"
                                pos="0.10014617443084717 0.0 4.948447793022559e-14"
                                quat="0.9999988838707776 4.8984655985788956e-23 -5.997066062640126e-22 0.0014940740273561408">
                                <joint type="hinge" name="joint_spine_05_x" axis="1 0 0"
                                    range="-15 15" class="stiff_high" />
                                <joint type="hinge" name="joint_spine_05_y" axis="0 1 0"
                                    range="-15 15" class="stiff_high" />
                                <joint type="hinge" name="joint_spine_05_z" axis="0 0 1"
                                    range="-15 20" class="stiff_high" />
                                <geom type="capsule" size="0.04"
                                    fromto="0.02 0.0 -0.09 0.02 0.0 0.09"
                                    density="1000" rgba="0.8 0.6 0.4 1" />
                                <!-- <geom type="capsule" size="0.025"
                                    fromto="0.0 0.0 0.0 0.043575406074523926 -0.0014655888080596924 0.023223280906677246"
                                    density="1000" rgba="0.8 0.6 0.4 1" />
                                <geom type="capsule" size="0.025"
                                    fromto="0.0 0.0 0.0 0.11313807964324951 0.001528620719909668 -2.3156799963891875e-14"
                                    density="1000" rgba="0.8 0.6 0.4 1" /> -->
                                <body name="clavicle_l"
                                    pos="0.04357445240020752 -0.001465722918510437 -0.023223210126161575"
                                    quat="0.6829487044068947 -0.04647542493978803 0.7289309223867121 0.008989572518595777">
                                    <joint type="hinge" name="joint_clavicle_l_x" axis="1 0 0"
                                        range="-2 2" class="stiff_high" />
                                    <joint type="hinge" name="joint_clavicle_l_y" axis="0 1 0"
                                        range="-15 5" class="stiff_high" />
                                    <joint type="hinge" name="joint_clavicle_l_z" axis="0 0 1"
                                        range="-10 15" class="stiff_high" />
                                    <geom type="capsule" size="0.03"
                                        fromto="0.0 0.0 0.0 0.1551109254360199 0.0 1.1920928955078125e-07"
                                        density="1000" rgba="0.8 0.6 0.4 1" />
                                    <body name="upperarm_l"
                                        pos="0.1551109254360199 0.0 1.1920928955078125e-07"
                                        quat="0.8436629524964866 -0.030462773804957186 0.5359774300853403 0.005747733129975311">
                                        <joint type="hinge" name="joint_upperarm_l_x" axis="1 0 0"
                                            range="-80 80" class="stiff_low" />
                                        <joint type="hinge" name="joint_upperarm_l_y" axis="0 1 0"
                                            range="-140 40" class="stiff_low" />
                                        <joint type="hinge" name="joint_upperarm_l_z" axis="0 0 1"
                                            range="-50 150" class="stiff_low" />
                                        <geom type="capsule" size="0.035"
                                            fromto="0.0 0.0 0.0 0.2665136456489563 0.0 0.0"
                                            density="1000" rgba="0.8 0.6 0.4 1" />
                                        <body name="lowerarm_l" pos="0.2665136456489563 0.0 0.0"
                                            quat="0.9686191585618291 -0.039058773569420266 0.009890000242635354 0.24526215723064054">
                                            <joint type="hinge" name="joint_lowerarm_l" axis="0 0 1"
                                                range="-25 126" />
                                            <geom type="capsule" size="0.025"
                                                fromto="0.0 0.0 0.0 0.223491370677948 0.0 0.0"
                                                density="1000" rgba="0.8 0.6 0.4 1" />
                                            <body name="hand_l" pos="0.223491370677948 0.0 0.0"
                                                quat="0.6975021136747361 0.7081025972512259 0.050679966201636166 -0.09753488714887959">
                                                <joint type="hinge" name="joint_hand_l_x"
                                                    axis="1 0 0" range="-45 45" class="stiff_low" />
                                                <joint type="hinge" name="joint_hand_l_y"
                                                    axis="0 1 0" range="-40 50" class="stiff_low" />
                                                <joint type="hinge" name="joint_hand_l_z"
                                                    axis="0 0 1" range="-20 30" class="stiff_low" />
                                                <geom type="sphere" size="0.025" density="1000"
                                                    pos="0.03 0 0" rgba="0.8 0.6 0.4 1" />
                                            </body>
                                        </body>
                                    </body>
                                </body>
                                <body name="clavicle_r"
                                    pos="0.043575406074523926 -0.0014655888080596924 0.023223280906677246"
                                    quat="0.008989572518595779 0.7289309223867122 0.04647542493978804 -0.6829487044068948">
                                    <joint type="hinge" name="joint_clavicle_r_x" axis="1 0 0"
                                        range="-2 2" class="stiff_high" />
                                    <joint type="hinge" name="joint_clavicle_r_y" axis="0 1 0"
                                        range="-15 5" class="stiff_high" />
                                    <joint type="hinge" name="joint_clavicle_r_z" axis="0 0 1"
                                        range="-10 15" class="stiff_high" />
                                    <geom type="capsule" size="0.03"
                                        fromto="0.0 0.0 0.0 -0.15511080622673035 -1.043081283569336e-07 -1.1920928955078125e-07"
                                        density="1000" rgba="0.8 0.6 0.4 1" />
                                    <body name="upperarm_r"
                                        pos="-0.15511080622673035 -1.043081283569336e-07 -1.1920928955078125e-07"
                                        quat="0.8436629524964866 -0.030462773804957186 0.5359774300853403 0.005747733129975311">
                                        <joint type="hinge" name="joint_upperarm_r_x" axis="1 0 0"
                                            range="-80 80" class="stiff_low" />
                                        <joint type="hinge" name="joint_upperarm_r_y" axis="0 1 0"
                                            range="-140 40" class="stiff_low" />
                                        <joint type="hinge" name="joint_upperarm_r_z" axis="0 0 1"
                                            range="-50 150" class="stiff_low" />
                                        <geom type="capsule" size="0.035"
                                            fromto="0.0 0.0 0.0 -0.26651448011398315 1.4901161193847656e-08 0.0"
                                            density="1000" rgba="0.8 0.6 0.4 1" />
                                        <body name="lowerarm_r"
                                            pos="-0.26651448011398315 1.4901161193847656e-08 0.0"
                                            quat="0.9686191585618291 -0.039058773569420266 0.009890000242635354 0.24526215723064054">
                                            <joint type="hinge" name="joint_lowerarm_r" axis="0 0 1"
                                                range="-25 126" />
                                            <geom type="capsule" size="0.025"
                                                fromto="0.0 0.0 0.0 -0.22349244356155396 5.960464477539063e-08 5.960464477539063e-08"
                                                density="1000" rgba="0.8 0.6 0.4 1" />
                                            <body name="hand_r"
                                                pos="-0.22349244356155396 5.960464477539063e-08 5.960464477539063e-08"
                                                quat="0.6975021136747361 0.7081025972512259 0.050679966201636166 -0.09753488714887959">
                                                <joint type="hinge" name="joint_hand_r_x"
                                                    axis="1 0 0" range="-45 45" class="stiff_low" />
                                                <joint type="hinge" name="joint_hand_r_y"
                                                    axis="0 1 0" range="-40 50" class="stiff_low" />
                                                <joint type="hinge" name="joint_hand_r_z"
                                                    axis="0 0 1" range="-20 30" class="stiff_low" />
                                                <geom type="sphere" size="0.025" density="1000"
                                                    pos="-0.03 0 0" rgba="0.8 0.6 0.4 1" />
                                            </body>
                                        </body>
                                    </body>
                                </body>
                                <body name="neck_01"
                                    pos="0.11313807964324951 0.001528620719909668 -2.3156799963891875e-14"
                                    quat="0.979739387686268 -7.759519592020889e-23 -2.5835163497133134e-22 0.20027663921720062">
                                    <joint type="hinge" name="joint_neck_01_x" axis="1 0 0"
                                        range="-10 10" class="stiff_medium" />
                                    <joint type="hinge" name="joint_neck_01_y" axis="0 1 0"
                                        range="-10 10" class="stiff_medium" />
                                    <joint type="hinge" name="joint_neck_01_z" axis="0 0 1"
                                        range="-30 30" class="stiff_medium" />
                                    <geom type="capsule" size="0.025"
                                        fromto="-0.02 0.0 0.0 0.0 0.0 0"
                                        density="1000" rgba="0.8 0.6 0.4 1" />
                                    <body name="neck_02"
                                        pos="0.04562675952911377 0.0 -4.467488933044411e-16"
                                        quat="0.9987275844696399 -2.976467671374415e-22 -1.0526860272186105e-21 -0.050430268881280355">
                                        <joint type="hinge" name="joint_neck_02_x" axis="1 0 0"
                                            range="-10 10" class="stiff_medium" />
                                        <joint type="hinge" name="joint_neck_02_y" axis="0 1 0"
                                            range="-10 10" class="stiff_medium" />
                                        <joint type="hinge" name="joint_neck_02_z" axis="0 0 1"
                                            range="-30 30" class="stiff_medium" />
                                        <geom type="capsule" size="0.025"
                                            fromto="-0.01 0.0 0.0 0.01 0.0 0"
                                            density="1000" rgba="0.8 0.6 0.4 1" />
                                        <body name="head"
                                            pos="0.05391383171081543 0.0 -1.1647306535721548e-15"
                                            quat="0.9933705878621503 1.1486719478815802e-22 5.2460160773161134e-23 -0.11495597057311105">
                                            <joint type="hinge" name="joint_head_x" axis="1 0 0"
                                                range="-70 70" class="stiff_medium" />
                                            <joint type="hinge" name="joint_head_y" axis="0 1 0"
                                                range="-2 2" class="stiff_medium" />
                                            <joint type="hinge" name="joint_head_z" axis="0 0 1"
                                                range="-3 3" class="stiff_medium" />
                                            <geom type="sphere" size="0.07" density="1000"
                                                pos= "0.01 0 0" rgba="0.8 0.6 0.4 1" />
                                        </body>
                                    </body>
                                </body>
                            </body>
                        </body>
                    </body>
                </body>
            </body>
            <body name="thigh_l"
                pos="-0.025693297386169434 -0.00010870397090911865 -0.07750068604946136"
                quat="0.9982444406463985 -0.0012370699081630204 -0.0571613282515377 -0.015462500739568058">
                <joint type="hinge" name="joint_thigh_l_x" axis="1 0 0" range="-20 20"
                    class="stiff_medium" />
                <joint type="hinge" name="joint_thigh_l_y" axis="0 1 0" range="-40 50"
                    class="stiff_medium" />
                <joint type="hinge" name="joint_thigh_l_z" axis="0 0 1" range="-90 20"
                    class="stiff_medium" />
                <geom type="capsule" size="0.045"
                    fromto="0.0 0.0 0.0 -0.39048439264297485 9.313225746154785e-10 0.0"
                    density="1000" rgba="0.8 0.6 0.4 1" />
                <body name="calf_l" pos="-0.39048439264297485 9.313225746154785e-10 0.0"
                    quat="0.9991078205558419 -0.0006403708434938144 0.01606467134332963 0.03905226195714184">
                    <joint type="hinge" name="joint_calf_l" axis="0 0 1" range="0 151" />
                    <geom type="capsule" size="0.035"
                        fromto="0.0 0.0 0.0 -0.42561033368110657 0.0010773949325084686 0.004291892051696777"
                        density="1000" rgba="0.8 0.6 0.4 1" />
                    <body name="foot_l"
                        pos="-0.42561033368110657 0.0010773949325084686 0.004291892051696777"
                        quat="0.997500889014535 -0.04178496437378156 0.045505050896279806 -0.03428240817709532">
                        <joint type="hinge" name="joint_foot_l_x" axis="1 0 0" range="-5 5"
                            class="stiff_low" />
                        <joint type="hinge" name="joint_foot_l_y" axis="0 1 0" range="-10 10"
                            class="stiff_low" />
                        <joint type="hinge" name="joint_foot_l_z" axis="0 0 1" range="-15 15"
                            class="stiff_low" />
                        <geom name="foot_l" type="box" size="0.025 0.08 0.04" pos="-0.05 0.035 0"
                            density="1000" rgba="0.8 0.6 0.4 1" />
                        <body name="ball_l"
                            pos="-0.06159850209951401 0.13826413452625275 0.006806045770645142"
                            quat="0.7079318965314102 -6.0294361177534856e-05 6.043667739381551e-05 0.7062806967385125">
                            <joint type="hinge" name="joint_ball_l_x" axis="1 0 0" range="-2 2"
                                class="stiff_low" />
                            <joint type="hinge" name="joint_ball_l_y" axis="0 1 0" range="-2 2"
                                class="stiff_low" />
                            <joint type="hinge" name="joint_ball_l_z" axis="0 0 1" range="-10 10"
                                class="stiff_low" />
                            <geom name="ball_l" type="box" size="0.03 0.02 0.02" pos="0.025 -0.01 0"
                            density="1000" rgba="0.8 0.6 0.4 1" />
                        </body>
                    </body>
                </body>
            </body>
            <body name="thigh_r"
                pos="-0.02569139003753662 -0.00010874494910240173 0.07750062644481659"
                quat="0.015462500739568058 0.0571613282515377 -0.0012370699081630204 0.9982444406463985">
                <joint type="hinge" name="joint_thigh_r_x" axis="1 0 0" range="-20 20"
                    class="stiff_medium" />
                <joint type="hinge" name="joint_thigh_r_y" axis="0 1 0" range="-40 50"
                    class="stiff_medium" />
                <joint type="hinge" name="joint_thigh_r_z" axis="0 0 1" range="-90 20"
                    class="stiff_medium" />
                <geom type="capsule" size="0.045"
                    fromto="0.0 0.0 0.0 0.39048659801483154 -2.2351741790771484e-08 0.0"
                    density="1000" rgba="0.8 0.6 0.4 1" />
                <body name="calf_r" pos="0.39048659801483154 -2.2351741790771484e-08 0.0"
                    quat="0.9991078205558419 -0.0006403708434938144 0.01606467134332963 0.03905226195714184">
                    <joint type="hinge" name="joint_calf_r" axis="0 0 1" range="0 151" />
                    <geom type="capsule" size="0.035"
                        fromto="0.0 0.0 0.0 0.42561158537864685 -0.0010772906243801117 -0.004292041063308716"
                        density="1000" rgba="0.8 0.6 0.4 1" />
                    <body name="foot_r"
                        pos="0.42561158537864685 -0.0010772906243801117 -0.004292041063308716"
                        quat="0.997500889014535 -0.04178496437378156 0.045505050896279806 -0.03428240817709532">
                        <joint type="hinge" name="joint_foot_r_x" axis="1 0 0" range="-5 5"
                            class="stiff_low" />
                        <joint type="hinge" name="joint_foot_r_y" axis="0 1 0" range="-10 10"
                            class="stiff_low" />
                        <joint type="hinge" name="joint_foot_r_z" axis="0 0 1" range="-15 15"
                            class="stiff_low" />
                        <geom name="foot_r" type="box" size="0.025 0.08 0.04" pos="0.05 -0.035 0"
                        density="1000" rgba="0.8 0.6 0.4 1" />
                        <body name="ball_r"
                            pos="0.06159842759370804 -0.13826420903205872 -0.006806075572967529"
                            quat="0.7079318965314102 -6.0294361177534856e-05 6.043667739381551e-05 0.7062806967385125">
                            <joint type="hinge" name="joint_ball_r_x" axis="1 0 0" range="-2 2"
                                class="stiff_low" />
                            <joint type="hinge" name="joint_ball_r_y" axis="0 1 0" range="-2 2"
                                class="stiff_low" />
                            <joint type="hinge" name="joint_ball_r_z" axis="0 0 1" range="-10 10"
                                class="stiff_low" />
                            <geom name="ball_r"  type="box" size="0.03 0.02 0.02" pos="-0.025 0.01 0"
                            density="1000" rgba="0.8 0.6 0.4 1" />
                        </body>
                    </body>
                </body>
            </body>
        </body>
    </worldbody>
    <contact>
        <exclude name="ignr_clavicle1" body1="clavicle_l" body2="clavicle_r"/>
        <exclude name="ignr_spine_clavl" body1="clavicle_l" body2="spine_04"/>
        <exclude name="ignr_spine_clavr" body1="clavicle_r" body2="spine_04"/>
        <exclude name="ignr_neck_clavl" body1="clavicle_l" body2="neck_01"/>
        <exclude name="ignr_neck_clavr" body1="clavicle_r" body2="neck_01"/>
        <exclude name="ignr_spine_neck" body1="neck_02" body2="spine_05"/>
        <exclude name="ignr_neck_head" body1="neck_01" body2="head"/>
        <pair geom1="foot_l" geom2="floor"/>
        <pair geom1="foot_r" geom2="floor"/>
        <pair geom1="ball_l" geom2="floor"/>
        <pair geom1="ball_r" geom2="floor"/>
    </contact>
    <actuator>
        <motor name="motor_spine_01_x" gear="40" joint="joint_spine_01_x" />
        <motor name="motor_spine_01_y" gear="40" joint="joint_spine_01_y" />
        <motor name="motor_spine_01_z" gear="40" joint="joint_spine_01_z" />
        <motor name="motor_spine_02_x" gear="40" joint="joint_spine_02_x" />
        <motor name="motor_spine_02_y" gear="40" joint="joint_spine_02_y" />
        <motor name="motor_spine_02_z" gear="40" joint="joint_spine_02_z" />
        <motor name="motor_spine_03_x" gear="40" joint="joint_spine_03_x" />
        <motor name="motor_spine_03_y" gear="40" joint="joint_spine_03_y" />
        <motor name="motor_spine_03_z" gear="40" joint="joint_spine_03_z" />
        <motor name="motor_spine_04_x" gear="40" joint="joint_spine_04_x" />
        <motor name="motor_spine_04_y" gear="40" joint="joint_spine_04_y" />
        <motor name="motor_spine_04_z" gear="40" joint="joint_spine_04_z" />
        <motor name="motor_spine_05_x" gear="40" joint="joint_spine_05_x" />
        <motor name="motor_spine_05_y" gear="40" joint="joint_spine_05_y" />
        <motor name="motor_spine_05_z" gear="40" joint="joint_spine_05_z" />
        <motor name="motor_clavicle_l_x" gear="40" joint="joint_clavicle_l_x" />
        <motor name="motor_clavicle_l_y" gear="40" joint="joint_clavicle_l_y" />
        <motor name="motor_clavicle_l_z" gear="40" joint="joint_clavicle_l_z" />
        <motor name="motor_upperarm_l_x" gear="20" joint="joint_upperarm_l_x" />
        <motor name="motor_upperarm_l_y" gear="20" joint="joint_upperarm_l_y" />
        <motor name="motor_upperarm_l_z" gear="20" joint="joint_upperarm_l_z" />
        <motor name="motor_lowerarm_l" gear="40" joint="joint_lowerarm_l" />
        <motor name="motor_hand_l_x" gear="20" joint="joint_hand_l_x" />
        <motor name="motor_hand_l_y" gear="20" joint="joint_hand_l_y" />
        <motor name="motor_hand_l_z" gear="20" joint="joint_hand_l_z" />
        <motor name="motor_clavicle_r_x" gear="40" joint="joint_clavicle_r_x" />
        <motor name="motor_clavicle_r_y" gear="40" joint="joint_clavicle_r_y" />
        <motor name="motor_clavicle_r_z" gear="40" joint="joint_clavicle_r_z" />
        <motor name="motor_upperarm_r_x" gear="20" joint="joint_upperarm_r_x" />
        <motor name="motor_upperarm_r_y" gear="20" joint="joint_upperarm_r_y" />
        <motor name="motor_upperarm_r_z" gear="20" joint="joint_upperarm_r_z" />
        <motor name="motor_lowerarm_r" gear="40" joint="joint_lowerarm_r" />
        <motor name="motor_hand_r_x" gear="20" joint="joint_hand_r_x" />
        <motor name="motor_hand_r_y" gear="20" joint="joint_hand_r_y" />
        <motor name="motor_hand_r_z" gear="20" joint="joint_hand_r_z" />
        <motor name="motor_neck_01_x" gear="40" joint="joint_neck_01_x" />
        <motor name="motor_neck_01_y" gear="40" joint="joint_neck_01_y" />
        <motor name="motor_neck_01_z" gear="40" joint="joint_neck_01_z" />
        <motor name="motor_neck_02_x" gear="40" joint="joint_neck_02_x" />
        <motor name="motor_neck_02_y" gear="40" joint="joint_neck_02_y" />
        <motor name="motor_neck_02_z" gear="40" joint="joint_neck_02_z" />
        <motor name="motor_head_x" gear="40" joint="joint_head_x" />
        <motor name="motor_head_y" gear="40" joint="joint_head_y" />
        <motor name="motor_head_z" gear="40" joint="joint_head_z" />
        <motor name="motor_thigh_l_x" gear="120" joint="joint_thigh_l_x" />
        <motor name="motor_thigh_l_y" gear="120" joint="joint_thigh_l_y" />
        <motor name="motor_thigh_l_z" gear="120" joint="joint_thigh_l_z" />
        <motor name="motor_calf_l" gear="80" joint="joint_calf_l" />
        <motor name="motor_foot_l_x" gear="40" joint="joint_foot_l_x" />
        <motor name="motor_foot_l_y" gear="40" joint="joint_foot_l_y" />
        <motor name="motor_foot_l_z" gear="40" joint="joint_foot_l_z" />
        <motor name="motor_ball_l_x" gear="40" joint="joint_ball_l_x" />
        <motor name="motor_ball_l_y" gear="40" joint="joint_ball_l_y" />
        <motor name="motor_ball_l_z" gear="40" joint="joint_ball_l_z" />
        <motor name="motor_thigh_r_x" gear="120" joint="joint_thigh_r_x" />
        <motor name="motor_thigh_r_y" gear="120" joint="joint_thigh_r_y" />
        <motor name="motor_thigh_r_z" gear="120" joint="joint_thigh_r_z" />
        <motor name="motor_calf_r" gear="80" joint="joint_calf_r" />
        <motor name="motor_foot_r_x" gear="40" joint="joint_foot_r_x" />
        <motor name="motor_foot_r_y" gear="40" joint="joint_foot_r_y" />
        <motor name="motor_foot_r_z" gear="40" joint="joint_foot_r_z" />
        <motor name="motor_ball_r_x" gear="40" joint="joint_ball_r_x" />
        <motor name="motor_ball_r_y" gear="40" joint="joint_ball_r_y" />
        <motor name="motor_ball_r_z" gear="40" joint="joint_ball_r_z" />
    </actuator>
    <asset>
        <texture type="2d" name="checkerboard" builtin="checker" rgb1="0 0 0" rgb2="1 1 1"
            width="256" height="256" />
        <material name="floor_material" texture="checkerboard" texrepeat="10 10" rgba="1 1 1 1" />
    </asset>
</mujoco>

Confirmations
@KamatMayur KamatMayur added the bug Something isn't working label Nov 20, 2024
@KamatMayur
Copy link
Author

Alright so the following changes seemed to make it work!

  1. os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = ".5" # reduce to 50% of GPU from default of 75%
  2. in the xml file. changing the options to use iterations="1" and ls_iterations="4" and disabling the eulerdamp flag.
  3. Reducing the iterations and ls_iterations was producing instability so I had to additionally increase the armature value for all my joints from 0.1 to 0.3
  4. Disabling the collisions on all geoms by setting contype="0" conaffinity="0" and the explicitly mentioning the pair of bodies that will collide under the contact section using pair attribute.

These changes have worked so far although I still get this one warning or error I don't really know what it is
E1122 14:18:54.592463 24092 hlo_lexer.cc:443] Failed to parse int literal: 894515288310727292233 Although the code seems to run just fine.

@saran-t saran-t added MJX Using JAX to run on GPU and removed bug Something isn't working labels Nov 22, 2024
@btaba
Copy link
Collaborator

btaba commented Dec 22, 2024

Hi @KamatMayur , I loaded your model and created an mjx.Data. ncon is 406, you'll want to use max_geom_pairs and max_contact_points to reduce the memory, or use contype conaffinity to reduce contact pairs. We're working on better support for broadphase, but with the impl at HEAD you'll have to manually specify a static ncon.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
MJX Using JAX to run on GPU
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
3 participants
@saran-t @btaba @KamatMayur