コード例 #1
0
    def make_sampler_fn(
        cls,
        stage: str,
        force_cache_reset: bool,
        allowed_scenes: Optional[Sequence[str]],
        seed: int,
        epochs: Union[str, float, int],
        scene_to_allowed_rearrange_inds: Optional[Dict[str,
                                                       Sequence[int]]] = None,
        x_display: Optional[str] = None,
        sensors: Optional[Sequence[Sensor]] = None,
        only_one_unshuffle_per_walkthrough: bool = False,
        thor_controller_kwargs: Optional[Dict] = None,
        **kwargs,
    ) -> RearrangeTaskSampler:
        """Return a RearrangeTaskSampler."""
        if "mp_ctx" in kwargs:
            del kwargs["mp_ctx"]
        assert not cls.RANDOMIZE_START_ROTATION_DURING_TRAINING

        return RearrangeTaskSampler.from_fixed_dataset(
            run_walkthrough_phase=True,
            run_unshuffle_phase=True,
            stage=stage,
            allowed_scenes=allowed_scenes,
            scene_to_allowed_rearrange_inds=scene_to_allowed_rearrange_inds,
            rearrange_env_kwargs=dict(
                force_cache_reset=force_cache_reset,
                **cls.REARRANGE_ENV_KWARGS,
                controller_kwargs={
                    "x_display":
                    x_display,
                    **cls.THOR_CONTROLLER_KWARGS,
                    **({} if thor_controller_kwargs is None else thor_controller_kwargs),
                    "renderDepthImage":
                    any(isinstance(s, DepthSensor) for s in cls.SENSORS),
                },
            ),
            seed=seed,
            sensors=SensorSuite(cls.SENSORS)
            if sensors is None else SensorSuite(sensors),
            max_steps=cls.MAX_STEPS,
            discrete_actions=cls.actions(),
            require_done_action=cls.REQUIRE_DONE_ACTION,
            force_axis_aligned_start=cls.FORCE_AXIS_ALIGNED_START,
            unshuffle_runs_per_walkthrough=cls.
            TRAIN_UNSHUFFLE_RUNS_PER_WALKTHROUGH if
            (not only_one_unshuffle_per_walkthrough) and stage == "train" else
            None,
            epochs=epochs,
            **kwargs,
        )
コード例 #2
0
    def machine_params(self, mode="train", **kwargs):
        if mode == "train":
            workers_per_device = 1
            gpu_ids = ([] if not torch.cuda.is_available() else
                       self.TRAINING_GPUS * workers_per_device)
            nprocesses = (1 if not torch.cuda.is_available() else
                          evenly_distribute_count_into_bins(
                              self.NUM_PROCESSES, len(gpu_ids)))
        elif mode == "valid":
            nprocesses = 1
            gpu_ids = [] if not torch.cuda.is_available(
            ) else self.VALIDATION_GPUS
        elif mode == "test":
            nprocesses = 1
            gpu_ids = [] if not torch.cuda.is_available(
            ) else self.TESTING_GPUS
        else:
            raise NotImplementedError(
                "mode must be 'train', 'valid', or 'test'.")

        sensor_preprocessor_graph = (SensorPreprocessorGraph(
            source_observation_spaces=SensorSuite(
                self.SENSORS).observation_spaces,
            preprocessors=self.PREPROCESSORS,
        ) if mode == "train" or (
            (isinstance(nprocesses, int) and nprocesses > 0) or
            (isinstance(nprocesses, Sequence) and sum(nprocesses) > 0)) else
                                     None)

        return MachineParams(
            nprocesses=nprocesses,
            devices=gpu_ids,
            sensor_preprocessor_graph=sensor_preprocessor_graph,
        )
コード例 #3
0
    def create_model(cls, **kwargs) -> WalkthroughActorCriticResNetWithPassiveMap:
        map_sensor = cast(
            BinnedPointCloudMapTHORSensor,
            next(
                s for s in cls.SENSORS if isinstance(s, BinnedPointCloudMapTHORSensor)
            ),
        )
        map_kwargs = dict(
            frame_height=224,
            frame_width=224,
            vision_range_in_cm=map_sensor.vision_range_in_cm,
            resolution_in_cm=map_sensor.resolution_in_cm,
            map_size_in_cm=map_sensor.map_size_in_cm,
        )

        observation_space = (
            SensorSuite(cls.SENSORS).observation_spaces
            if kwargs.get("sensor_preprocessor_graph") is None
            else kwargs["sensor_preprocessor_graph"].observation_spaces
        )

        return WalkthroughActorCriticResNetWithPassiveMap(
            action_space=gym.spaces.Discrete(len(cls.actions())),
            observation_space=observation_space,
            rgb_uuid=cls.EGOCENTRIC_RGB_UUID,
            unshuffled_rgb_uuid=cls.UNSHUFFLED_RGB_UUID,
            semantic_map_channels=len(cls.ORDERED_OBJECT_TYPES),
            height_map_channels=3,
            map_kwargs=map_kwargs,
        )
コード例 #4
0
    def resnet_preprocessor_graph(cls, mode: str) -> SensorPreprocessorGraph:
        def create_resnet_builder(in_uuid: str, out_uuid: str):
            return ResNetPreprocessor(
                input_height=cls.THOR_CONTROLLER_KWARGS["height"],
                input_width=cls.THOR_CONTROLLER_KWARGS["width"],
                output_width=7,
                output_height=7,
                output_dims=512,
                pool=False,
                torchvision_resnet_model=torchvision.models.resnet18,
                input_uuids=[in_uuid],
                output_uuid=out_uuid,
            )

        img_uuids = [cls.EGOCENTRIC_RGB_UUID, cls.UNSHUFFLED_RGB_UUID]
        return SensorPreprocessorGraph(
            source_observation_spaces=SensorSuite([
                sensor for sensor in cls.SENSORS if
                (mode == "train" or not isinstance(sensor, ExpertActionSensor))
            ]).observation_spaces,
            preprocessors=[
                create_resnet_builder(sid, f"{sid}_resnet")
                for sid in img_uuids
            ],
        )
コード例 #5
0
    def create_model(cls, **kwargs) -> nn.Module:
        def get_sensor_uuid(stype: Type[Sensor]) -> Optional[str]:
            s = next((s for s in cls.SENSORS if isinstance(s, stype)), None,)
            return None if s is None else s.uuid

        walkthrougher_should_ignore_action_mask = [
            any(k in a for k in ["drop", "open", "pickup"]) for a in cls.actions()
        ]

        if cls.CNN_PREPROCESSOR_TYPE_AND_PRETRAINING is None:
            return TwoPhaseRearrangeActorCriticSimpleConvRNN(
                action_space=gym.spaces.Discrete(len(cls.actions())),
                observation_space=SensorSuite(cls.SENSORS).observation_spaces,
                rgb_uuid=cls.EGOCENTRIC_RGB_UUID,
                unshuffled_rgb_uuid=cls.UNSHUFFLED_RGB_UUID,
                in_walkthrough_phase_uuid=get_sensor_uuid(InWalkthroughPhaseSensor),
                is_walkthrough_phase_embedding_dim=cls.IS_WALKTHROUGH_PHASE_EMBEDING_DIM,
                rnn_type=cls.RNN_TYPE,
                walkthrougher_should_ignore_action_mask=walkthrougher_should_ignore_action_mask,
                done_action_index=cls.actions().index("done"),
            )
        else:
            return ResNetTwoPhaseRearrangeActorCriticRNN(
                action_space=gym.spaces.Discrete(len(cls.actions())),
                observation_space=kwargs[
                    "sensor_preprocessor_graph"
                ].observation_spaces,
                rgb_uuid=cls.EGOCENTRIC_RGB_RESNET_UUID,
                unshuffled_rgb_uuid=cls.UNSHUFFLED_RGB_RESNET_UUID,
                in_walkthrough_phase_uuid=get_sensor_uuid(InWalkthroughPhaseSensor),
                is_walkthrough_phase_embedding_dim=cls.IS_WALKTHROUGH_PHASE_EMBEDING_DIM,
                rnn_type=cls.RNN_TYPE,
                walkthrougher_should_ignore_action_mask=walkthrougher_should_ignore_action_mask,
                done_action_index=cls.actions().index("done"),
            )
コード例 #6
0
ファイル: gym_tutorial.py プロジェクト: elliotthwang/allenact
 def create_model(cls, **kwargs) -> nn.Module:
     return MemorylessActorCritic(
         input_uuid="gym_box_data",
         action_space=gym.spaces.Box(
             -1.0, 1.0, (2, )),  # 2 actors, each in the range [-1.0, 1.0]
         observation_space=SensorSuite(cls.SENSORS).observation_spaces,
         action_std=0.5,
     )
コード例 #7
0
 def create_model(cls, **kwargs) -> nn.Module:
     return MiniGridSimpleConv(
         action_space=gym.spaces.Discrete(len(MiniGridTask.class_action_names())),
         observation_space=SensorSuite(cls.SENSORS).observation_spaces,
         num_objects=cls.SENSORS[0].num_objects,
         num_colors=cls.SENSORS[0].num_colors,
         num_states=cls.SENSORS[0].num_states,
     )
コード例 #8
0
 def create_model(cls, **kwargs) -> nn.Module:
     return ConditionedMiniGridSimpleConvRNN(
         action_space=gym.spaces.Dict(higher=gym.spaces.Discrete(2),
                                      lower=gym.spaces.Discrete(2)),
         observation_space=SensorSuite(cls.SENSORS).observation_spaces,
         num_objects=cls.SENSORS[0].num_objects,
         num_colors=cls.SENSORS[0].num_colors,
         num_states=cls.SENSORS[0].num_states,
     )
コード例 #9
0
 def create_model(cls, **kwargs) -> nn.Module:
     sensors = cls.get_sensors()
     return BabyAIRecurrentACModel(
         action_space=gym.spaces.Discrete(
             len(BabyAITask.class_action_names())),
         observation_space=SensorSuite(sensors).observation_spaces,
         use_instr=cls.USE_INSTR,
         use_memory=True,
         arch=cls.ARCH,
     )
コード例 #10
0
    def __init__(
        self,
        world_dim: int,
        world_radius: int,
        sensors: Union[SensorSuite, List[Sensor]],
        max_steps: int,
        max_tasks: Optional[int] = None,
        num_unique_seeds: Optional[int] = None,
        task_seeds_list: Optional[List[int]] = None,
        deterministic_sampling: bool = False,
        seed: Optional[int] = None,
        **kwargs,
    ):
        self.env = LightHouseEnvironment(world_dim=world_dim,
                                         world_radius=world_radius)

        self._last_sampled_task: Optional[FindGoalLightHouseTask] = None
        self.sensors = (SensorSuite(sensors)
                        if not isinstance(sensors, SensorSuite) else sensors)
        self.max_steps = max_steps
        self.max_tasks = max_tasks
        self.num_tasks_generated = 0
        self.deterministic_sampling = deterministic_sampling

        self.num_unique_seeds = num_unique_seeds
        self.task_seeds_list = task_seeds_list
        assert (self.num_unique_seeds is
                None) or (0 < self.num_unique_seeds
                          ), "`num_unique_seeds` must be a positive integer."

        self.num_unique_seeds = num_unique_seeds
        self.task_seeds_list = task_seeds_list
        if self.task_seeds_list is not None:
            if self.num_unique_seeds is not None:
                assert self.num_unique_seeds == len(
                    self.task_seeds_list
                ), "`num_unique_seeds` must equal the length of `task_seeds_list` if both specified."
            self.num_unique_seeds = len(self.task_seeds_list)
        elif self.num_unique_seeds is not None:
            self.task_seeds_list = list(range(self.num_unique_seeds))

        assert (not deterministic_sampling) or (
            self.num_unique_seeds is not None
        ), "Cannot use deterministic sampling when `num_unique_seeds` is `None`."

        if (not deterministic_sampling) and self.max_tasks:
            get_logger().warning(
                "`deterministic_sampling` is `False` but you have specified `max_tasks < inf`,"
                " this might be a mistake when running testing.")

        self.seed: int = int(
            seed if seed is not None else np.random.randint(0, 2**31 - 1))
        self.np_seeded_random_gen: Optional[np.random.RandomState] = None
        self.set_seed(self.seed)
コード例 #11
0
 def create_model(cls, **kwargs) -> nn.Module:
     return ObjectNavBaselineActorCritic(
         action_space=gym.spaces.Discrete(
             len(ObjectNaviThorGridTask.class_action_names())),
         observation_space=SensorSuite(cls.SENSORS).observation_spaces,
         rgb_uuid=cls.SENSORS[0].uuid,
         depth_uuid=None,
         goal_sensor_uuid="goal_object_type_ind",
         hidden_size=512,
         object_type_embedding_dim=8,
     )
コード例 #12
0
 def create_model(cls, **kwargs) -> nn.Module:
     return NavToPartnerActorCriticSimpleConvRNN(
         action_space=gym.spaces.Tuple(
             [
                 gym.spaces.Discrete(len(NavToPartnerTask.class_action_names())),
                 gym.spaces.Discrete(len(NavToPartnerTask.class_action_names())),
             ]
         ),
         observation_space=SensorSuite(cls.SENSORS).observation_spaces,
         hidden_size=512,
     )
コード例 #13
0
ファイル: task.py プロジェクト: apoorvkh/allenact
 def __init__(self, env: EnvType, sensors: Union[SensorSuite,
                                                 Sequence[Sensor]],
              task_info: Dict[str, Any], max_steps: int, **kwargs) -> None:
     self.env = env
     self.sensor_suite = (SensorSuite(sensors) if
                          not isinstance(sensors, SensorSuite) else sensors)
     self.task_info = task_info
     self.max_steps = max_steps
     self.observation_space = self.sensor_suite.observation_spaces
     self._num_steps_taken = 0
     self._total_reward: Union[float, List[float]] = 0.0
コード例 #14
0
    def __init__(
        self,
        env_builder: Union[str, Callable[..., MiniGridEnv]],
        sensors: Union[SensorSuite, List[Sensor]],
        max_tasks: Optional[int] = None,
        num_unique_seeds: Optional[int] = None,
        task_seeds_list: Optional[List[int]] = None,
        deterministic_sampling: bool = False,
        extra_task_kwargs: Optional[Dict] = None,
        **kwargs,
    ):
        super(BabyAITaskSampler, self).__init__()
        self.sensors = (SensorSuite(sensors)
                        if not isinstance(sensors, SensorSuite) else sensors)
        self.max_tasks = max_tasks
        self.num_unique_seeds = num_unique_seeds
        self.deterministic_sampling = deterministic_sampling
        self.extra_task_kwargs = (extra_task_kwargs
                                  if extra_task_kwargs is not None else {})

        self._last_env_seed: Optional[int] = None
        self._last_task: Optional[BabyAITask] = None

        assert (self.num_unique_seeds is
                None) or (0 < self.num_unique_seeds
                          ), "`num_unique_seeds` must be a positive integer."

        self.num_unique_seeds = num_unique_seeds
        self.task_seeds_list = task_seeds_list
        if self.task_seeds_list is not None:
            if self.num_unique_seeds is not None:
                assert self.num_unique_seeds == len(
                    self.task_seeds_list
                ), "`num_unique_seeds` must equal the length of `task_seeds_list` if both specified."
            self.num_unique_seeds = len(self.task_seeds_list)
        elif self.num_unique_seeds is not None:
            self.task_seeds_list = list(range(self.num_unique_seeds))

        if (not deterministic_sampling) and self.max_tasks:
            get_logger().warning(
                "`deterministic_sampling` is `False` but you have specified `max_tasks < inf`,"
                " this might be a mistake when running testing.")

        if isinstance(env_builder, str):
            self.env = gym.make(env_builder)
        else:
            self.env = env_builder()

        self.np_seeded_random_gen, _ = seeding.np_random(
            random.randint(0, 2**31 - 1))
        self.num_tasks_generated = 0
コード例 #15
0
    def resnet_preprocessor_graph(cls, mode: str) -> SensorPreprocessorGraph:
        def create_resnet_builder(in_uuid: str, out_uuid: str):
            cnn_type, pretraining_type = cls.CNN_PREPROCESSOR_TYPE_AND_PRETRAINING
            if pretraining_type == "imagenet":
                assert cnn_type in [
                    "RN18",
                    "RN50",
                ], "Only allow using RN18/RN50 with `imagenet` pretrained weights."
                return ResNetPreprocessor(
                    input_height=cls.THOR_CONTROLLER_KWARGS["height"],
                    input_width=cls.THOR_CONTROLLER_KWARGS["width"],
                    output_width=7,
                    output_height=7,
                    output_dims=512 if "18" in cnn_type else 2048,
                    pool=False,
                    torchvision_resnet_model=getattr(
                        torchvision.models,
                        f"resnet{cnn_type.replace('RN', '')}"),
                    input_uuids=[in_uuid],
                    output_uuid=out_uuid,
                )
            elif pretraining_type == "clip":
                from allenact_plugins.clip_plugin.clip_preprocessors import (
                    ClipResNetPreprocessor, )
                import clip

                # Let's make sure we download the clip model now
                # so we don't download it on every spawned process
                clip.load(cnn_type, "cpu")

                return ClipResNetPreprocessor(
                    rgb_input_uuid=in_uuid,
                    clip_model_type=cnn_type,
                    pool=False,
                    output_uuid=out_uuid,
                )
            else:
                raise NotImplementedError

        img_uuids = [cls.EGOCENTRIC_RGB_UUID, cls.UNSHUFFLED_RGB_UUID]
        return SensorPreprocessorGraph(
            source_observation_spaces=SensorSuite([
                sensor for sensor in cls.sensors() if
                (mode == "train" or not isinstance(sensor, ExpertActionSensor))
            ]).observation_spaces,
            preprocessors=[
                create_resnet_builder(sid, f"{sid}_resnet")
                for sid in img_uuids
            ],
        )
コード例 #16
0
 def create_model(cls, **kwargs) -> nn.Module:
     sensors = cls.get_sensors()
     return BabyAIRecurrentACModel(
         action_space=gym.spaces.Discrete(
             len(BabyAITask.class_action_names())),
         observation_space=SensorSuite(sensors).observation_spaces,
         use_instr=cls.USE_INSTR,
         use_memory=True,
         arch=cls.ARCH,
         instr_dim=256,
         lang_model="attgru",
         memory_dim=2048,
         include_auxiliary_head=cls.INCLUDE_AUXILIARY_HEAD,
     )
コード例 #17
0
    def machine_params(self, mode="train", **kwargs):
        sampler_devices: Sequence[int] = []
        if mode == "train":
            workers_per_device = 1
            gpu_ids = (
                []
                if not torch.cuda.is_available()
                else self.TRAIN_GPU_IDS * workers_per_device
            )
            nprocesses = (
                1
                if not torch.cuda.is_available()
                else evenly_distribute_count_into_bins(self.NUM_PROCESSES, len(gpu_ids))
            )
            sampler_devices = self.SAMPLER_GPU_IDS
        elif mode == "valid":
            nprocesses = 1
            gpu_ids = [] if not torch.cuda.is_available() else self.VALID_GPU_IDS
        elif mode == "test":
            nprocesses = 5 if torch.cuda.is_available() else 1
            gpu_ids = [] if not torch.cuda.is_available() else self.TEST_GPU_IDS
        else:
            raise NotImplementedError("mode must be 'train', 'valid', or 'test'.")

        sensors = [*self.SENSORS]
        if mode != "train":
            sensors = [s for s in sensors if not isinstance(s, ExpertActionSensor)]

        sensor_preprocessor_graph = (
            SensorPreprocessorGraph(
                source_observation_spaces=SensorSuite(sensors).observation_spaces,
                preprocessors=self.preprocessors(),
            )
            if mode == "train"
            or (
                (isinstance(nprocesses, int) and nprocesses > 0)
                or (isinstance(nprocesses, Sequence) and sum(nprocesses) > 0)
            )
            else None
        )

        return MachineParams(
            nprocesses=nprocesses,
            devices=gpu_ids,
            sampler_devices=sampler_devices
            if mode == "train"
            else gpu_ids,  # ignored with > 1 gpu_ids
            sensor_preprocessor_graph=sensor_preprocessor_graph,
        )
コード例 #18
0
 def create_model(cls, **kwargs) -> nn.Module:
     if not cls.USE_RESNET_CNN:
         return RearrangeActorCriticSimpleConvRNN(
             action_space=gym.spaces.Discrete(len(cls.actions())),
             observation_space=SensorSuite(cls.SENSORS).observation_spaces,
             rgb_uuid=cls.EGOCENTRIC_RGB_UUID,
             unshuffled_rgb_uuid=cls.UNSHUFFLED_RGB_UUID,
         )
     else:
         return ResNetRearrangeActorCriticRNN(
             action_space=gym.spaces.Discrete(len(cls.actions())),
             observation_space=kwargs["sensor_preprocessor_graph"].
             observation_spaces,
             rgb_uuid=cls.EGOCENTRIC_RGB_RESNET_UUID,
             unshuffled_rgb_uuid=cls.UNSHUFFLED_RGB_RESNET_UUID,
         )
コード例 #19
0
    def create_model(cls, **kwargs) -> nn.Module:
        """We define our `ActorCriticModel` agent using a lightweight
        implementation with separate MLPs for actors and critic,
        MemorylessActorCritic.

        Since this is a model for continuous control, note that the
        superclass of our model is `ActorCriticModel[GaussianDistr]`
        instead of `ActorCriticModel[CategoricalDistr]`, since we'll use
        a Gaussian distribution to sample actions.
        """
        action_space = gym.spaces.Box(-1.0, 1.0, (2, ), "float32")
        return MemorylessActorCritic(
            input_uuid="gym_mujoco_data",
            action_space=action_space,  # specific action_space
            observation_space=SensorSuite(cls.SENSORS).observation_spaces,
            action_std=0.5,
        )
コード例 #20
0
 def create_model(cls, **kwargs) -> nn.Module:
     if cls.CNN_PREPROCESSOR_TYPE_AND_PRETRAINING is None:
         return RearrangeActorCriticSimpleConvRNN(
             action_space=gym.spaces.Discrete(len(cls.actions())),
             observation_space=SensorSuite(
                 cls.sensors()).observation_spaces,
             rgb_uuid=cls.EGOCENTRIC_RGB_UUID,
             unshuffled_rgb_uuid=cls.UNSHUFFLED_RGB_UUID,
         )
     else:
         return ResNetRearrangeActorCriticRNN(
             action_space=gym.spaces.Discrete(len(cls.actions())),
             observation_space=kwargs["sensor_preprocessor_graph"].
             observation_spaces,
             rgb_uuid=cls.EGOCENTRIC_RGB_RESNET_UUID,
             unshuffled_rgb_uuid=cls.UNSHUFFLED_RGB_RESNET_UUID,
         )
コード例 #21
0
 def make_sampler_fn(
     cls,
     stage: str,
     force_cache_reset: bool,
     allowed_scenes: Optional[Sequence[str]],
     seed: int,
     scene_to_allowed_rearrange_inds: Optional[Dict[str, Sequence[int]]] = None,
     x_display: Optional[str] = None,
     sensors: Optional[Sequence[Sensor]] = None,
     thor_controller_kwargs: Optional[Dict] = None,
     **kwargs,
 ) -> RearrangeTaskSampler:
     """Return an RearrangeTaskSampler."""
     sensors = cls.SENSORS if sensors is None else sensors
     if "mp_ctx" in kwargs:
         del kwargs["mp_ctx"]
     return RearrangeTaskSampler.from_fixed_dataset(
         run_walkthrough_phase=True,
         run_unshuffle_phase=False,
         stage=stage,
         allowed_scenes=allowed_scenes,
         scene_to_allowed_rearrange_inds=scene_to_allowed_rearrange_inds,
         rearrange_env_kwargs=dict(
             force_cache_reset=force_cache_reset,
             **cls.REARRANGE_ENV_KWARGS,
             controller_kwargs={
                 "x_display": x_display,
                 **cls.THOR_CONTROLLER_KWARGS,
                 "renderDepthImage": any(
                     isinstance(s, DepthSensor) for s in sensors
                 ),
                 **(
                     {} if thor_controller_kwargs is None else thor_controller_kwargs
                 ),
             },
         ),
         seed=seed,
         sensors=SensorSuite(sensors),
         max_steps=cls.MAX_STEPS,
         discrete_actions=cls.actions(),
         require_done_action=cls.REQUIRE_DONE_ACTION,
         force_axis_aligned_start=cls.FORCE_AXIS_ALIGNED_START,
         randomize_start_rotation=stage == "train"
         and cls.RANDOMIZE_START_ROTATION_DURING_TRAINING,
         **kwargs,
     )
コード例 #22
0
    def machine_params(self, mode="train", **kwargs):
        sampler_devices: Sequence[torch.device] = []
        devices: Sequence[torch.device]
        if mode == "train":
            workers_per_device = 1
            devices = ([torch.device("cpu")] if not torch.cuda.is_available()
                       else cast(Tuple, self.train_gpu_ids) *
                       workers_per_device)
            nprocesses = evenly_distribute_count_into_bins(
                self.num_train_processes, max(len(devices), 1))
            sampler_devices = self.sampler_devices
        elif mode == "valid":
            nprocesses = 1
            devices = ([torch.device("cpu")]
                       if not torch.cuda.is_available() else self.val_gpu_ids)
        elif mode == "test":
            nprocesses = 10 if torch.cuda.is_available() else 1
            devices = ([torch.device("cpu")]
                       if not torch.cuda.is_available() else self.test_gpu_ids)
        else:
            raise NotImplementedError(
                "mode must be 'train', 'valid', or 'test'.")

        sensors = [*self.SENSORS]
        if mode != "train":
            sensors = [
                s for s in sensors if not isinstance(s, ExpertActionSensor)
            ]

        sensor_preprocessor_graph = (SensorPreprocessorGraph(
            source_observation_spaces=SensorSuite(sensors).observation_spaces,
            preprocessors=self.preprocessors(),
        ) if mode == "train" or (
            (isinstance(nprocesses, int) and nprocesses > 0) or
            (isinstance(nprocesses, Sequence) and sum(nprocesses) > 0)) else
                                     None)

        return MachineParams(
            nprocesses=nprocesses,
            devices=devices,
            sampler_devices=sampler_devices
            if mode == "train" else devices,  # ignored with > 1 gpu_ids
            sensor_preprocessor_graph=sensor_preprocessor_graph,
        )
コード例 #23
0
    def machine_params(self, mode="train", **kwargs):
        sampler_devices: Sequence[int] = []
        if mode == "train":
            workers_per_device = 1
            gpu_ids = ([] if not torch.cuda.is_available() else
                       self.TRAIN_GPU_IDS * workers_per_device)
            nprocesses = (1 if not torch.cuda.is_available() else
                          self.split_num_processes(len(gpu_ids)))
            sampler_devices = self.TRAIN_GPU_IDS
        elif mode == "valid":
            nprocesses = 1
            gpu_ids = [] if not torch.cuda.is_available(
            ) else self.VALID_GPU_IDS
        elif mode == "test":
            nprocesses = 7
            gpu_ids = [] if not torch.cuda.is_available(
            ) else self.TEST_GPU_IDS
        else:
            raise NotImplementedError(
                "mode must be 'train', 'valid', or 'test'.")

        # Disable parallelization for validation process
        if mode == "valid":
            for prep in self.PREPROCESSORS:
                prep.kwargs["parallel"] = False

        sensor_preprocessor_graph = (SensorPreprocessorGraph(
            source_observation_spaces=SensorSuite(
                self.SENSORS).observation_spaces,
            preprocessors=self.PREPROCESSORS,
        ) if mode == "train" or (
            (isinstance(nprocesses, int) and nprocesses > 0) or
            (isinstance(nprocesses, Sequence) and sum(nprocesses) > 0)) else
                                     None)

        return MachineParams(
            nprocesses=nprocesses,
            devices=gpu_ids,
            sampler_devices=sampler_devices
            if mode == "train" else gpu_ids,  # ignored with > 1 gpu_ids
            sensor_preprocessor_graph=sensor_preprocessor_graph,
        )
コード例 #24
0
    def create_model(cls, **kwargs) -> nn.Module:
        map_kwargs = dict(
            frame_height=224,
            frame_width=224,
            vision_range_in_cm=cls.MAP_INFO["vision_range_in_cm"],
            resolution_in_cm=cls.MAP_INFO["resolution_in_cm"],
            map_size_in_cm=cls.MAP_INFO["map_size_in_cm"],
        )

        observation_space = (
            SensorSuite(cls.sensors()).observation_spaces
            if kwargs.get("sensor_preprocessor_graph") is None
            else kwargs["sensor_preprocessor_graph"].observation_spaces
        )

        semantic_map_channels = len(cls.ORDERED_OBJECT_TYPES)
        height_map_channels = 3
        map_kwargs["n_map_channels"] = height_map_channels + semantic_map_channels
        frozen_map = ActiveNeuralSLAM(**map_kwargs, use_resnet_layernorm=True)

        pretrained_map_ckpt_path = os.path.join(
            ABS_PATH_OF_REARRANGE_TOP_LEVEL_DIR,
            "pretrained_model_ckpts",
            "pretrained_active_neural_slam_via_walkthrough_75m.pt",
        )
        multiprocessing_safe_download_file_from_url(
            url="https://prior-model-weights.s3.us-east-2.amazonaws.com/embodied-ai/rearrangement/walkthrough/pretrained_active_neural_slam_via_walkthrough_75m.pt",
            save_path=pretrained_map_ckpt_path,
        )
        frozen_map.load_state_dict(
            torch.load(pretrained_map_ckpt_path, map_location="cpu",)
        )

        return OnePhaseRearrangeActorCriticFrozenMap(
            map=frozen_map,
            action_space=gym.spaces.Discrete(len(cls.actions())),
            observation_space=observation_space,
            rgb_uuid=cls.EGOCENTRIC_RGB_UUID,
            unshuffled_rgb_uuid=cls.UNSHUFFLED_RGB_UUID,
            semantic_map_channels=semantic_map_channels,
            height_map_channels=height_map_channels,
        )
    def create_model(cls, **kwargs) -> nn.Module:
        def get_sensor_uuid(stype: Type[Sensor]) -> Optional[str]:
            s = next(
                (s for s in cls.SENSORS if isinstance(s, stype)),
                None,
            )
            return None if s is None else s.uuid

        walkthrougher_should_ignore_action_mask = [
            any(k in a for k in ["drop", "open", "pickup"])
            for a in cls.actions()
        ]

        map_kwargs = dict(
            frame_height=224,
            frame_width=224,
            vision_range_in_cm=cls.MAP_INFO["vision_range_in_cm"],
            resolution_in_cm=cls.MAP_INFO["resolution_in_cm"],
            map_size_in_cm=cls.MAP_INFO["map_size_in_cm"],
        )

        observation_space = (
            SensorSuite(cls.SENSORS).observation_spaces
            if kwargs.get("sensor_preprocessor_graph") is None else
            kwargs["sensor_preprocessor_graph"].observation_spaces)

        semantic_map_channels = len(cls.ORDERED_OBJECT_TYPES)
        height_map_channels = 3
        map_kwargs[
            "n_map_channels"] = height_map_channels + semantic_map_channels
        frozen_map = ActiveNeuralSLAM(**map_kwargs, use_resnet_layernorm=True)

        pretrained_map_ckpt_path = os.path.join(
            ABS_PATH_OF_REARRANGE_TOP_LEVEL_DIR,
            "pretrained_model_ckpts",
            "pretrained_active_neural_slam_via_walkthrough_75m.pt",
        )
        multiprocessing_safe_download_file_from_url(
            url=
            "https://prior-model-weights.s3.us-east-2.amazonaws.com/embodied-ai/rearrangement/walkthrough/pretrained_active_neural_slam_via_walkthrough_75m.pt",
            save_path=pretrained_map_ckpt_path,
        )
        frozen_map.load_state_dict(
            torch.load(
                pretrained_map_ckpt_path,
                map_location="cpu",
            ))

        return TwoPhaseRearrangeActorCriticFrozenMap(
            map=frozen_map,
            semantic_map_channels=semantic_map_channels,
            height_map_channels=height_map_channels,
            action_space=gym.spaces.Discrete(len(cls.actions())),
            observation_space=observation_space,
            rgb_uuid=cls.EGOCENTRIC_RGB_UUID,
            in_walkthrough_phase_uuid=get_sensor_uuid(
                InWalkthroughPhaseSensor),
            is_walkthrough_phase_embedding_dim=cls.
            IS_WALKTHROUGH_PHASE_EMBEDING_DIM,
            rnn_type=cls.RNN_TYPE,
            walkthrougher_should_ignore_action_mask=
            walkthrougher_should_ignore_action_mask,
            done_action_index=cls.actions().index("done"),
        )
コード例 #26
0
    def __init__(
        self,
        env_class: Callable[..., Union[MiniGridEnv]],
        sensors: Union[SensorSuite, List[Sensor]],
        env_info: Optional[Dict[str, Any]] = None,
        max_tasks: Optional[int] = None,
        num_unique_seeds: Optional[int] = None,
        task_seeds_list: Optional[List[int]] = None,
        deterministic_sampling: bool = False,
        cache_graphs: Optional[bool] = False,
        task_class: Callable[..., MiniGridTask] = MiniGridTask,
        repeat_failed_task_for_min_steps: int = 0,
        extra_task_kwargs: Optional[Dict] = None,
        **kwargs,
    ):
        super(MiniGridTaskSampler, self).__init__()
        self.sensors = (SensorSuite(sensors)
                        if not isinstance(sensors, SensorSuite) else sensors)
        self.max_tasks = max_tasks
        self.num_unique_seeds = num_unique_seeds
        self.cache_graphs = cache_graphs
        self.deterministic_sampling = deterministic_sampling
        self.repeat_failed_task_for_min_steps = repeat_failed_task_for_min_steps
        self.extra_task_kwargs = (extra_task_kwargs
                                  if extra_task_kwargs is not None else {})

        self._last_env_seed: Optional[int] = None
        self._last_task: Optional[MiniGridTask] = None
        self._number_of_steps_taken_with_task_seed = 0

        assert (not deterministic_sampling
                ) or repeat_failed_task_for_min_steps <= 0, (
                    "If `deterministic_sampling` is True then we require"
                    " `repeat_failed_task_for_min_steps <= 0`")
        assert (not self.cache_graphs) or self.num_unique_seeds is not None, (
            "When caching graphs you must specify"
            " a number of unique tasks to sample from.")
        assert (self.num_unique_seeds is
                None) or (0 < self.num_unique_seeds
                          ), "`num_unique_seeds` must be a positive integer."

        self.num_unique_seeds = num_unique_seeds
        self.task_seeds_list = task_seeds_list
        if self.task_seeds_list is not None:
            if self.num_unique_seeds is not None:
                assert self.num_unique_seeds == len(
                    self.task_seeds_list
                ), "`num_unique_seeds` must equal the length of `task_seeds_list` if both specified."
            self.num_unique_seeds = len(self.task_seeds_list)
        elif self.num_unique_seeds is not None:
            self.task_seeds_list = list(range(self.num_unique_seeds))
        if num_unique_seeds is not None and repeat_failed_task_for_min_steps > 0:
            raise NotImplementedError(
                "`repeat_failed_task_for_min_steps` must be <=0 if number"
                " of unique seeds is not None.")

        assert (
            not self.cache_graphs
        ) or self.num_unique_seeds <= 1000, "Too many tasks (graphs) to cache"
        assert (not deterministic_sampling) or (
            self.num_unique_seeds is not None
        ), "Cannot use deterministic sampling when `num_unique_seeds` is `None`."

        if (not deterministic_sampling) and self.max_tasks:
            get_logger().warning(
                "`deterministic_sampling` is `False` but you have specified `max_tasks < inf`,"
                " this might be a mistake when running testing.")

        self.env = env_class(**env_info)
        self.task_class = task_class

        self.np_seeded_random_gen, _ = seeding.np_random(
            random.randint(0, 2**31 - 1))

        self.num_tasks_generated = 0
コード例 #27
0
    def __init__(
        self,
        gym_env_type: str = "LunarLanderContinuous-v2",
        sensors: Optional[Union[SensorSuite, List[Sensor]]] = None,
        max_tasks: Optional[int] = None,
        num_unique_seeds: Optional[int] = None,
        task_seeds_list: Optional[List[int]] = None,
        deterministic_sampling: bool = False,
        task_selector: Callable[[str], type] = task_selector,
        repeat_failed_task_for_min_steps: int = 0,
        extra_task_kwargs: Optional[Dict] = None,
        seed: Optional[int] = None,
        **kwargs,
    ):
        super().__init__()

        self.gym_env_type = gym_env_type

        self.sensors: SensorSuite
        if sensors is None:
            self.sensors = SensorSuite([sensor_selector(self.gym_env_type)])
        else:
            self.sensors = (SensorSuite(sensors) if
                            not isinstance(sensors, SensorSuite) else sensors)

        self.max_tasks = max_tasks
        self.num_unique_seeds = num_unique_seeds
        self.deterministic_sampling = deterministic_sampling
        self.repeat_failed_task_for_min_steps = repeat_failed_task_for_min_steps
        self.extra_task_kwargs = (extra_task_kwargs
                                  if extra_task_kwargs is not None else {})

        self._last_env_seed: Optional[int] = None
        self._last_task: Optional[GymTask] = None
        self._number_of_steps_taken_with_task_seed = 0

        assert (not deterministic_sampling
                ) or repeat_failed_task_for_min_steps <= 0, (
                    "If `deterministic_sampling` is True then we require"
                    " `repeat_failed_task_for_min_steps <= 0`")
        assert (self.num_unique_seeds is
                None) or (0 < self.num_unique_seeds
                          ), "`num_unique_seeds` must be a positive integer."

        self.num_unique_seeds = num_unique_seeds
        self.task_seeds_list = task_seeds_list
        if self.task_seeds_list is not None:
            if self.num_unique_seeds is not None:
                assert self.num_unique_seeds == len(
                    self.task_seeds_list
                ), "`num_unique_seeds` must equal the length of `task_seeds_list` if both specified."
            self.num_unique_seeds = len(self.task_seeds_list)
        elif self.num_unique_seeds is not None:
            self.task_seeds_list = list(range(self.num_unique_seeds))
        if num_unique_seeds is not None and repeat_failed_task_for_min_steps > 0:
            raise NotImplementedError(
                "`repeat_failed_task_for_min_steps` must be <=0 if number"
                " of unique seeds is not None.")

        assert (not deterministic_sampling) or (
            self.num_unique_seeds is not None
        ), "Cannot use deterministic sampling when `num_unique_seeds` is `None`."

        if (not deterministic_sampling) and self.max_tasks:
            get_logger().warning(
                "`deterministic_sampling` is `False` but you have specified `max_tasks < inf`,"
                " this might be a mistake when running testing.")

        if seed is not None:
            self.set_seed(seed)
        else:
            self.np_seeded_random_gen, _ = seeding.np_random(
                random.randint(0, 2**31 - 1))

        self.num_tasks_generated = 0
        self.task_type = task_selector(self.gym_env_type)
        self.env: GymEnvironment = GymEnvironment(self.gym_env_type)