def _get_simulation_reward_with_done()

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

• 19 lines of code
• 8 McCabe index (conditional complexity)

    def _get_simulation_reward_with_done(self, info: dict) -> Tuple[float, bool]:
        reward, done = super()._get_simulation_reward_with_done(info)

        # If there is a gripper to table contact, apply table_collision penalty
        table_penalty = self.parameters.simulation_params.penalty.get(
            "table_collision", 0.0
        )

        if self.mujoco_simulation.get_gripper_table_contact():
            reward -= table_penalty

        # Add a large negative penalty for "breaking" the wrist camera by hitting it
        # with the table or another object.
        if info.get("wrist_cam_contacts", {}).get("any", False):
            reward -= self.parameters.simulation_params.penalty.get(
                "wrist_collision", 0.0
            )

        if "objects_off_table" in info and info["objects_off_table"].any():
            # If any object is off the table, we will end the episode
            done = True
            # Add a penalty to letting an object go off the table
            reward -= self.parameters.simulation_params.penalty.get(
                "objects_off_table", 0.0
            )
        if any(self.observe()["safety_stop"]):
            reward -= self.parameters.simulation_params.penalty.get("safety_stop", 0.0)
        return reward, done