google-deepmind / dm_control

Google DeepMind's software stack for physics-based simulation and Reinforcement Learning environments, using MuJoCo.
Apache License 2.0
3.75k stars 665 forks source link

Attach a class-defined body to another class-defined body using joints #480

Open ajoardar007 opened 3 months ago

ajoardar007 commented 3 months ago

Hello, In order attach multiple links together, each having its own collision/visual mesh, inertia, joint, etc. I decided to construct a class that defines such a link, and it is then called multiple times, where each of its objects are initialized accordingly.

The issue that i am facing is that, when attaching such a child-body (Child) to the parent body, where the parent-body (Parent-2) has a joint w.r.t. its parent (Parent-1), I am unable to place the Child inside the XMLdefinition of Parent-2 where the joint is defined, preventing the effects of the joint (Joint 1-to-2) between Parent-1 and Parent-2 to effect the Child body.

Below, I define a class for constructing bodies with joints:

class Link_w_Joints:
    #===================#
    #  C O N S T A N T  #
    #===================#

    #===============================#
    #  I N I T I A L I Z A T I O N  #
    #===============================#
    # joint axis default value [0, 0, 1]: https://mujoco.readthedocs.io/en/stable/XMLreference.html#body-joint
    def __init__(self, col_meshes, com_pos, dclass_col, inertia, joint_damping, joint_frictionloss, joint_name, joint_range, mass, name, viz_meshes, dclass_viz = None, joint_axis = [0, 0, 1]):
        # Make a MujoCo Root
        self.mjcf_model = mjcf.RootElement(model=name)
        # Add a body element:
        self.link_body = self.mjcf_model.worldbody.add('body',
                                                        name = name)
        # Add Inertial properties
        self.link_body.add('inertial',
                           pos = com_pos, 
                           mass = mass,
                           fullinertia = inertia)
        # Add Joints
        self.link_body.add('joint',
                           name = joint_name,
                           range = joint_range, 
                           damping = joint_damping,
                           axis = joint_axis,
                           frictionloss = joint_frictionloss,
                           dclass = 'rand')
        # Add the Geometry (https://github.com/google-deepmind/dm_control/blob/ed424509c0a0e8ddf7a43824924de483026ad9cc/dm_control/locomotion/soccer/humanoid.py#L50)
        # For visualization
        self.link_body.add('geom',
                           dclass = name+'\\viz' if dclass_viz is None else dclass_viz,
                           mesh = viz_meshes)
        # For Collision
        for x in range(len(col_meshes)):
            self.link_body.add('geom',
                               dclass = dclass_col,
                               mesh = col_meshes[x])

and below, I call the class Link_w_Joints multiple-times and attach successive-links (from ground to tool) to the preceding-ones:

class WAM:
    #===================#
    #  C O N S T A N T  #
    #===================#
    _body_name = ['base_link',
                  'shoulder_yaw_link',
                  'shoulder_pitch_link']
    _bhand_name = 'bhand'
    _dclass_col = 'wam\\col'
    _camera_pos = [0, 0.085, 0.02]
    _camera_name = "camera"
    _camera_fovy = 55
    _col_mesh_names = [['shoulder_link_convex_decomposition_p1',
                        'shoulder_link_convex_decomposition_p2',
                        'shoulder_link_convex_decomposition_p3'],
                       ['shoulder_pitch_link_convex']]
    _com_pos = [[-0.00443422, -0.00066489, -0.12189039],
                [-0.00236983, -0.01542114, 0.03105614]]
    _inertial = [[0.13488033, 0.09046330, 0.11328369, -0.00012485, 0.00213041, -0.00068555],
                [0.02140958, 0.01558906, 0.01377875, -0.00002461, 0.00027172, 0.00181920]]
    _joint_damping = [1000,
                       1000]
    _joint_frictionloss = 1000
    _joint_name = ['wam\\J1',
                    'wam\\J2']
    _joint_range = [[-2.6, 2.6],
                     [-2.0, 2.0]]
    _mass = [10.76768767,
              3.8749]
    _pos = [[0, 0, 0.346],
            [0, 0, 0]]
    # quant=w x y z 
    _quat = [[1, 0, 0, 0],
            [0.707107, -0.707107, 0, 0]]
    _site_name = 'site'
    _viz_mesh_names = ['shoulder_link_fine',
                       'shoulder_pitch_link_fine']
    _wam_name = 'wam'
    #===============================#
    #  I N I T I A L I Z A T I O N  #
    #===============================#
    def __init__(self, name):
        # Make a MujoCo Root
        self.mjcf_model = mjcf.RootElement(model=name)
        # Add 'base_link' element:
        self.base_link = self.mjcf_model.worldbody.add('body',
                                                        name = self._wam_name+'\\'+self._body_name[0],
                                                        pos = (0, 0, 0))

        # Attach site for 'shoulder_yaw_link'
        self.shoulder_yaw_link_site = self.base_link.add('site',
                                                          pos = self._pos[0],
                                                          quat = self._quat[0],
                                                          dclass = self._site_name)
        # Create link 'shoulder_yaw_link'
        self.shoulder_yaw_link = Link_w_Joints(col_meshes = self._col_mesh_names[0],
                                                com_pos = self._com_pos[0],
                                                dclass_col = 'wam\\col',
                                                inertia = self._inertial[0], 
                                                joint_damping = self._joint_damping[0],
                                                joint_frictionloss = self._joint_frictionloss,
                                                joint_name = self._joint_name[0],
                                                joint_range = self._joint_range[0],
                                                mass = self._mass[0],
                                                name = self._wam_name+'\\'+self._body_name[1],
                                                viz_meshes = self._viz_mesh_names[0])
        # Attach 'shoulder_yaw_link'
        self.shoulder_yaw_link_site.attach(self.shoulder_yaw_link.mjcf_model)

        # Attach site for 'shoulder_pitch_link'
        self.shoulder_pitch_link_site = self.shoulder_yaw_link.mjcf_model.worldbody.add('site',
                                                                                                                                         pos = self._pos[1],
                                                                                                                                         quat = self._quat[1],
                                                                                                                                         dclass = self._site_name)
        # Create link 'shoulder_pitch_link'
        self.shoulder_pitch_link = Link_w_Joints(col_meshes = self._col_mesh_names[1],
                                               com_pos = self._com_pos[1],
                                               dclass_col = 'wam\\col',
                                               inertia = self._inertial[1], 
                                               joint_damping = self._joint_damping[1],
                                               joint_frictionloss = self._joint_frictionloss,
                                               joint_name = self._joint_name[1],
                                               joint_range = self._joint_range[1],
                                               mass = self._mass[1],
                                               name = self._wam_name+'\\'+self._body_name[2],
                                               viz_meshes = self._viz_mesh_names[1])
        # Attach 'shoulder_pitch_link'
        self.shoulder_pitch_link_site.attach(self.shoulder_pitch_link.mjcf_model)

I have already identified the cause of said problem:

  1. When attaching a body to a site, the worldbody-element of the Child is converted into a body-element and inserted into Parent-2 link,
  2. Because the same thing had happened when attaching Parent-2 to Parent-1, the Child does not end up inside the body-element that holds the joint-definition, but in the outer-one that was defined when the worldbody-element of the child-body is converted into a body-element.

Below, you will find the output of the above script which shows the above-described behavior:

<body name="wam_7dof_bhand_0\wam\base_link" pos="0 0 0">
          <body name="wam_7dof_bhand_0\wam\shoulder_yaw_link\" pos="0 0 0.34599999999999997" quat="1 0 0 0">
            <body name="wam_7dof_bhand_0\wam\shoulder_yaw_link\wam\shoulder_yaw_link">
                <inertial pos="-0.00443422 -0.00066489000000000001 -0.12189039" mass="10.767687670000001" fullinertia="0.13488032999999999 0.090463299999999996 0.11328369000000001 -0.00012485 0.0021304100000000001 -0.00068555000000000005"/>
                <joint name="wam\J1" class="rand" axis="0 0 1" range="-2.6000000000000001 2.6000000000000001" damping="1000" frictionloss="1000"/>
                <geom name="//unnamed_geom_18" class="wam\shoulder_yaw_link\viz" mesh="shoulder_link_fine"/>
                <geom name="//unnamed_geom_19" class="wam\col" mesh="shoulder_link_convex_decomposition_p1"/>
                <geom name="//unnamed_geom_20" class="wam\col" mesh="shoulder_link_convex_decomposition_p2"/>
                <geom name="//unnamed_geom_21" class="wam\col" mesh="shoulder_link_convex_decomposition_p3"/>
            </body>
            <body name="wam_7dof_bhand_0\wam\shoulder_yaw_link\wam\shoulder_pitch_link\" pos="0 0 0" quat="0.70710700000000004 -0.70710700000000004 0 0">
              <body name="wam_7dof_bhand_0\wam\shoulder_yaw_link\wam\shoulder_pitch_link\wam\shoulder_pitch_link">
                  <inertial pos="-0.0023698299999999999 -0.01542114 0.031056139999999999" mass="3.8748999999999998" fullinertia="0.021409580000000001 0.01558906 0.013778749999999999 -2.461e-05 0.00027171999999999998 0.0018192"/>
                  <joint name="wam\J2" class="rand" axis="0 0 1" range="-2 2" damping="1000" frictionloss="1000"/>
                  <geom name="//unnamed_geom_22" class="wam\shoulder_pitch_link\viz" mesh="shoulder_pitch_link_fine"/>
                  <geom name="//unnamed_geom_23" class="wam\col" mesh="shoulder_pitch_link_convex"/>
              </body>

            </body>
        </body>
 </body>

What is desired is the following output (the problematic-lines have been commented <-- -->):

<body name="wam_7dof_bhand_0\wam\base_link" pos="0 0 0">
          <body name="wam_7dof_bhand_0\wam\shoulder_yaw_link\" pos="0 0 0.34599999999999997" quat="1 0 0 0">
            <--body name="wam_7dof_bhand_0\wam\shoulder_yaw_link\wam\shoulder_yaw_link"-->
                <inertial pos="-0.00443422 -0.00066489000000000001 -0.12189039" mass="10.767687670000001" fullinertia="0.13488032999999999 0.090463299999999996 0.11328369000000001 -0.00012485 0.0021304100000000001 -0.00068555000000000005"/>
                <joint name="wam\J1" class="rand" axis="0 0 1" range="-2.6000000000000001 2.6000000000000001" damping="1000" frictionloss="1000"/>
                <geom name="//unnamed_geom_18" class="wam\shoulder_yaw_link\viz" mesh="shoulder_link_fine"/>
                <geom name="//unnamed_geom_19" class="wam\col" mesh="shoulder_link_convex_decomposition_p1"/>
                <geom name="//unnamed_geom_20" class="wam\col" mesh="shoulder_link_convex_decomposition_p2"/>
                <geom name="//unnamed_geom_21" class="wam\col" mesh="shoulder_link_convex_decomposition_p3"/>
            <--/body-->
            <body name="wam_7dof_bhand_0\wam\shoulder_yaw_link\wam\shoulder_pitch_link\" pos="0 0 0" quat="0.70710700000000004 -0.70710700000000004 0 0">
              <--body name="wam_7dof_bhand_0\wam\shoulder_yaw_link\wam\shoulder_pitch_link\wam\shoulder_pitch_link"-->
                  <inertial pos="-0.0023698299999999999 -0.01542114 0.031056139999999999" mass="3.8748999999999998" fullinertia="0.021409580000000001 0.01558906 0.013778749999999999 -2.461e-05 0.00027171999999999998 0.0018192"/>
                  <joint name="wam\J2" class="rand" axis="0 0 1" range="-2 2" damping="1000" frictionloss="1000"/>
                  <geom name="//unnamed_geom_22" class="wam\shoulder_pitch_link\viz" mesh="shoulder_pitch_link_fine"/>
                  <geom name="//unnamed_geom_23" class="wam\col" mesh="shoulder_pitch_link_convex"/>
              <--/body-->

            </body>
        </body>
 </body>

Can we make changes that does not create a new body-element, and, instead, appends the the site's pos and quat information to the child's body-element, which is defined inside the child's worldbody element? This will also address the issue here as well

ajoardar007 commented 3 months ago

In the last two-lines of the Attachment-frame topic, it is mentioned that joint is a valid child of attachment-frames.

Apparently, "The easiest way to do this is to hold a reference to the object returned by a call to attach":

frame_1 = self.shoulder_pitch_link_site.attach(self.shoulder_pitch_link.mjcf_model)

However, I am unable to attach a joint at the attachment-frame:

# Add joint to the attachment frame
        frame_1.add('joint',
                    name = self._joint_name[1],
                    range = self._joint_range[1], 
                    damping = self._joint_damping[1],
                    axis = self._joint_axis,
                    type = self._joint_type,
                    frictionloss = self._joint_frictionloss,
                    dclass = 'rand')

while I get the following error:

Traceback (most recent call last):
  File "/home/arnab/.local/lib/python3.8/site-packages/dm_control/mjcf/element.py", line 593, in _check_valid_child
    return self._spec.children[element_name]
KeyError: 'joint'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/arnab/UWARL_catkin_ws/src/uwarl-mujoco-summit-wam-sim/src/main_rob_dup_1.py", line 965, in <module>
    summit_body = Summit(prefix = prefix)
  File "/home/arnab/UWARL_catkin_ws/src/uwarl-mujoco-summit-wam-sim/src/main_rob_dup_1.py", line 916, in __init__
    wam = WAM(name = self._wam_name+'_'+prefix)
  File "/home/arnab/UWARL_catkin_ws/src/uwarl-mujoco-summit-wam-sim/src/main_rob_dup_1.py", line 471, in __init__
    self.shoulder_yaw_link_site.add('joint',
  File "/home/arnab/.local/lib/python3.8/site-packages/dm_control/mjcf/element.py", line 630, in add
    return self.insert(element_name, position=None, **kwargs)
  File "/home/arnab/.local/lib/python3.8/site-packages/dm_control/mjcf/element.py", line 647, in insert
    child_spec = self._check_valid_child(element_name)
  File "/home/arnab/.local/lib/python3.8/site-packages/dm_control/mjcf/element.py", line 595, in _check_valid_child
    raise AttributeError(  # pylint: disable=raise-missing-from
AttributeError: <joint> is not a valid child of <site>

Could you please advise as to how I may add a joint-element at the attachment-frame between a parent and a child?