def _act()

in robogym/envs/rearrange/common/base.py [0:0]

• 34 lines of code
• 7 McCabe index (conditional complexity)

    def _act(self, action):
        if self.constants.teleport_to_goal and self.t > 10:
            # Override the policy by teleporting directly to the goal state (used to generate
            # a balanced distribution for training a goal classifier). Add some noise to the
            # objects' states; we tune the scale of the noise so that with probability p=0.5,
            # all objects are within the success_threshold. This will result in our episodes
            # containing a ~p/(1+p) fraction of goal-achieved states.
            target_pos = self.mujoco_simulation.get_target_pos(pad=False)
            target_quat = self.mujoco_simulation.get_target_quat(pad=False)

            num_objects = target_pos.shape[0]
            num_randomizations = num_objects * (
                ("obj_pos" in self.constants.success_threshold)
                + ("obj_rot" in self.constants.success_threshold)
            )
            assert num_randomizations > 0
            success_prob = 0.5 ** (1 / num_randomizations)

            # Add Gaussian noise to x and y position.
            if "obj_pos" in self.constants.success_threshold:
                # Tune the noise so that the position is within success_threshold with
                # probability success_prob. Note for example that noise_scale -> 0 when
                # success_prob -> 1, and noise_scale -> infinity when success_prob -> 0.
                noise_scale = np.ones_like(target_pos)
                noise_scale *= self.constants.success_threshold["obj_pos"]
                noise_scale /= np.sqrt(-2 * np.log(1 - success_prob))
                noise_scale[:, 2] = 0.0  # Don't add noise to the z-axis
                target_pos = np.random.normal(loc=target_pos, scale=noise_scale)

            # Add Gaussian noise to rotation about z-axis.
            if "obj_rot" in self.constants.success_threshold:
                # Tune the noise so that the rotation is within success_threshold with
                # probability success_prob. Note for example that noise_scale -> 0 when
                # success_prob -> 1, and noise_scale -> infinity when success_prob -> 0.
                noise_scale = self.constants.success_threshold["obj_rot"]
                noise_scale /= scipy.special.ndtri(
                    success_prob + (1 - success_prob) / 2
                )
                noise_quat = rotation.quat_from_angle_and_axis(
                    angle=np.random.normal(
                        loc=np.zeros((num_objects,)), scale=noise_scale
                    ),
                    axis=np.array([[0, 0, 1.0]] * num_objects),
                )
                target_quat = rotation.quat_mul(target_quat, noise_quat)

            self.mujoco_simulation.set_object_pos(target_pos)
            self.mujoco_simulation.set_object_quat(target_quat)
            self.mujoco_simulation.forward()
        else:
            self._set_action(action)