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,
)
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,
)
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,
)
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
],
)
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"),
)
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,
)
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,
)
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,
)
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,
)
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)
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,
)
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,
)
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
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
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
],
)
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,
)
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,
)
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,
)
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,
)
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,
)
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,
)
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,
)
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,
)
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"),
)
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
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)