def _randomize_cube_initial_position()

in robogym/envs/dactyl/full_perpendicular.py [0:0]

• 35 lines of code
• 5 McCabe index (conditional complexity)

    def _randomize_cube_initial_position(self):
        """ Draw a random initial position for a cube """
        # Original env had this, but I'm not really sure it's needed
        for i in range(self.constants.reset_initial_steps):
            ctrl = self.mujoco_simulation.shadow_hand.denormalize_position_control(
                self.mujoco_simulation.shadow_hand.zero_control()
            )
            self.mujoco_simulation.shadow_hand.set_position_control(ctrl)
            self.mujoco_simulation.step()

        cube_translation = (
            self._random_state.randn(3) * self.parameters.cube_position_wiggle_std
        )
        self.mujoco_simulation.add_qpos("cube_position", cube_translation)

        cube_orientation = rotation.uniform_quat(self._random_state)
        self.mujoco_simulation.set_qpos("cube_rotation", cube_orientation)

        self._scramble_cube()

        if self.constants.scramble_face_angles:
            self._scramble_face_angles()

        if self.constants.randomize_face_angles:
            # Face angles
            random_face_angle = self._random_state.uniform(
                -np.pi / 4, np.pi / 4, size=2
            )
            # Face axes
            random_axis = self._random_state.randint(3)

            self.mujoco_simulation.cube_model.rotate_face(
                random_axis, 0, random_face_angle[0]
            )

            self.mujoco_simulation.cube_model.rotate_face(
                random_axis, 1, random_face_angle[1]
            )

        # Need to call this after the qpos is modified
        self.mujoco_simulation.forward()

        action = self._random_state.uniform(-1.0, 1.0, self.action_space.shape[0])

        for _ in range(self.parameters.n_random_initial_steps):
            ctrl = self.mujoco_simulation.shadow_hand.denormalize_position_control(
                action
            )
            self.mujoco_simulation.shadow_hand.set_position_control(ctrl)
            self.mujoco_simulation.step()