Skip to content

submission_comet

CometWrapper

Bases: EnvironmentWrapper

Parameters:

Name Type Description Default
env Environment

The environment to wrap.

 required
Source code in OmniGibson/omnigibson/learning/wrappers/challenge_submissions/submission_comet.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class CometWrapper(EnvironmentWrapper):
    """
    Args:
        env (og.Environment): The environment to wrap.
    """

    def __init__(self, env: Environment):
        super().__init__(env=env)
        # Note that from eval.py we only set rgb modality, here we include more (depth + seg_instance_id)
        # Here, we change the camera resolution and head camera aperture to match the one we used in data collection
        robot = env.robots[0]
        # 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
                robot.sensors[sensor_name].image_height = HEAD_RESOLUTION[0]
                robot.sensors[sensor_name].image_width = HEAD_RESOLUTION[1]
            else:
                robot.sensors[sensor_name].image_height = WRIST_RESOLUTION[0]
                robot.sensors[sensor_name].image_width = WRIST_RESOLUTION[1]
            # add depth and segmentation
            # robot.sensors[sensor_name].add_modality("depth_linear")
            # robot.sensors[sensor_name].add_modality("seg_semantic")
            # robot.sensors[sensor_name].add_modality("seg_instance_id")
        # reload observation space
        env.load_observation_space()
        logger.info("Reloaded observation space!")