def init_sim()

in bisk/tasks/polebalance.py [0:0]

• 94 lines of code
• 4 McCabe index (conditional complexity)

    def init_sim(self, root: mjcf.RootElement, frameskip: int = 5):
        try:
            from matplotlib import pyplot as plt

            cmap = plt.get_cmap('rainbow')
        except:
            cmap = lambda x: [1, 0, 0, 1]

        size = 0.05
        if self.robot in {'humanoid', 'humanoidpc'}:
            size = 0.02
            head = root.find('body', 'robot/head')
            headg = head.find('geom', 'head')
            zpos = headg.size[0]
            pole = head.add('body', name='pole-0', pos=[0, 0, zpos])
        elif self.robot in {'halfcheetah'}:
            torso = root.find('body', 'robot/torso')
            pole = torso.add('body', name='pole-0', pos=[0, 0, 0])
        elif self.robot in {'walker'}:
            torso = root.find('body', 'robot/torso')
            torsog = torso.find('geom', 'torso')
            pole = torso.add('body', name='pole-0', pos=[0, 0, torsog.size[1]])
        else:
            try:
                torso = root.find('body', 'robot/torso')
                zpos = torso.find('geom', 'torso').size[2]
                pole = torso.add('body', name='pole-0', pos=[0, 0, zpos])
            except:
                raise NotImplementedError(
                    f'Don\'t know how to place poles on a {self.robot} robot'
                )

        if self.robot in {'halfcheetah', 'walker'}:
            # HalfCheetah model is defined in radians
            limit = np.pi if self.robot == 'halfcheetah' else 180
            pole.add(
                'joint',
                name='pole-0_rot',
                type='hinge',
                damping=0.1,
                stiffness=0,
                axis='0 1 0',
                pos=[0, 0, 0],
                range=[-limit, limit],
            )
        else:
            pole.add(
                'joint',
                name='pole-0_rot',
                damping=0.1,
                type='ball',
                pos=[0, 0, 0],
                range=[0, 90],
            )
        pole.add(
            'geom',
            name='pole-0_geom',
            type='capsule',
            fromto=[0, 0, 0, 0, 0, self.pole_length],
            size=[size],
            mass=self.pole_mass,
            rgba=cmap(0),
        )

        for i in range(1, self.n_poles):
            pole = pole.add(
                'body', name=f'pole-{i}', pos=[0, 0, self.pole_length]
            )
            if self.robot in {'halfcheetah', 'walker'}:
                limit = np.pi if self.robot == 'halfcheetah' else 180
                pole.add(
                    'joint',
                    name=f'pole-{i}_rot',
                    type='hinge',
                    damping=0.1,
                    stiffness=0,
                    axis='0 1 0',
                    pos=[0, 0, 0],
                    range=[-limit, limit],
                )
            else:
                pole.add(
                    'joint',
                    name=f'pole-{i}_rot',
                    type='ball',
                    damping=0.1,
                    pos=[0, 0, 0],
                    range=[0, 90],
                )
            pole.add(
                'geom',
                name=f'pole-{i}_geom',
                type='capsule',
                fromto=[0, 0, 0, 0, 0, self.pole_length],
                size=[size],
                mass=self.pole_mass,
                rgba=cmap((i + 1) / self.n_poles),
            )

        super().init_sim(root, frameskip)