How can I build the model dynamically like dm_control does? NO mjcf NO URDF.

[baixianger]

I have to build my environment from Brax. But the thing is my model has a lot of physical parameters to change. Is there a way in Brax to build the model from code, instead of from the XML file?

Here i what i did in dm_control mjcf submodule. I want to simulate a whip, so I cut the whip into many serial joints and geoms. But i don't konw how many joints i need. if I implement it in XML, every time I want to make a change, I have to nest the xml elements agian and again.

class Whip(composer.Entity):
    """A Whip Entity."""
    def _build(self, *args, **kwargs):
        """Initializes the whip."""
        self._model = self._make_model(**kwargs)
        self._whip_begin = self._model.find('body', 'B0')
        self._whip_end = self._model.find('body', 'whip_end')
        self._whip_joints = self._model.find_all('joint')

    def _build_observables(self):
        """Returns the observables for the whip."""
        return WhipObservables(self)

    @property
    def mjcf_model(self):
        return self._model

    @property
    def whip_end(self):
        """Returns the whip end."""
        return self._whip_end

    @property
    def whip_begin(self):
        """Returns the whip begin."""
        return self._whip_begin

    def update_whip_color(self, physics):
        """Updates the activation of the whip."""
        physics.bind(self._model.find_all('geom')).rgba = [1, 0, 0, 1]

    def _make_model(self, wtype=0):
        """Returns an MJCF model of a Whip."""
        if wtype not in [0, 1]:
            raise ValueError("whip type must be 0 or 1")

        model = mjcf.RootElement(model='whip')
        # set compiler
        model.compiler.angle = 'radian'
        model.compiler.autolimits = True
        model.option.integrator = "implicit"
        model.option.impratio = "10"

        whip_material = model.asset.add('material', name="white", rgba="1 1 1 1")
        whip = model.worldbody.add('body', name="whip")
        whip.add('geom', type="sphere", size=[0.045,], material=whip_material)
        temp_body = whip.add('body', name='B0', pos=[0, 0, 0.045])

        for i in range(_COUNT):
            temp_body.add('joint', name=f'J0_{i}', **whip_params['joint0'])
            temp_body.add('joint', name=f'J1_{i}', **whip_params['joint1'])
            if wtype == 0:
                temp_body.add('geom',  name=f'G0_{i}',
                              material=whip_material, **whip_params['sphere'])
                temp_body.add('geom',  name=f'G1_{i}',
                              material=whip_material, **whip_params['cylinder'])
            elif wtype == 1:
                temp_body.add('geom',  name=f'G0_{i}',
                              material=whip_material, **whip_params['capsule'])
            name = f'B{i+1}' if i != _COUNT - 1 else 'whip_end'
            temp_body = temp_body.add('body',  name=name, pos=[0, 0, _UNIT_L])

        return model

Do brax provide the same methods to create model by code? Thank you advance.

[RolandZhu]

Hi, hope I'm not too late. I have implemented some functions in v1 to generate the simulation environment procedurally, which might be helpful for you. See here.

[baixianger]

Hi, hope I'm not too late. I have implemented some functions in v1 to generate the simulation environment procedurally, which might be helpful for you. See here.

That's so cool. That would be a major help for me. Thank you!

[btaba]

Hi @baixianger , we haven't implemented a "composer" class in v2, but if you're interested in implementing it, the main fields that need constructing are here in System

https://github.com/google/brax/blob/3b72c952f7c990e8e9e355f736e77f10f408bc14/brax/io/mjcf.py#L461C2-L464

and these https://github.com/google/brax/blob/3b72c952f7c990e8e9e355f736e77f10f408bc14/brax/io/mjcf.py#L475-L478