Skip to content

heavy_robot_wrapper

HeavyRobotWrapper

Bases: EnvironmentWrapper

Parameters:

Name Type Description Default
env Environment

The environment to wrap.

required
Source code in OmniGibson/omnigibson/learning/wrappers/heavy_robot_wrapper.py
class HeavyRobotWrapper(EnvironmentWrapper):
    """
    Args:
        env (og.Environment): The environment to wrap.
    """

    def __init__(self, env: Environment):
        super().__init__(env=env)
        # Here, we modify the robot observation to  use 224 * 224 resolution
        # For a complete list of available modalities, see VisionSensor.ALL_MODALITIES
        # We also change the robot base mass to 250kg to match the configuration during data collection.
        robot = env.robots[0]
        with og.sim.stopped():
            robot.base_footprint_link.mass = 250.0  # increase base mass to 250kg
        # Update robot sensors:
        for camera_id, camera_name in ROBOT_CAMERA_NAMES["R1Pro"].items():
            sensor_name = camera_name.split("::")[1]
            if camera_id == "head":
                robot.sensors[sensor_name].horizontal_aperture = 40.0  # this is what we used in data collection
            robot.sensors[sensor_name].image_height = 224
            robot.sensors[sensor_name].image_width = 224
        # reload observation space
        env.load_observation_space()
        logger.info("Reloaded observation space!")