def create_model(cls, **kwargs) -> nn.Module:
return NavToPartnerActorCriticSimpleConvRNN(
action_space=gym.spaces.Discrete(
len(NavToPartnerTask.class_action_names())),
observation_space=SensorSuite(cls.SENSORS).observation_spaces,
hidden_size=512,
)
def __init__(
self,
source_ids: List[str],
all_preprocessors: List[Union[Preprocessor, Builder[Preprocessor]]],
all_sensors: List[Sensor],
) -> None:
"""Initializer.
# Parameters
source_ids : The sensors and preprocessors that will be included in the set.
all_preprocessors : The entire list of preprocessors to be executed.
all_sensors : The entire list of sensors.
"""
self.graph = PreprocessorGraph(all_preprocessors)
self.source_ids = source_ids
assert len(set(self.source_ids)) == len(
self.source_ids), "No duplicated uuids allowed in source_ids"
sensor_spaces = SensorSuite(all_sensors).observation_spaces
preprocessor_spaces = self.graph.observation_spaces
spaces: OrderedDict[str, gym.Space] = OrderedDict()
for uuid in self.source_ids:
assert (
uuid in sensor_spaces.spaces
or uuid in preprocessor_spaces.spaces
), "uuid {} missing from sensor suite and preprocessor graph".format(
uuid)
if uuid in sensor_spaces.spaces:
spaces[uuid] = sensor_spaces[uuid]
else:
spaces[uuid] = preprocessor_spaces[uuid]
self.observation_spaces = SpaceDict(spaces=spaces)
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 ObjectNavBaselineActorCritic(
action_space=gym.spaces.Discrete(len(ObjectNavTask.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:
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, 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,
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 __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 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 = 10 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 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 __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
