search

google / brax

Massively parallel rigidbody physics simulation on accelerator hardware.
Apache License 2.0
2.14k stars 234 forks source link

Error in Importing XML Model with Extended DOFs in Brax #426

Closed bheijden closed 7 months ago

bheijden commented 7 months ago

Description: Encountering a KeyError when importing an XML model that includes an object with 4 degrees of freedom (DOFs). The issue arises when adding additional joints to the object in the pusher.xml, specifically a z-slider and a yaw-hinge, increasing the DOFs to four (x, y, z, yaw).

MWE: To replicate the issue, I modified the pusher.xml by adding two joints to the object: a z-slider and a yaw-hinge. Below is the Python snippet used to load the modified XML and to print the degrees of freedom:

Expected Behavior: The expected output should correctly reflect the number of degrees of freedom of the object, which in this case has been increased to four.

Actual Behavior: The code throws a KeyError: '4', indicating an issue with handling the extended DOFs in the XML model.

Additional Information: The issue seems to be specific to handling models with DOFs beyond the standard x, y, and z. Any guidance or fix for this issue would be greatly appreciated. Are only 1,2,3, and 6 DOFs currently supported?

Environment: Brax version: 0.9.3 Python version: 3.8 OS: Ubuntu 20.04

from brax.io import mjcf
m = mjcf.load('./pusher_4dof.xml')
print(f"degrees of freedom: {m.qd_size()}")  # Throws KeyError: '4'
<mujoco model="arm3d">
  <compiler inertiafromgeom="true" angle="radian" coordinate="local"/>
  <option timestep="0.01" gravity="0 0 0" iterations="20" integrator="Euler" />

  <default>
    <joint armature='0.04' damping="1" limited="true"/>
    <geom friction=".8 .1 .1" density="300" margin="0.002" condim="1" contype="0" conaffinity="0"/>
  </default>

  <custom>
    <!-- brax custom params -->
    <numeric data="3000" name="constraint_stiffness"/>
    <numeric data="2000" name="constraint_limit_stiffness"/>
    <numeric data="10" name="constraint_ang_damping"/>
    <numeric data="50" name="constraint_vel_damping"/>
    <numeric data="0" name="ang_damping"/>
    <numeric data="1" name="spring_mass_scale"/>
        <numeric data="1" name="spring_inertia_scale"/>
        <numeric data="10" name="solver_maxls"/>
  </custom>

  <worldbody>
    <light diffuse=".5 .5 .5" pos="0 0 3" dir="0 0 -1"/>
    <geom name="table" type="plane" pos="0 0.5 -0.325" size="1 1 0.1" contype="1" conaffinity="0"/>

    <body name="r_shoulder_pan_link" pos="0 -0.6 0">
      <geom name="e1" type="sphere" rgba="0.6 0.6 0.6 1" pos="-0.06 0.05 0.2" size="0.05" />
      <geom name="e2" type="sphere" rgba="0.6 0.6 0.6 1" pos=" 0.06 0.05 0.2" size="0.05" />
      <geom name="e1p" type="sphere" rgba="0.1 0.1 0.1 1" pos="-0.06 0.09 0.2" size="0.03" />
      <geom name="e2p" type="sphere" rgba="0.1 0.1 0.1 1" pos=" 0.06 0.09 0.2" size="0.03" />
      <geom name="sp" type="capsule" fromto="0 0 -0.4 0 0 0.2" size="0.1" />
      <joint name="r_shoulder_pan_joint" type="hinge" pos="0 0 0" axis="0 0 1" range="-2.2854 1.714602" damping="1.0" />

      <body name="r_shoulder_lift_link" pos="0.1 0 0">
        <geom name="sl" type="capsule" fromto="0 -0.1 0 0 0.1 0" size="0.1" />
        <joint name="r_shoulder_lift_joint" type="hinge" pos="0 0 0" axis="0 1 0" range="-0.5236 1.3963" damping="1.0" />

        <body name="r_upper_arm_roll_link" pos="0 0 0">
          <geom name="uar" type="capsule" fromto="-0.1 0 0 0.1 0 0" size="0.02" />
          <joint name="r_upper_arm_roll_joint" type="hinge" pos="0 0 0" axis="1 0 0" range="-1.5 1.7" damping="0.1" />

          <body name="r_upper_arm_link" pos="0 0 0">
            <geom name="ua" type="capsule" fromto="0 0 0 0.4 0 0" size="0.06" />

            <body name="r_elbow_flex_link" pos="0.4 0 0">
              <geom name="ef" type="capsule" fromto="0 -0.02 0 0.0 0.02 0" size="0.06" />
              <joint name="r_elbow_flex_joint" type="hinge" pos="0 0 0" axis="0 1 0" range="-2.3213 0" damping="0.1" />

              <body name="r_forearm_roll_link" pos="0 0 0">
                <geom name="fr" type="capsule" fromto="-0.1 0 0 0.1 0 0" size="0.02" />
                <joint name="r_forearm_roll_joint" type="hinge" limited="true" pos="0 0 0" axis="1 0 0" damping=".1" range="-1.5 1.5"/>

                <body name="r_forearm_link" pos="0 0 0">
                  <geom name="fa" type="capsule" fromto="0 0 0 0.291 0 0" size="0.05" />

                  <body name="r_wrist_flex_link" pos="0.321 0 0">
                    <geom name="wf" type="capsule" fromto="0 -0.02 0 0 0.02 0" size="0.01" />
                    <joint name="r_wrist_flex_joint" type="hinge" pos="0 0 0" axis="0 1 0" range="-1.094 0" damping=".1" />

                    <body name="r_wrist_roll_link" pos="0 0 0">
                      <joint name="r_wrist_roll_joint" type="hinge" pos="0 0 0" limited="true" axis="1 0 0" damping="0.1" range="-1.5 1.5"/>
                      <body name="tips_arm" pos="0 0 0">
                        <geom name="tip_arml" type="sphere" pos="0.1 -0.1 0." size="0.01" />
                        <geom name="tip_armr" type="sphere" pos="0.1 0.1 0." size="0.01" />
                      </body>
                      <geom name="r_wrist_roll1" type="capsule" fromto="0 -0.1 0. 0.0 +0.1 0" size="0.02" contype="0" conaffinity="1" />
                      <geom name="r_wrist_roll2" type="capsule" fromto="0 -0.1 0. 0.1 -0.1 0" size="0.02" contype="0" conaffinity="1" />
                      <geom name="r_wrist_roll3" type="capsule" fromto="0 +0.1 0. 0.1 +0.1 0." size="0.02" contype="0" conaffinity="1" />
                    </body>
                  </body>
                </body>
              </body>
            </body>
          </body>
        </body>
      </body>
    </body>

    <!--<body name="object" pos="0.55 -0.3 -0.275" >-->
    <body name="object" pos="0.45 -0.05 -0.275" >
      <!-- Density changed to 1 from 0.00001 for brax. -->
      <geom rgba="1 1 1 0" type="sphere" size="0.05 0.05 0.05" density="1" conaffinity="0"/>
      <geom name="object_geom" rgba="1 1 1 1" type="capsule" size="0.05 0.05 0.05" density="1" contype="1" conaffinity="0"/>
      <joint name="obj_slidey" type="slide" pos="0 0 0" axis="0 1 0" range="-10.3213 10.3" damping="0.5"/>
      <joint name="obj_slidex" type="slide" pos="0 0 0" axis="1 0 0" range="-10.3213 10.3" damping="0.5"/>
      <joint name="goal_slidez" type="slide" pos="0 0 0" axis="0 0 1" range="-10.3213 10.3" damping="0.5"/> <!-- EXTRA DOF -->
      <joint name="goal_yaw" type="hinge" pos="0 0 0" axis="0 0 1" limited="false"  damping="0.5"/> <!-- EXTRA DOF -->
    </body>

    <body name="goal" pos="0.45 -0.05 -0.3230">
      <geom rgba="1 0 0 1" type="capsule" size="0.08 0.001 0.1" density='0.00001' contype="0" conaffinity="0"/>
      <joint name="goal_slidey" type="slide" pos="0 0 0" axis="0 1 0" range="-10.3213 10.3" damping="0.5"/>
      <joint name="goal_slidex" type="slide" pos="0 0 0" axis="1 0 0" range="-10.3213 10.3" damping="0.5"/
    </body>
  </worldbody>

  <actuator>
    <motor joint="r_shoulder_pan_joint" ctrlrange="-2.0 2.0" ctrllimited="true" />
    <motor joint="r_shoulder_lift_joint" ctrlrange="-2.0 2.0" ctrllimited="true" />
    <motor joint="r_upper_arm_roll_joint" ctrlrange="-2.0 2.0" ctrllimited="true" />
    <motor joint="r_elbow_flex_joint" ctrlrange="-2.0 2.0" ctrllimited="true" />
    <motor joint="r_forearm_roll_joint" ctrlrange="-2.0 2.0" ctrllimited="true" />
    <motor joint="r_wrist_flex_joint" ctrlrange="-2.0 2.0" ctrllimited="true" />
    <motor joint="r_wrist_roll_joint" ctrlrange="-2.0 2.0" ctrllimited="true"/>
  </actuator>

</mujoco>
btaba commented 7 months ago

Hi @bheijden , brax doesn't support more than 3 stacked joints. Can you add goal_yaw to another body?

This is supported in MuJoCo-MJX FYI