diff --git a/habitat_extensions/config/default.py b/habitat_extensions/config/default.py index 6e47e5c6..7deeabc8 100644 --- a/habitat_extensions/config/default.py +++ b/habitat_extensions/config/default.py @@ -44,6 +44,13 @@ # ---------------------------------------------------------------------------- _C.TASK.VLN_ORACLE_PROGRESS_SENSOR = CN() _C.TASK.VLN_ORACLE_PROGRESS_SENSOR.TYPE = "VLNOracleProgressSensor" + +# ----------------------------------------------------------------------------- +# VLN ORACLE DISTANCE LEFT SENSOR +# ---------------------------------------------------------------------------- +_C.TASK.VLN_ORACLE_DISTANCE_LEFT_SENSOR = CN() +_C.TASK.VLN_ORACLE_DISTANCE_LEFT_SENSOR.TYPE = "VLNOracleDistanceLeftSensor" + # ---------------------------------------------------------------------------- # PANO ANGLE FEATURE SENSOR # ---------------------------------------------------------------------------- diff --git a/habitat_extensions/config/vlnce_task.yaml b/habitat_extensions/config/vlnce_task.yaml index dddc88fb..e23d772e 100644 --- a/habitat_extensions/config/vlnce_task.yaml +++ b/habitat_extensions/config/vlnce_task.yaml @@ -23,7 +23,8 @@ TASK: SENSORS: [ INSTRUCTION_SENSOR, SHORTEST_PATH_SENSOR, - VLN_ORACLE_PROGRESS_SENSOR + VLN_ORACLE_PROGRESS_SENSOR, + VLN_ORACLE_DISTANCE_LEFT_SENSOR ] INSTRUCTION_SENSOR_UUID: instruction POSSIBLE_ACTIONS: [STOP, MOVE_FORWARD, TURN_LEFT, TURN_RIGHT] diff --git a/habitat_extensions/config/vlnce_task_aug.yaml b/habitat_extensions/config/vlnce_task_aug.yaml index 72b8e72c..558b7523 100644 --- a/habitat_extensions/config/vlnce_task_aug.yaml +++ b/habitat_extensions/config/vlnce_task_aug.yaml @@ -26,7 +26,8 @@ TASK: SENSORS: [ INSTRUCTION_SENSOR, SHORTEST_PATH_SENSOR, - VLN_ORACLE_PROGRESS_SENSOR + VLN_ORACLE_PROGRESS_SENSOR, + VLN_ORACLE_DISTANCE_LEFT_SENSOR ] INSTRUCTION_SENSOR_UUID: instruction POSSIBLE_ACTIONS: [STOP, MOVE_FORWARD, TURN_LEFT, TURN_RIGHT] diff --git a/habitat_extensions/sensors.py b/habitat_extensions/sensors.py index e2674f10..8bd8f030 100644 --- a/habitat_extensions/sensors.py +++ b/habitat_extensions/sensors.py @@ -86,6 +86,41 @@ def get_observation(self, *args: Any, episode, **kwargs: Any) -> float: [(distance_from_start - distance_to_target) / distance_from_start] ) +@registry.register_sensor +class VLNOracleDistanceLeftSensor(Sensor): + """Distance left towards goal""" + + cls_uuid: str = "distance_left" + + def __init__( + self, sim: Simulator, config: Config, *args: Any, **kwargs: Any + ) -> None: + self._sim = sim + super().__init__(config=config) + + def _get_uuid(self, *args: Any, **kwargs: Any) -> str: + return self.cls_uuid + + def _get_sensor_type(self, *args: Any, **kwargs: Any) -> SensorTypes: + return SensorTypes.MEASUREMENT + + def _get_observation_space(self, *args: Any, **kwargs: Any) -> Space: + return spaces.Box(low=0.0, high=1.0, shape=(1,), dtype=np.float) + + def get_observation(self, *args: Any, episode, **kwargs: Any) -> float: + distance_to_target = self._sim.geodesic_distance( + self._sim.get_agent_state().position.tolist(), + episode.goals[0].position, + ) + + # just in case the agent ends up somewhere it shouldn't + if not np.isfinite(distance_to_target): + distance_to_target = 0.0 + + + return np.array( + [distance_to_target] + ) @registry.register_sensor class AngleFeaturesSensor(Sensor): diff --git a/run.py b/run.py index 31bc4cbf..e87191aa 100644 --- a/run.py +++ b/run.py @@ -8,7 +8,7 @@ import torch from habitat import logger from habitat_baselines.common.baseline_registry import baseline_registry - +import gc import habitat_extensions # noqa: F401 import vlnce_baselines # noqa: F401 from vlnce_baselines.config.default import get_config @@ -30,7 +30,8 @@ def main(): "--exp-config", type=str, required=True, - help="path to config yaml containing info about experiment", + help="path to config yaml containing info about experiment. " + "If this is a directory, run all yaml file contained in the dir.", ) parser.add_argument( "opts", @@ -40,7 +41,19 @@ def main(): ) args = parser.parse_args() - run_exp(**vars(args)) + if os.path.isdir(args.exp_config): + conf_parameter = args.exp_config + if os.path.isdir(conf_parameter): + print("Running several config files from:", conf_parameter) + for file in sorted(os.listdir(conf_parameter)): + if file.endswith(".yaml") or file.endswith(".yml"): + file_path = os.path.join(conf_parameter, file) + print("exp_config", file_path) + run_exp(exp_config=file_path, run_type=args.run_type, opts=args.opts) + else: + print("Not a valid config file:", file) + else: + run_exp(**vars(args)) def run_exp(exp_config: str, run_type: str, opts=None) -> None: @@ -52,11 +65,16 @@ def run_exp(exp_config: str, run_type: str, opts=None) -> None: opts: list of strings of additional config options. """ config = get_config(exp_config, opts) - logger.info(f"config: {config}") logdir = "/".join(config.LOG_FILE.split("/")[:-1]) + if not logdir: + logdir = "logs" + os.makedirs(logdir, exist_ok=True) + config_file_root__name = logdir+"/"+exp_config.split("/")[-1].split(".")[0] if logdir: os.makedirs(logdir, exist_ok=True) - logger.add_filehandler(config.LOG_FILE) + log_file = config_file_root__name + "_" + config.LOG_FILE + logger.add_filehandler(log_file) + logger.info(f"config: {config}") random.seed(config.TASK_CONFIG.SEED) np.random.seed(config.TASK_CONFIG.SEED) @@ -87,6 +105,9 @@ def run_exp(exp_config: str, run_type: str, opts=None) -> None: elif run_type == "inference": trainer.inference() + # avoids to write to all previous files if running in a loop + logger.removeHandler(logger.handlers[-1]) + gc.collect() if __name__ == "__main__": main() diff --git a/vlnce_baselines/__init__.py b/vlnce_baselines/__init__.py index d7b01c87..9cb30a1b 100644 --- a/vlnce_baselines/__init__.py +++ b/vlnce_baselines/__init__.py @@ -2,6 +2,9 @@ dagger_trainer, ddppo_waypoint_trainer, recollect_trainer, + decision_transformer_trainer ) from vlnce_baselines.common import environments -from vlnce_baselines.models import cma_policy, seq2seq_policy, waypoint_policy +from vlnce_baselines.models import (cma_policy, seq2seq_policy, + waypoint_policy, + decision_transformer_policy) diff --git a/vlnce_baselines/common/base_il_trainer.py b/vlnce_baselines/common/base_il_trainer.py index 7ab46046..6e9a1563 100644 --- a/vlnce_baselines/common/base_il_trainer.py +++ b/vlnce_baselines/common/base_il_trainer.py @@ -65,8 +65,8 @@ def _initialize_policy( action_space=action_space, ) self.policy.to(self.device) - - self.optimizer = torch.optim.Adam( + # torch.optim.RAdam or torch.optim.Adam for example + self.optimizer = eval(config.IL.optimizer)( self.policy.parameters(), lr=self.config.IL.lr ) if load_from_ckpt: @@ -191,16 +191,21 @@ def _pause_envs( ): # pausing envs with no new episode if len(envs_to_pause) > 0: + # That can avoid nasty bugs when creating new Trainers... + envs_to_pause = sorted(envs_to_pause) state_index = list(range(envs.num_envs)) for idx in reversed(envs_to_pause): state_index.pop(idx) envs.pause_at(idx) - # indexing along the batch dimensions - recurrent_hidden_states = recurrent_hidden_states[state_index] - not_done_masks = not_done_masks[state_index] - prev_actions = prev_actions[state_index] - + # indexing along the batch dimensions => because we removed the environement to pause in + # the previous step from the state_index list, we just keep everything related to the active environments + if recurrent_hidden_states is not None: + recurrent_hidden_states = recurrent_hidden_states[state_index] + if not_done_masks is not None: + not_done_masks = not_done_masks[state_index] + if prev_actions is not None: + prev_actions = prev_actions[state_index] for k, v in batch.items(): batch[k] = v[state_index] diff --git a/vlnce_baselines/common/env_utils.py b/vlnce_baselines/common/env_utils.py index 381c50e9..69d3c3c2 100644 --- a/vlnce_baselines/common/env_utils.py +++ b/vlnce_baselines/common/env_utils.py @@ -93,6 +93,7 @@ def construct_envs( env_fn_args=tuple(zip(configs, env_classes)), auto_reset_done=auto_reset_done, workers_ignore_signals=workers_ignore_signals, + multiprocessing_start_method=config.MULTIPROCESSING ) return envs diff --git a/vlnce_baselines/common/environments.py b/vlnce_baselines/common/environments.py index d707ac32..42b38cc0 100644 --- a/vlnce_baselines/common/environments.py +++ b/vlnce_baselines/common/environments.py @@ -12,6 +12,26 @@ from habitat_extensions.utils import generate_video, navigator_video_frame +@baseline_registry.register_env(name="VLNCEDecisionTransformerEnv") +class VLNCEDecisionTransformerEnv(habitat.RLEnv): + def __init__(self, config: Config, dataset: Optional[Dataset] = None): + super().__init__(config.TASK_CONFIG, dataset) + + def get_reward_range(self) -> Tuple[float, float]: + # We don't use the Habitat Framework to create rewards, they are + # created with the trajectories. + return (0.0, 0.0) + + def get_reward(self, observations: Observations) -> float: + return 0.0 + + def get_done(self, observations: Observations) -> bool: + return self._env.episode_over + + def get_info(self, observations: Observations) -> Dict[Any, Any]: + return self.habitat_env.get_metrics() + + @baseline_registry.register_env(name="VLNCEDaggerEnv") class VLNCEDaggerEnv(habitat.RLEnv): def __init__(self, config: Config, dataset: Optional[Dataset] = None): diff --git a/vlnce_baselines/config/default.py b/vlnce_baselines/config/default.py index 6fdb7b1e..34635238 100644 --- a/vlnce_baselines/config/default.py +++ b/vlnce_baselines/config/default.py @@ -24,7 +24,10 @@ _C.VIDEO_DIR = "data/videos/debug" _C.TENSORBOARD_DIR = "data/tensorboard_dirs/debug" _C.RESULTS_DIR = "data/checkpoints/pretrained/evals" - +# Enables debugging for Pycharm. Default value is "forkserver". +# https://youtrack.jetbrains.com/issue/PY-52273/Debugger-multiprocessing-hangs-pycharm-20213 +_C.MULTIPROCESSING = "forkserver" # Set to 'spawn' when debugging with Pycharm, +_C.use_pbar = True # ---------------------------------------------------------------------------- # EVAL CONFIG # ---------------------------------------------------------------------------- @@ -37,7 +40,9 @@ _C.EVAL.EVAL_NONLEARNING = False _C.EVAL.NONLEARNING = CN() _C.EVAL.NONLEARNING.AGENT = "RandomAgent" - +_C.EVAL.VAL_SEEN_SMALL = "val_seen_80_ep" # only used when ran in train_complete mode +_C.EVAL.VAL_SEEN = "val_seen" # only used when ran in train_complete mode +_C.EVAL.VAL_UNSEEN = "val_unseen" # only used when ran in train_complete mode # ---------------------------------------------------------------------------- # INFERENCE CONFIG # ---------------------------------------------------------------------------- @@ -57,10 +62,13 @@ # IMITATION LEARNING CONFIG # ---------------------------------------------------------------------------- _C.IL = CN() +_C.IL.optimizer = "torch.optim.Adam" +_C.IL.dataload_workers = 1 _C.IL.lr = 2.5e-4 _C.IL.batch_size = 5 # number of network update rounds per iteration _C.IL.epochs = 4 +_C.IL.preload_dataloader_size = 100 # if true, uses class-based inflection weighting _C.IL.use_iw = True # inflection coefficient for RxR training set GT trajectories (guide): 1.9 @@ -69,9 +77,12 @@ # load an already trained model for fine tuning _C.IL.load_from_ckpt = False _C.IL.ckpt_to_load = "data/checkpoints/ckpt.0.pth" +_C.IL.continue_ckpt_naming = True # if True, loads the optimizer state, epoch, and step_id from the ckpt dict. _C.IL.is_requeue = False - +_C.IL.checkpoint_frequency = 1 # regulates the frequency (epochs % checkpoint_frequency == 0) to save the model. +_C.IL.mean_loss_to_save_checkpoint = 0.40 +_C.IL.mean_loss_to_stop_training = 0.06 # ---------------------------------------------------------------------------- # IL: RECOLLECT TRAINER CONFIG # ---------------------------------------------------------------------------- @@ -116,6 +127,28 @@ "data/trajectories_dirs/debug/trajectories.lmdb" ) _C.IL.DAGGER.drop_existing_lmdb_features = True + +# ---------------------------------------------------------------------------- +# IL: DAGGER / DECISION TRANSFORMER CONFIG +# ---------------------------------------------------------------------------- + +_C.IL.DECISION_TRANSFORMER = CN() +_C.IL.DECISION_TRANSFORMER.episode_horizon = 183 +_C.IL.DECISION_TRANSFORMER.use_perfect_episode_only_for_dagger = True +_C.IL.DECISION_TRANSFORMER.use_oracle_actions = False +_C.IL.DECISION_TRANSFORMER.reward_type = "POINT_GOAL_NAV_REWARD" # POINT_GOAL_NAV_REWARD or SPARSE_REWARD +_C.IL.DECISION_TRANSFORMER.sensor_uuid = "distance_left" # USed to calculate the Return To Go +_C.IL.DECISION_TRANSFORMER.recompute_reward = True +_C.IL.DECISION_TRANSFORMER.POINT_GOAL_NAV_REWARD = CN() +_C.IL.DECISION_TRANSFORMER.POINT_GOAL_NAV_REWARD.step_penalty = -0.01 +_C.IL.DECISION_TRANSFORMER.POINT_GOAL_NAV_REWARD.success = 1.0 +_C.IL.DECISION_TRANSFORMER.SPARSE_REWARD = CN() +_C.IL.DECISION_TRANSFORMER.POINT_GOAL_NAV_REWARD.step_penalty = -0.01 +_C.IL.DECISION_TRANSFORMER.POINT_GOAL_NAV_REWARD.success = 1.0 +_C.IL.DECISION_TRANSFORMER.NDTW_REWARD = CN() +_C.IL.DECISION_TRANSFORMER.NDTW_REWARD.step_penalty = -0.01 +_C.IL.DECISION_TRANSFORMER.NDTW_REWARD.success = 1.0 + # ---------------------------------------------------------------------------- # RL CONFIG # ---------------------------------------------------------------------------- @@ -284,6 +317,58 @@ _C.MODEL.WAYPOINT.discrete_offsets = 7 _C.MODEL.WAYPOINT.offset_temperature = 1.0 +# ---------------------------------------------------------------------------- +# DECISION TRANSFORMER CONFIG +# ---------------------------------------------------------------------------- +_C.MODEL.DECISION_TRANSFORMER = CN() +_C.MODEL.DECISION_TRANSFORMER.use_re_zero = False # https://arxiv.org/abs/2003.04887 +_C.MODEL.DECISION_TRANSFORMER.hidden_dim = 128 +# the max in the training split. +_C.MODEL.DECISION_TRANSFORMER.episode_horizon = _C.IL.DECISION_TRANSFORMER.episode_horizon +_C.MODEL.DECISION_TRANSFORMER.reward_type = "POINT_GOAL_NAV_REWARD" # POINT_GOAL_NAV_REWARD or SPARSE_REWARD +_C.MODEL.DECISION_TRANSFORMER.return_to_go_inference = 1.0 +_C.MODEL.DECISION_TRANSFORMER.spatial_output = False # If set to false, depth and rgb feature are averaged +_C.MODEL.DECISION_TRANSFORMER.model_type = None +_C.MODEL.DECISION_TRANSFORMER.n_layer = 2 +_C.MODEL.DECISION_TRANSFORMER.n_head = 1 +_C.MODEL.DECISION_TRANSFORMER.n_embd = _C.MODEL.DECISION_TRANSFORMER.hidden_dim +_C.MODEL.DECISION_TRANSFORMER.use_transformer_encoded_instruction = False +# these options must be filled in externally +_C.MODEL.DECISION_TRANSFORMER.vocab_size = 4 +_C.MODEL.DECISION_TRANSFORMER.step_size = 3 #We multiply by three because at each time step, we use [reward, action, state]. +_C.MODEL.DECISION_TRANSFORMER.block_size = _C.MODEL.DECISION_TRANSFORMER.episode_horizon *_C.MODEL.DECISION_TRANSFORMER.step_size +_C.MODEL.DECISION_TRANSFORMER.allowed_models = ["DecisionTransformerNet", + "DecisionTransformerEnhancedNet", + "FullDecisionTransformerNet", + "FullDecisionTransformerSingleVisionStateNet"] +_C.MODEL.DECISION_TRANSFORMER.allowed_rewards = ["point_nav_reward_to_go", "sparse_reward_to_go", + "point_nav_reward", "sparse_reward", "ndtw_reward", + "ndtw_reward_to_go"] +_C.MODEL.DECISION_TRANSFORMER.exclude_past_action_for_prediction = True +_C.MODEL.DECISION_TRANSFORMER.normalize_depth = False # Needs to be done during dataset creation +_C.MODEL.DECISION_TRANSFORMER.normalize_rgb = False +# dropout hyperparameters +_C.MODEL.DECISION_TRANSFORMER.embd_pdrop = 0.1 +_C.MODEL.DECISION_TRANSFORMER.resid_pdrop = 0.1 +_C.MODEL.DECISION_TRANSFORMER.attn_pdrop = 0.1 +_C.MODEL.DECISION_TRANSFORMER.activation_action_drop = 0.3 +_C.MODEL.DECISION_TRANSFORMER.activation_instruction_drop = 0.0 +_C.MODEL.DECISION_TRANSFORMER.activation_rgb_drop = 0.0 +_C.MODEL.DECISION_TRANSFORMER.activation_depth_drop = 0.0 +_C.MODEL.DECISION_TRANSFORMER.ENCODER = CN() +_C.MODEL.DECISION_TRANSFORMER.ENCODER.n_layer = 2 +_C.MODEL.DECISION_TRANSFORMER.ENCODER.n_head = 1 +_C.MODEL.DECISION_TRANSFORMER.ENCODER.use_sentence_encoding = True +# Only for FullDecisionTransformerNet +_C.MODEL.DECISION_TRANSFORMER.ENCODER.use_rgb_state_embeddings = True +_C.MODEL.DECISION_TRANSFORMER.ENCODER.use_depth_state_embeddings = True +_C.MODEL.DECISION_TRANSFORMER.ENCODER.use_output_rgb_instructions = True +_C.MODEL.DECISION_TRANSFORMER.ENCODER.use_output_depth_instructions = True +_C.MODEL.DECISION_TRANSFORMER.ENCODER.use_output_rgb = True +_C.MODEL.DECISION_TRANSFORMER.ENCODER.use_output_depth = True +# Only for FullDecisionTransformerSingleVisionStateNet +_C.MODEL.DECISION_TRANSFORMER.ENCODER.use_output_state_instructions = True +_C.MODEL.DECISION_TRANSFORMER.ENCODER.use_output_state = True def purge_keys(config: CN, keys: List[str]) -> None: for k in keys: diff --git a/vlnce_baselines/config/r2r_baselines/README.md b/vlnce_baselines/config/r2r_baselines/README.md index 4d28e168..c305f3ff 100644 --- a/vlnce_baselines/config/r2r_baselines/README.md +++ b/vlnce_baselines/config/r2r_baselines/README.md @@ -36,3 +36,21 @@ gdown https://drive.google.com/uc?id=1xIxh5eUkjGzSL_3AwBqDQlNXjkFrpcg4 # Seq2Seq_DA (135MB) gdown https://drive.google.com/uc?id=14y7dXkAEwB_q81cDCow8JNKD2aAAPxbW ``` + + +## Experimental + +A trainer based on the [Decision Transformer](https://arxiv.org/abs/2106.01345) has been added. + +The results are underwhelming (below the best Seq2Seq) but it constitutes a good starting point for +anybody wanting to test Transformer in VLN-CE. + +Pretrained models with corresponding training file under : + + +[Decision Transformer Agent](https://drive.google.com/file/d/1-E1l5g7DM36m3HYx8b4b4CNBC8d-OS83/view?usp=sharing) + +[Enhanced Decision Transformer Agent](https://drive.google.com/file/d/1b2hpkHpiZIc2CBsaLzZWCa7qurfKDDsu/view?usp=sharing) + +[Full Decision Transformer Agent](https://drive.google.com/file/d/1rS2_yo9_z35zzaHW4CtByorZ-jpDpht_/view?usp=sharing) + diff --git a/vlnce_baselines/decision_transformer_trainer.py b/vlnce_baselines/decision_transformer_trainer.py new file mode 100644 index 00000000..87e04471 --- /dev/null +++ b/vlnce_baselines/decision_transformer_trainer.py @@ -0,0 +1,1572 @@ +import gc +import random +import warnings + +import lmdb +import msgpack_numpy +import numpy as np +from vlnce_baselines.common.base_il_trainer import BaseVLNCETrainer +from vlnce_baselines.common.env_utils import construct_envs + +from torch import Tensor +import re +import json +import os +import time +import warnings +from collections import defaultdict +import torch +import torch.nn.functional as F +import tqdm +from habitat import Config, logger +from habitat.utils.visualizations.utils import append_text_to_image +from habitat_baselines.common.baseline_registry import baseline_registry +from habitat_baselines.common.environments import get_env_class +from habitat_baselines.common.obs_transformers import ( + apply_obs_transforms_batch, +) +from habitat_baselines.common.tensorboard_utils import TensorboardWriter +from habitat_baselines.rl.ddppo.algo.ddp_utils import is_slurm_batch_job +from habitat_baselines.utils.common import batch_obs + +from habitat_extensions.utils import generate_video, observations_to_image +from vlnce_baselines.common.env_utils import construct_envs_auto_reset_false +from vlnce_baselines.common.utils import extract_instruction_tokens + +with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=FutureWarning) + import tensorflow as tf # noqa: F401 +import jsonlines +from typing import Any, Dict, List, Optional, Tuple + + +class ObservationsDict(dict): + def pin_memory(self): + for k, v in self.items(): + self[k] = v.pin_memory() + + return self + + +# Trick to create extra start token directly in the collate_fn +# we don t need to recreate the whole dataset... +global EXTRA_START_TOKEN_ID +global STOP_ACTION_TOKEN_ID +EXTRA_START_TOKEN_ID = 4 +STOP_ACTION_TOKEN_ID = 0 + + +def _is_correct_previous_actions(batch): + """ + Somehow, I detected a bug. Some actions are not shifted correctly + prev_actions_batch[i+1] is not always equal to corrected_actions_batch[i] + :param batch: + :return: + """ + return sum([(batch[i][1][1:] == batch[i][2][:-1]).sum() == len(batch[i][1][1:]) for i in range(len(batch))]) == len( + batch) + + +def collate_fn_check_batch(batch): + """Each sample in batch: ( + obs, + prev_actions, + oracle_actions, + inflec_weight, + ) + """ + if not _is_correct_previous_actions(batch): + raise Exception( + "Dataset has not been created correctly! Prev actions and corrected actions not shifted accordingly!") + + +def _block_shuffle(lst, block_size): + blocks = [lst[i: i + block_size] for i in range(0, len(lst), block_size)] + random.shuffle(blocks) + + return [ele for block in blocks for ele in block] + + +class IWTrajectoryDataset(torch.utils.data.IterableDataset): + def __init__( + self, + lmdb_features_dir, + use_iw, + inflection_weight_coef=1.0, + lmdb_map_size=1e9, + batch_size=1, + preload_size=128 + ): + super().__init__() + assert preload_size > 0 + self.lmdb_features_dir = lmdb_features_dir + self.lmdb_map_size = lmdb_map_size + self.preload_size = batch_size * preload_size + self._preload = [] + self.batch_size = batch_size + + if use_iw: + self.inflec_weights = torch.tensor([1.0, inflection_weight_coef]) + else: + self.inflec_weights = torch.tensor([1.0, 1.0]) + + with lmdb.open( + self.lmdb_features_dir, + map_size=int(self.lmdb_map_size), + readonly=True, + lock=False, + ) as lmdb_env: + self.length = lmdb_env.stat()["entries"] + + def _load_next(self): + if len(self._preload) == 0: + if len(self.load_ordering) == 0: + raise StopIteration + + new_preload = [] + lengths = [] + with lmdb.open( + self.lmdb_features_dir, + map_size=int(self.lmdb_map_size), + readonly=True, + lock=False, + readahead=False, + meminit=True + ) as lmdb_env, lmdb_env.begin(buffers=True) as txn: + for _ in range(self.preload_size): + if len(self.load_ordering) == 0: + break + entry = txn.get(str(self.load_ordering.pop()).encode()) + unpacked = msgpack_numpy.unpackb(entry,raw=False,) + new_preload.append(unpacked) + + lengths.append(len(new_preload[-1][0])) + + sort_priority = list(range(len(lengths))) + random.shuffle(sort_priority) + + sorted_ordering = list(range(len(lengths))) + sorted_ordering.sort(key=lambda k: (lengths[k], sort_priority[k])) + + for idx in _block_shuffle(sorted_ordering, self.batch_size): + self._preload.append(new_preload[idx]) + + return self._preload.pop() + + def __next__(self): + obs, prev_actions, oracle_actions = self._load_next() + + for k, v in obs.items(): + obs[k] = torch.from_numpy(np.copy(v)) + + prev_actions = torch.from_numpy(np.copy(prev_actions)) + oracle_actions = torch.from_numpy(np.copy(oracle_actions)) + + inflections = torch.cat( + [ + torch.tensor([1], dtype=torch.long), + (oracle_actions[1:] != oracle_actions[:-1]).long(), + ] + ) + + return ( + obs, + prev_actions, + oracle_actions, + self.inflec_weights[inflections], + ) + + def __iter__(self): + worker_info = torch.utils.data.get_worker_info() + if worker_info is None: + start = 0 + end = self.length + else: + per_worker = int(np.ceil(self.length / worker_info.num_workers)) + + start = per_worker * worker_info.id + end = min(start + per_worker, self.length) + + # Reverse so we can use .pop() + self.load_ordering = list( + reversed( + _block_shuffle(list(range(start, end)), self.preload_size) + ) + ) + + return self + + +@baseline_registry.register_trainer(name="decision_transformer") +class DecisionTransformerTrainer(BaseVLNCETrainer): + def __init__(self, config=None): + self.lmdb_features_dir = config.IL.DAGGER.lmdb_features_dir.format( + split=config.TASK_CONFIG.DATASET.SPLIT + ) + # + self.rewards = {"point_nav_reward": { + "step_penalty": config.IL.DECISION_TRANSFORMER.POINT_GOAL_NAV_REWARD.step_penalty, + "success": config.IL.DECISION_TRANSFORMER.POINT_GOAL_NAV_REWARD.success}, + "sparse_reward": { + "step_penalty": config.IL.DECISION_TRANSFORMER.SPARSE_REWARD.step_penalty, + "success": config.IL.DECISION_TRANSFORMER.SPARSE_REWARD.success}, + "ndtw_reward": { + "step_penalty": config.IL.DECISION_TRANSFORMER.NDTW_REWARD.step_penalty, + "success": config.IL.DECISION_TRANSFORMER.NDTW_REWARD.success}, + } + device = "cuda:" + str(config.TORCH_GPU_ID) + + self.rgb_depth_stats = { + "mean_rgb": torch.as_tensor(np.asarray([0.533, 0.498, 0.453]), + dtype=torch.float32, device=device), + "std_rgb": torch.as_tensor(np.asarray([0.183, 0.185, 0.2020]), + dtype=torch.float32, device=device), + "mean_depth": torch.as_tensor(np.asarray([0.222]), dtype=torch.float32, device=device), + "std_depth": torch.as_tensor(np.asarray([0.18]), dtype=torch.float32, device=device)} + + + + super().__init__(config) + + def _create_feature_hooks(self): + self.rgb_features = None + self.rgb_hook = None + + def hook_builder(tgt_tensor): + def hook(m, i, o): + tgt_tensor.set_(o.cpu()) + + return hook + + if not self.config.MODEL.RGB_ENCODER.trainable: + self.rgb_features = torch.zeros((1,), device="cpu") + self.rgb_hook = self.policy.net.rgb_encoder.cnn.register_forward_hook( + hook_builder(self.rgb_features) + ) + + self.depth_features = None + self.depth_hook = None + if not self.config.MODEL.DEPTH_ENCODER.trainable: + self.depth_features = torch.zeros((1,), device="cpu") + self.depth_hook = self.policy.net.depth_encoder.visual_encoder.register_forward_hook( + hook_builder(self.depth_features) + ) + + def _release_hook(self): + + if self.rgb_hook is not None: + self.rgb_hook.remove() + if self.depth_hook is not None: + self.depth_hook.remove() + self.rgb_features = torch.zeros((1,), device="cpu") + self.depth_features = torch.zeros((1,), device="cpu") + + def _calculate_return_to_go(self, traj_obs: dict, reward_type: str, observation_type: str, scaling_factor=1.0, + destination_key=None): + """ + Calculate the return to go. For a given step, sum of all rewards to come + :param traj_obs: + :param reward_type: + :param observation_type: + :param destination_key: if not given, will try to derive a destination name from observation_type (should start with raw_) + :param scaling_factor: scale the rewards down, proportinally to the sequence length + :return: + """ + assert (reward_type in self.rewards.keys()) + assert (observation_type in traj_obs.keys()) + rewards = traj_obs[observation_type] + # work around when transforming the values read in the database... + # that avoid to have a second fonction for the collate_fn when recalculating on the fly + isTensor = type(rewards) is torch.Tensor + if isTensor: + rewards = rewards.numpy() + rewards = rewards + self.rewards[reward_type]["step_penalty"] + rewards[-1] = rewards[-1] + self.rewards[reward_type]["success"] + # Just save the simple rewards for each time steps, not accumulated if needed + rewards = np.float32(rewards.squeeze()) + # In some cases, when using dagger, The agent stops on first action. Hence, + # we need to transform the scalar in array... + if type(rewards) is np.float32: + rewards = np.array([rewards]) + if isTensor: + simple_reward = torch.from_numpy(rewards) + else: + simple_reward = rewards + traj_obs[reward_type] = simple_reward + rewards = np.flip(np.flip(rewards).cumsum()) + if destination_key is None: + assert observation_type.startswith("raw_") + destination_key = observation_type.split("raw_")[1] + "_to_go" + reward_to_go = np.float32(rewards / scaling_factor) + if isTensor: + reward_to_go = torch.from_numpy(reward_to_go) + traj_obs[destination_key] = reward_to_go + + def _calculate_rewards(self, traj_obs, scaling_factor): + self._calculate_return_to_go(traj_obs, "point_nav_reward", "raw_point_nav_reward", scaling_factor) + self._calculate_return_to_go(traj_obs, "sparse_reward", "raw_sparse_reward", scaling_factor) + self._calculate_return_to_go(traj_obs, "ndtw_reward", "raw_ndtw_reward", scaling_factor) + + def _make_dirs(self) -> None: + self._make_ckpt_dir() + os.makedirs(self.lmdb_features_dir, exist_ok=True) + if self.config.EVAL.SAVE_RESULTS: + self._make_results_dir() + + def _prepare_observation(self, observations): + ''' + From the observation created by the environment, creates dictionaries of features, + with shape : number of environment * all the remaning features. + :param observations: + :return: + ''' + observations = extract_instruction_tokens( + observations, self.config.TASK_CONFIG.TASK.INSTRUCTION_SENSOR_UUID + ) + batch = batch_obs(observations, self.device) + batch = apply_obs_transforms_batch(batch, self.obs_transforms) + + return observations, batch + + def _modify_batch_for_transformer(self, episodes: list, batch: ObservationsDict, rgb_features: Tensor, + depth_features: Tensor, envs, prev_actions, rgb_key, depth_key): + """ + This function help to prepare the input needed by the transformer model. + Habitat Sim provides one image / observation per time step, we need the whole serie for the transformer model. + Hence, the batch parameter is modified to have a whole sequence of rgb, depth and instructions + Moreover, the previous actions are returned by this function (only important at the first timestep) + :param episodes: + :param batch: + :param rgb_features: + :param depth_features: + :param envs: + :param prev_actions: + :param rgb_key: + :param depth_key: + :return: + """ + current_rgb = rgb_features.unsqueeze(dim=1) + current_depth = depth_features.unsqueeze(dim=1) + + if not self._are_episodes_empty(episodes): + # preparing the past images as a sequence + rgb_seq = self._create_sequence(episodes, rgb_key) + depth_seq = self._create_sequence(episodes, depth_key) + # adding the current image to the end of the sequence + rgb_seq = torch.cat((rgb_seq, current_rgb), dim=1) + depth_seq = torch.cat((depth_seq, current_depth), dim=1) + else: + # we unsqueeze at dim = 1 to create a shape of of batch, sequence (of size 1 at the beginning), and all other dim + rgb_seq = current_rgb + depth_seq = current_depth + prev_actions = torch.zeros( + envs.num_envs, + 1, + device=self.device, + dtype=torch.long, + ) + if self.config.MODEL.DECISION_TRANSFORMER.use_extra_start_token: + prev_actions = prev_actions + EXTRA_START_TOKEN_ID + # store the last images + for i in range(envs.num_envs): + episodes[i].append({rgb_key: rgb_seq[i][-1], depth_key: depth_seq[i][-1]}) + seq_length = rgb_seq.shape[1] + # just repeat the instructions for each time step + batch["instruction"] = batch["instruction"].unsqueeze(dim=1).repeat(1, seq_length, 1) + # setting it here to not trigger the hook another time and increase processing time + batch[rgb_key] = rgb_seq.to(self.device) + batch[depth_key] = depth_seq.to(self.device) + + return prev_actions + + def _create_sequence(self, episodes: list, feature_key: str): + ''' + Returns a tensor corresponding to the sequence of features. It has a shape + number of environments * sequence length * all the remaining dimensions + :param episodes: list of active environments; each environment contains the sequence of observations + :param feature_key: "rgb_features", "depth_features" + :return: + ''' + return torch.stack( + [torch.stack([obs[feature_key] for obs in ep], dim=0) for env, ep in enumerate(episodes) if len(ep) > 0], + dim=0) + + def _are_episodes_empty(self, episodes: list): + ''' + Check if the current list of steps is empty or not. + :param episodes: + :return: + ''' + return sum([len(e) > 0 for e in episodes]) == 0 + + def _filter_envs_episodes(self, + envs_to_pause, + envs, + episodes=None, + + ): + """ + Same logic as _pause_envs in the BaseTrainer, but does not pause the episodes. + Must be called before calling _pause_envs(), on lists that can fit in + _pause_envs() + Args: + envs_to_pause: + envs: + episodes: + + Returns: + + """ + # pausing envs with no new episode + if len(envs_to_pause) > 0: + state_index = list(range(envs.num_envs)) + for idx in reversed(envs_to_pause):#The envs to paused as to be done so in reverse order, otherwise, you mess up the index... + state_index.pop(idx) + if episodes is not None: + episodes = [episodes[i] for i in state_index] + return episodes + def _update_dataset(self, data_it): + """ + Cache the whole dataset. Data Aggregation can be used, the trained model + can output some action for time steps at a given probability. As the whole Task rely on an Oracle that + can output the best decision to reach the next trajectory node, even a bad decision of the model can be recovered + (imagine backtracking...), hence effectively implementing DAGGER. + :param data_it: + :return: + """ + if torch.cuda.is_available(): + with torch.cuda.device(self.device): + torch.cuda.empty_cache() + + envs = construct_envs(self.config, get_env_class(self.config.ENV_NAME)) + expert_uuid = self.config.IL.DAGGER.expert_policy_sensor_uuid + distance_left_uuid = self.config.IL.DECISION_TRANSFORMER.sensor_uuid + hidden_states = torch.zeros(envs.num_envs, 1, + dtype=torch.float) # more of a placeholder, we don t need it for the transformer + # prev_actions = torch.zeros( + # envs.num_envs, + # 1, + # device=self.device, + # dtype=torch.long, + # ) + prev_actions = None + not_done_masks = torch.zeros( + envs.num_envs, 1, dtype=torch.uint8, device=self.device + ) + + observations = envs.reset() + observations, batch = self._prepare_observation(observations) + # initialize at dim 1 for sequences of frames etc... + + episodes = [[] for _ in range(envs.num_envs)] + episode_features = [[] for _ in range(envs.num_envs)] + + skips = [False for _ in range(envs.num_envs)] + # Populate dones with False initially + dones = [False for _ in range(envs.num_envs)] + + # https://arxiv.org/pdf/1011.0686.pdf + # Theoretically, any beta function is fine so long as it converges to + # zero as data_it -> inf. The paper suggests starting with beta = 1 and + # exponential decay. + p = self.config.IL.DAGGER.p + # in Python 0.0 ** 0.0 == 1.0, but we want 0.0 + beta = 0.0 if p == 0.0 else p ** data_it + + ensure_unique_episodes = beta == 1.0 + + self._create_feature_hooks() + + # That needs to be turned to eval when doing Dagger! Not on the original implementation + self.policy.eval() + + rgb_encoder = self.policy.net.rgb_encoder + depth_encoder = self.policy.net.depth_encoder + + collected_eps = 0 + ep_ids_collected = [] + if ensure_unique_episodes: + ep_ids_collected = set() + + dataset_episodes = sum(envs.number_of_episodes) + print("Numbers of episodes in the split:", dataset_episodes) + if (self.config.IL.DAGGER.update_size > dataset_episodes and ensure_unique_episodes): + collect_size = dataset_episodes + print("Ensure unique episodes") + else: + print("Unique episodes not enforced") + collect_size = self.config.IL.DAGGER.update_size + + print(f"To be collected: {collect_size} ") + horizon = 1 + agent_action = False + + def _detect_wrong_episode(transposed_ep): + return not (transposed_ep[1][1:] == transposed_ep[2][:-1]).sum() == len(transposed_ep[1][1:]) + + collected_eps_for_real = 0 + with tqdm.tqdm( + total=collect_size, dynamic_ncols=True + ) as pbar, lmdb.open( + self.lmdb_features_dir, + map_size=int(self.config.IL.DAGGER.lmdb_map_size), + ) as lmdb_env, torch.no_grad(): + start_id = lmdb_env.stat()["entries"] + txn = lmdb_env.begin(write=True) + last_episodes = envs.current_episodes() + while collected_eps < collect_size: + envs_to_pause = [] + current_episodes = envs.current_episodes() + # if the max steps of the transform model is reached, + # and the agent does not call the stop action, force the agent to ignore the episode + if horizon == self.config.IL.DECISION_TRANSFORMER.episode_horizon: + episode_end = torch.where(actions == STOP_ACTION_TOKEN_ID, True, False) + for i in range(envs.num_envs): + if not episode_end[i] and i not in envs_to_pause: + skips[i] = True + envs_to_pause.append(i) + + for i in range(envs.num_envs): + + if dones[i] and not skips[i]: + ep = episodes[i] + traj_obs = batch_obs( + [step[0] for step in ep], + device=torch.device("cpu"), + ) + del traj_obs[expert_uuid] + for k, v in traj_obs.items(): + traj_obs[k] = v.numpy() + if self.config.IL.DAGGER.lmdb_fp16: + traj_obs[k] = traj_obs[k].astype(np.float16) + # First step: calculate the difference between 2 consecutive time steps. + # We add the initial distance to the goal to calculate the first differential reward + # traj_obs["point_nav_reward_to_go"] = np.diff( + # np.concatenate(([current_episodes[i].info["geodesic_distance"]], traj_obs[distance_left_uuid])), axis=0) + # We add the final distance to the goal once again because on th elast step, + # the STOP action is called + + traj_obs["raw_point_nav_reward"] = np.diff( + np.concatenate((traj_obs[distance_left_uuid], [traj_obs[distance_left_uuid][-1]])), + axis=0) * -1.0 + # PReparing entries for sparse rewards + traj_obs["raw_sparse_reward"] = np.zeros_like(traj_obs["raw_point_nav_reward"]) + scaling_factor = traj_obs[distance_left_uuid].size # Scaling by the episode length + del traj_obs[distance_left_uuid] + traj_obs["raw_ndtw_reward"] = np.array([step[3] for step in ep], dtype=np.float16) + self._calculate_rewards(traj_obs, scaling_factor) + transposed_ep = [ + traj_obs, + np.array([step[1] for step in ep], dtype=np.int64), + np.array([step[2] for step in ep], dtype=np.int64), + ] + + is_episode_perfect = True + + if self.config.IL.DECISION_TRANSFORMER.use_perfect_episode_only_for_dagger: + is_episode_perfect = infos[i]["success"] == 1.0 + + # don t add anything that seems weird + if is_episode_perfect: + txn.put( + str(start_id + collected_eps_for_real).encode(), + msgpack_numpy.packb( + transposed_ep, use_bin_type=True + ), + ) + collected_eps_for_real += 1 + # incrementinmg here outside the if block is not a bug + # If we can't add successfull episodes (while using dagger), + # we still need a way to exit the update function... + collected_eps += 1 + pbar.update() + + + if ( + collected_eps_for_real > 0 and collected_eps_for_real + % self.config.IL.DAGGER.lmdb_commit_frequency + ) == 0: + txn.commit() + txn = lmdb_env.begin(write=True) + + if ensure_unique_episodes: + if (not last_episodes[i].episode_id in ep_ids_collected): + ep_ids_collected.add(current_episodes[i].episode_id) + + # In opposition to the RNN logic, where only one state per time step is handled, + # We need this to force all sequences in the current batch to finish... + if dones[i]: + if i not in envs_to_pause: + envs_to_pause.append(i) + + envs_to_pause = sorted(envs_to_pause) + + episode_features = self._filter_envs_episodes(envs_to_pause, envs, episode_features) + + try: + ( + envs, + hidden_states, + not_done_masks, + prev_actions, + batch, + episodes, + ) = self._pause_envs( + envs_to_pause, + envs, + hidden_states, + not_done_masks, + prev_actions, + batch, + episodes, # A trick, I am using what is thought for the RGB features to reduce this list as well + ) + except Exception as e: + logger.warning(f"Something went wrong! Dagger It {data_it}") + for j in range(len(current_episodes)): + logger.warning(f"Current Episode culprit: {current_episodes[j].episode_id} , env {j}") + for j in range(len(last_episodes)): + logger.warning(f"Last Episode culprit: {last_episodes[j].episode_id} , env {j}") + logger.warning(f"Current horizon:{horizon}") + logger.warning(envs_to_pause) + logger.warning(f"Num envs :{envs.num_envs}") + save_file = f"ckpt.{data_it * self.config.IL.epochs}.pth" + self.save_checkpoint( + save_file + ) + logger.warning(f"Saved : {save_file}") + raise e + + if envs.num_envs == 0: + envs.resume_all() + observations = envs.reset() + # This piece of code enforce to load only episode + # not previously collected. + to_init = min((collect_size - len(ep_ids_collected)), envs.num_envs) + if ensure_unique_episodes and to_init > 0: + initialized = 0 + while initialized < to_init: + if initialized > 0: + initialized = 0 + for env, e in enumerate(envs.current_episodes()): + if e.episode_id in ep_ids_collected: + observations[env] = envs.reset_at(env)[0] + else: + initialized += 1 + current_episodes = envs.current_episodes() + episodes = [[] for _ in range(envs.num_envs)] + episode_features = [[] for _ in range(envs.num_envs)] + prev_actions = None + observations, batch = self._prepare_observation(observations) + self.rgb_features = self.rgb_features.set_(torch.zeros((1,), device="cpu")) + self.depth_features = self.depth_features.set_(torch.zeros((1,), device="cpu")) + + rgb_encoder(batch) + depth_encoder(batch) + for i in range(envs.num_envs): + if self.rgb_features is not None: + observations[i]["rgb_features"] = self.rgb_features[i] + del observations[i]["rgb"] + + if self.depth_features is not None: + observations[i]["depth_features"] = self.depth_features[i] + del observations[i]["depth"] + + prev_actions = self._modify_batch_for_transformer(episode_features, batch, self.rgb_features, + self.depth_features, envs, + prev_actions, + "rgb_features", "depth_features") + batch_size = prev_actions.shape[0] + horizon = prev_actions.shape[1] + perform_dagger = (torch.rand((batch_size, 1), dtype=torch.float) < beta).to(self.device) + # only perform dagger when the random process allows it (should lower the + # processing time...) + if perform_dagger.sum() < batch_size: + agent_action = True + actions, _ = self.policy.act( + batch, + hidden_states, + prev_actions, + not_done_masks, + deterministic=False, + ) + else: + actions = torch.ones_like(batch[expert_uuid].long()) + # actions.shape[0] == number of active enviroments + hidden_states = torch.zeros(actions.shape[0], 1, dtype=torch.float) + + actions = torch.where( + perform_dagger, + batch[expert_uuid].long(), + actions, + ) + + if self.config.IL.DECISION_TRANSFORMER.use_oracle_actions: + next_actions = batch[expert_uuid] #This is maybe a big bug for Dagger. Because if we do that like this, the sequences won't be aligned anymore + else: + next_actions = actions + # We gathered images, actions, and sensor feedback for current timestep + # time to save the timestep + + for i in range(envs.num_envs): + episodes[i].append( + ( + observations[i], + prev_actions[i, -1].item(), # this is a sequence of actions, we take the last action + next_actions[i].item(), + ) + ) + + skips = batch[ + expert_uuid].long() == -1 # looks like the short path sensor return -1 if there is a problem,, hence you need to skip an environment + actions = torch.where( + skips, torch.zeros_like(actions), actions + ) + skips = skips.squeeze(-1).to(device="cpu", non_blocking=True) + # add the last actions to the sequence of previous actions + prev_actions = torch.cat([prev_actions, actions], dim=1).to(self.device) + # prev_actions.copy_(actions) + + # When we step, environments can be reloaded automatically. + # we need to cache the previous list of episodes to be able to add them correctly in the part + # where done and not skip is applied. + last_episodes = current_episodes + outputs = envs.step([a[0].item() for a in actions]) + + observations, _, dones, infos = [list(x) for x in zip(*outputs)] + + # Just add ndtw, if you need it as Reward + for i in range(envs.num_envs): + obs, prev_act, next_act = episodes[i][-1] + episodes[i][-1] = (obs, prev_act, next_act, infos[i]["ndtw"]) + + observations, batch = self._prepare_observation(observations) + + not_done_masks = torch.tensor( + [[0] if done else [1] for done in dones], + dtype=torch.uint8, + device=self.device, + ) + + txn.commit() + + envs.close() + envs = None + + self._release_hook() + if agent_action: + print("Dataset Creation with some agent actions.") + # That needs to be turned back on... + self.policy.train() + + def inference( + self, + ): + """Evaluates a single checkpoint. + + Args: + checkpoint_path: path of checkpoint + writer: tensorboard writer object + checkpoint_index: index of the current checkpoint + """ + checkpoint_path = self.config.INFERENCE.CKPT_PATH + logger.info(f"checkpoint_path: {checkpoint_path}") + + if self.config.INFERENCE.USE_CKPT_CONFIG: + config = self._setup_eval_config( + self.load_checkpoint(checkpoint_path, map_location="cpu")[ + "config" + ] + ) + else: + config = self.config.clone() + + config.defrost() + config.TASK_CONFIG.DATASET.SPLIT = self.config.INFERENCE.SPLIT + config.TASK_CONFIG.DATASET.ROLES = ["guide"] + config.TASK_CONFIG.DATASET.LANGUAGES = config.INFERENCE.LANGUAGES + config.TASK_CONFIG.ENVIRONMENT.ITERATOR_OPTIONS.SHUFFLE = False + config.TASK_CONFIG.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT_STEPS = ( + -1 + ) + config.IL.ckpt_to_load = config.INFERENCE.CKPT_PATH + config.TASK_CONFIG.TASK.MEASUREMENTS = [] + config.TASK_CONFIG.TASK.SENSORS = [ + s for s in config.TASK_CONFIG.TASK.SENSORS if "INSTRUCTION" in s + ] + config.ENV_NAME = "VLNCEInferenceEnv" + config.freeze() + + envs = construct_envs_auto_reset_false( + config, get_env_class(config.ENV_NAME) + ) + + observation_space, action_space = self._get_spaces(config, envs=envs) + + self._initialize_policy( + config, + load_from_ckpt=True, + observation_space=observation_space, + action_space=action_space, + ) + self.policy.eval() + + self._create_feature_hooks() + rgb_encoder = self.policy.net.rgb_encoder + depth_encoder = self.policy.net.depth_encoder + + observations = envs.reset() + observations, batch = self._prepare_observation(observations) + + prev_actions = None + not_done_masks = torch.zeros( + envs.num_envs, 1, dtype=torch.uint8, device=self.device + ) + + episode_predictions = defaultdict(list) + + episodes = [[] for _ in range(envs.num_envs)] + + # episode ID --> instruction ID for rxr predictions format + instruction_ids: Dict[str, int] = {} + + episode_already_predicted = [] + + def _populate_episode_with_starting_states(): + # populate episode_predictions with the starting state + current_episodes = envs.current_episodes() + for i in range(envs.num_envs): + ep_id = current_episodes[i].episode_id + if ep_id not in episode_already_predicted: + episode_predictions[current_episodes[i].episode_id].append( + envs.call_at(i, "get_info", {"observations": {}}) + ) + if config.INFERENCE.FORMAT == "rxr": + ep_id = current_episodes[i].episode_id + k = current_episodes[i].instruction.instruction_id + instruction_ids[ep_id] = int(k) + + _populate_episode_with_starting_states() + + num_eps = sum(envs.count_episodes()) + pbar = tqdm.tqdm(total=num_eps) if hasattr(config, "use_pbar") and config.use_pbar else None + + # if all envs finishes at the same time, the operation is equal to 1. + # if all env are still processing, the operation is simply zero... + has_env_finished_early = lambda envs_that_needs_to_wait: sum(envs_that_needs_to_wait.values()) / len( + envs_that_needs_to_wait) > 0 + + while envs.num_envs > 0 and len(episode_already_predicted) < num_eps: + + current_episodes = envs.current_episodes() + # caching the outputs of the cnn on one image only + self._normalize_depth(batch) + rgb_encoder(batch) + depth_encoder(batch) + del batch["rgb"] + del batch["depth"] + rgb_key = "rgb_features" + depth_key = "depth_features" + prev_actions = self._modify_batch_for_transformer(episodes, batch, self.rgb_features, self.depth_features, + envs, + prev_actions, rgb_key, depth_key) + + with torch.no_grad(): + actions, hidden_states = self.policy.act( + batch, + None, + prev_actions, + not_done_masks, + deterministic=not config.EVAL.SAMPLE, + ) + prev_actions = torch.cat([prev_actions, actions], dim=1).to(self.device) + # prev_actions.copy_(actions) + + horizon = prev_actions.shape[1] + # if the max steps of the transform model is reached, force to end the game + if horizon == self.config.IL.DECISION_TRANSFORMER.episode_horizon: + actions[:, -1] = 0 + + outputs = envs.step([a[0].item() for a in actions]) + observations, _, dones, infos = [list(x) for x in zip(*outputs)] + # need to use a deep copy, otherwise, observations would be the same as cleaned_observations + observations, batch = self._prepare_observation(observations) + not_done_masks = torch.tensor( + [[0] if done else [1] for done in dones], + dtype=torch.uint8, + device=self.device, + ) + + # reset envs and observations if necessary + envs_that_needs_to_wait = {} + for i in range(envs.num_envs): + ep_id = current_episodes[i].episode_id + if ep_id not in episode_already_predicted: + episode_predictions[ep_id].append(infos[i]) + # This helps us to generate the transformer sequence + # episodes[i].append((cleaned_observations[i], prev_actions[i, -1].item())) + if not dones[i]: + envs_that_needs_to_wait[i] = False + continue + envs_that_needs_to_wait[i] = True + episodes[i] = [] + episode_already_predicted.append(ep_id) + # observations[i] = envs.reset_at(i)[0] + # This step is usually done in self._prepare_observation(observations) + # but now, because we amenbd only one observation, we need to take care of this step manually... + # observations[i][self.config.TASK_CONFIG.TASK.INSTRUCTION_SENSOR_UUID] = \ + # observations[i][self.config.TASK_CONFIG.TASK.INSTRUCTION_SENSOR_UUID]["tokens"] + # self.rgb_features = self.rgb_features.set_(torch.zeros((1,), device="cpu")) + # self.depth_features = self.depth_features.set_(torch.zeros((1,), device="cpu")) + # observations, batch = self._prepare_observation(observations) + + if pbar: + pbar.update() + + envs_to_pause = [] + next_episodes = envs.current_episodes() + + for i in range(envs.num_envs): + if next_episodes[i].episode_id in episode_already_predicted: + envs_to_pause.append(i) + elif has_env_finished_early(envs_that_needs_to_wait): + if envs_that_needs_to_wait[i] and i not in envs_to_pause: + envs_to_pause.append(i) + episodes = self._filter_envs_episodes(envs_to_pause, envs, episodes) + ( + envs, + hidden_states, + not_done_masks, + prev_actions, + batch, + _, + ) = self._pause_envs( + envs_to_pause, + envs, + hidden_states, + not_done_masks, + prev_actions, + batch, + None, + ) + + # at this stage, if we dont have any env left, + # that means that all prediction within the same "batch" + # are finished, we can wake all envs now. + if envs.num_envs < 1: + envs.resume_all() + episodes = [[] for _ in range(envs.num_envs)] + observations = envs.reset() + prev_actions = None + observations, batch = self._prepare_observation(observations) + _populate_episode_with_starting_states() + + envs.close() + gc.collect() + self._release_hook() + if pbar: + pbar.close() + + if config.INFERENCE.FORMAT == "r2r": + with open(config.INFERENCE.PREDICTIONS_FILE, "w") as f: + json.dump(episode_predictions, f, indent=2) + + logger.info( + f"Predictions saved to: {config.INFERENCE.PREDICTIONS_FILE}" + ) + else: # use 'rxr' format for rxr-habitat leaderboard + predictions_out = [] + + for k, v in episode_predictions.items(): + + # save only positions that changed + path = [v[0]["position"]] + for p in v[1:]: + if path[-1] != p["position"]: + path.append(p["position"]) + + predictions_out.append( + { + "instruction_id": instruction_ids[k], + "path": path, + } + ) + + predictions_out.sort(key=lambda x: x["instruction_id"]) + with jsonlines.open( + config.INFERENCE.PREDICTIONS_FILE, mode="w" + ) as writer: + writer.write_all(predictions_out) + + logger.info( + f"Predictions saved to: {config.INFERENCE.PREDICTIONS_FILE}" + ) + + def _eval_checkpoint( + self, + checkpoint_path: str, + writer: TensorboardWriter, + checkpoint_index: int = 0, + ) -> None: + """Evaluates a single checkpoint. + + Args: + checkpoint_path: path of checkpoint + writer: tensorboard writer object + checkpoint_index: index of the current checkpoint + """ + logger.info(f"checkpoint_path: {checkpoint_path}") + + config = self.config.clone() + if self.config.EVAL.USE_CKPT_CONFIG: + ckpt = self.load_checkpoint(checkpoint_path, map_location="cpu") + config = self._setup_eval_config(ckpt) + + split = config.EVAL.SPLIT + + config.defrost() + config.TASK_CONFIG.DATASET.SPLIT = split + config.TASK_CONFIG.DATASET.ROLES = ["guide"] + config.TASK_CONFIG.DATASET.LANGUAGES = config.EVAL.LANGUAGES + config.TASK_CONFIG.TASK.NDTW.SPLIT = split + config.TASK_CONFIG.ENVIRONMENT.ITERATOR_OPTIONS.SHUFFLE = False + config.TASK_CONFIG.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT_STEPS = ( + -1 + ) + config.IL.ckpt_to_load = checkpoint_path + config.use_pbar = not is_slurm_batch_job() + + if len(config.VIDEO_OPTION) > 0: + config.TASK_CONFIG.TASK.MEASUREMENTS.append("TOP_DOWN_MAP_VLNCE") + + config.freeze() + + if config.EVAL.SAVE_RESULTS: + model_file = checkpoint_path.split("/")[-1] + base_name = f"ckpt_{checkpoint_index}" + if model_file != base_name: + base_name = model_file + fname = os.path.join( + config.RESULTS_DIR, + f"stats_{base_name}_{split}.json", + ) + if os.path.exists(fname): + logger.info(f"skipping {base_name} -- evaluation exists.") + return + + envs = construct_envs_auto_reset_false( + config, get_env_class(config.ENV_NAME) + ) + observation_space, action_space = self._get_spaces(config, envs=envs) + + self._initialize_policy( + config, + load_from_ckpt=True, + observation_space=observation_space, + action_space=action_space, + ) + self.policy.eval() + + self._create_feature_hooks() + rgb_encoder = self.policy.net.rgb_encoder + depth_encoder = self.policy.net.depth_encoder + + observations = envs.reset() + observations, batch = self._prepare_observation(observations) + + hidden_states = torch.zeros(envs.num_envs, 1, dtype=torch.float) + + prev_actions = None + not_done_masks = torch.zeros( + envs.num_envs, 1, dtype=torch.uint8, device=self.device + ) + + stats_episodes = {} + + rgb_frames = [[] for _ in range(envs.num_envs)] + episodes = [[] for _ in range(envs.num_envs)] + if len(config.VIDEO_OPTION) > 0: + os.makedirs(config.VIDEO_DIR, exist_ok=True) + + num_eps = sum(envs.number_of_episodes) + if config.EVAL.EPISODE_COUNT > -1: + num_eps = min(config.EVAL.EPISODE_COUNT, num_eps) + + pbar = tqdm.tqdm(total=num_eps) if config.use_pbar else None + log_str = ( + f"[Ckpt: {checkpoint_index}]" + " [Episodes evaluated: {evaluated}/{total}]" + " [Time elapsed (s): {time}]" + ) + start_time = time.time() + + # if all envs finishes at the same time, the operation is equal to 1. + # if all env are still processing, the operation is simply zero... + has_env_finished_early = lambda envs_that_needs_to_wait: sum(envs_that_needs_to_wait.values()) / len( + envs_that_needs_to_wait) > 0 + + while envs.num_envs > 0 and len(stats_episodes) < num_eps: + current_episodes = envs.current_episodes() + + # caching the outputs of the cnn on one image only + self._normalize_depth(batch) + rgb_encoder(batch) + depth_encoder(batch) + del batch["rgb"] + del batch["depth"] + rgb_key = "rgb_features" + depth_key = "depth_features" + prev_actions = self._modify_batch_for_transformer(episodes, batch, self.rgb_features, self.depth_features, + envs, + prev_actions, rgb_key, depth_key) + + with torch.no_grad(): + actions, hidden_states = self.policy.act( + batch, + None, + prev_actions, + not_done_masks, + deterministic=not config.EVAL.SAMPLE, + ) + prev_actions = torch.cat([prev_actions, actions], dim=1).to(self.device) + # prev_actions.copy_(actions) + + horizon = prev_actions.shape[1] + # if the max steps of the transform model is reached, force to end the game + if horizon == self.config.IL.DECISION_TRANSFORMER.episode_horizon: + actions[:, -1] = 0 + + outputs = envs.step([a[0].item() for a in actions]) + observations, _, dones, infos = [list(x) for x in zip(*outputs)] + # need to use a deep copy, otherwise, observations would be the same as cleaned_observations + observations, batch = self._prepare_observation(observations) + not_done_masks = torch.tensor( + [[0] if done else [1] for done in dones], + dtype=torch.uint8, + device=self.device, + ) + + # reset envs and observations if necessary + envs_that_needs_to_wait = {} + for i in range(envs.num_envs): + if len(config.VIDEO_OPTION) > 0: + frame = observations_to_image(observations[i], infos[i]) + frame = append_text_to_image( + frame, current_episodes[i].instruction.instruction_text + ) + rgb_frames[i].append(frame) + # This helps us to generate the transformer sequence + # episodes[i].append((cleaned_observations[i], prev_actions[i, -1].item())) + if not dones[i]: + envs_that_needs_to_wait[i] = False + continue + envs_that_needs_to_wait[i] = True + episodes[i] = [] + ep_id = current_episodes[i].episode_id + stats_episodes[ep_id] = infos[i] + # observations[i] = envs.reset_at(i)[0] + # This step is usually done in self._prepare_observation(observations) + # but now, because we amenbd only one observation, we need to take care of this step manually... + # observations[i][self.config.TASK_CONFIG.TASK.INSTRUCTION_SENSOR_UUID] = \ + # observations[i][self.config.TASK_CONFIG.TASK.INSTRUCTION_SENSOR_UUID]["tokens"] + # self.rgb_features = self.rgb_features.set_(torch.zeros((1,), device="cpu")) + # self.depth_features = self.depth_features.set_(torch.zeros((1,), device="cpu")) + # observations, batch = self._prepare_observation(observations) + + if config.use_pbar: + pbar.update() + else: + logger.info( + log_str.format( + evaluated=len(stats_episodes), + total=num_eps, + time=round(time.time() - start_time), + ) + ) + + if len(config.VIDEO_OPTION) > 0: + generate_video( + video_option=config.VIDEO_OPTION, + video_dir=config.VIDEO_DIR, + images=rgb_frames[i], + episode_id=ep_id, + checkpoint_idx=checkpoint_index, + metrics={"spl": stats_episodes[ep_id]["spl"]}, + tb_writer=writer, + ) + del stats_episodes[ep_id]["top_down_map_vlnce"] + rgb_frames[i] = [] + + envs_to_pause = [] + next_episodes = envs.current_episodes() + + for i in range(envs.num_envs): + if next_episodes[i].episode_id in stats_episodes: + envs_to_pause.append(i) + if has_env_finished_early(envs_that_needs_to_wait): + if envs_that_needs_to_wait[i] and i not in envs_to_pause: + envs_to_pause.append(i) + + episodes = self._filter_envs_episodes(envs_to_pause, envs, episodes) + ( + envs, + hidden_states, + not_done_masks, + prev_actions, + batch, + rgb_frames, + ) = self._pause_envs( + envs_to_pause, + envs, + hidden_states, + not_done_masks, + prev_actions, + batch, + rgb_frames, + ) + + # at this stage, if we dont have any env left, + # that means that all prediction within the same "batch" + # are finished, we can wake all envs now. + if envs.num_envs < 1: + envs.resume_all() + episodes = [[] for _ in range(envs.num_envs)] + rgb_frames = [[] for _ in range(envs.num_envs)] + observations = envs.reset() + prev_actions = None + observations, batch = self._prepare_observation(observations) + + envs.close() + gc.collect() + self._release_hook() + if config.use_pbar: + pbar.close() + + aggregated_stats = {} + num_episodes = len(stats_episodes) + for k in next(iter(stats_episodes.values())).keys(): + aggregated_stats[k] = ( + sum(v[k] for v in stats_episodes.values()) / num_episodes + ) + + if config.EVAL.SAVE_RESULTS: + with open(fname, "w") as f: + json.dump(aggregated_stats, f, indent=4) + + logger.info(f"Episodes evaluated: {num_episodes}") + checkpoint_num = checkpoint_index + 1 + for k, v in aggregated_stats.items(): + logger.info(f"{k}: {v:.6f}") + writer.add_scalar(f"eval_{split}_{k}", v, checkpoint_num) + + def train(self) -> None: + """Main method for training DAgger.""" + if self.config.IL.DAGGER.preload_lmdb_features: + try: + lmdb.open(self.lmdb_features_dir, readonly=True) + except lmdb.Error as err: + logger.error( + "Cannot open database for teacher forcing preload." + ) + raise err + else: + with lmdb.open( + self.lmdb_features_dir, + map_size=int(self.config.IL.DAGGER.lmdb_map_size), + ) as lmdb_env, lmdb_env.begin(write=True) as txn: + txn.drop(lmdb_env.open_db()) + + EPS = self.config.IL.DAGGER.expert_policy_sensor + if EPS not in self.config.TASK_CONFIG.TASK.SENSORS: + self.config.TASK_CONFIG.TASK.SENSORS.append(EPS) + + self.config.defrost() + + # if doing teacher forcing, don't switch the scene until it is complete + if self.config.IL.DAGGER.p == 1.0: + self.config.TASK_CONFIG.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT_STEPS = ( + -1 + ) + self.config.freeze() + + def collate_fn(batch): + """Each sample in batch: ( + obs, + prev_actions, + oracle_actions, + inflec_weight, + ) + """ + + def _pad_helper(t, max_len, fill_val=0): + pad_amount = max_len - t.size(0) + if pad_amount == 0: + return t + + pad = torch.full_like(t[0:1], fill_val).expand( + pad_amount, *t.size()[1:] + ) + return torch.cat([t, pad], dim=0) + + if not _is_correct_previous_actions(batch) and not self.config.IL.DECISION_TRANSFORMER.use_oracle_actions: + raise Exception( + "Dataset has not been created correctly! Prev actions and corrected actions not shifted accordingly!") + transposed = list(zip(*batch)) + observations_batch = list(transposed[0]) + + if self.config.IL.DECISION_TRANSFORMER.recompute_reward: + for o in observations_batch: + scaling_factor = len(o["raw_sparse_reward"]) + self._calculate_rewards(o, scaling_factor) + + prev_actions_batch = list(transposed[1]) + corrected_actions_batch = list(transposed[2]) + weights_batch = list(transposed[3]) # to make it batch * seq length + batch_size = len(prev_actions_batch) + + new_observations_batch = defaultdict(list) + for sensor in observations_batch[0]: + for bid in range(batch_size): + new_observations_batch[sensor].append( + observations_batch[bid][sensor] + ) + + observations_batch = new_observations_batch + + max_traj_len = max(ele.size(0) for ele in prev_actions_batch) + for bid in range(batch_size): + for sensor in observations_batch: + fill = 0.0 if "_reward" in sensor else 1.0 + # Workaround when the reward is only a single scalar... + if len(observations_batch[sensor][bid].shape) == 0: + observations_batch[sensor][bid] = observations_batch[sensor][bid].unsqueeze(-1) + observations_batch[sensor][bid] = _pad_helper( + observations_batch[sensor][bid], max_traj_len, fill_val=fill + ) + + prev_actions_batch[bid] = _pad_helper( + prev_actions_batch[bid], max_traj_len + ) + corrected_actions_batch[bid] = _pad_helper( + corrected_actions_batch[bid], max_traj_len + ) + weights_batch[bid] = _pad_helper(weights_batch[bid], max_traj_len) + + stack_dimension = 0 + + for sensor in observations_batch: + observations_batch[sensor] = torch.stack(observations_batch[sensor], dim=stack_dimension) + if "_reward" in sensor: + observations_batch[sensor] = observations_batch[sensor].unsqueeze(-1) + + prev_actions_batch = torch.stack(prev_actions_batch, dim=stack_dimension) + corrected_actions_batch = torch.stack(corrected_actions_batch, dim=stack_dimension) + + weights_batch = torch.stack(weights_batch, dim=stack_dimension) + not_done_masks = torch.ones_like( + corrected_actions_batch, dtype=torch.uint8 + ) + not_done_masks[0] = 0 + + if self.config.MODEL.DECISION_TRANSFORMER.use_extra_start_token: + # The environment only use actions from 0 to 3, the 4 is just a + # a dummy token to indicate the beginning of a sequence. + prev_actions_batch[:, 0] = EXTRA_START_TOKEN_ID + else: + prev_actions_batch[:, 0] = STOP_ACTION_TOKEN_ID # this is zero. + # shape batch size time max episode length + timesteps = torch.arange(0, max_traj_len).repeat(batch_size, 1) + observations_batch["timesteps"] = timesteps + observations_batch = ObservationsDict(observations_batch) + + return ( + observations_batch, + prev_actions_batch, + not_done_masks, + corrected_actions_batch, + weights_batch + ) + + observation_space, action_space = self._get_spaces(self.config) + + self._initialize_policy( + self.config, + self.config.IL.load_from_ckpt, + observation_space=observation_space, + action_space=action_space, + ) + # Seems to bottleneck on Dataloader access if I have more than 1 worker + workers = self.config.IL.dataload_workers + + # Tries to name the next checkpoints correctly based on the loaded file + start_epoch = 0 + if self.config.IL.load_from_ckpt and self.config.IL.continue_ckpt_naming: + checkpoint_name = self.config.IL.ckpt_to_load.split("/")[-1] + epochs = re.findall(r"\d+", checkpoint_name) + if len(epochs) > 0: + start_epoch = int(epochs[0]) + 1 + + with TensorboardWriter( + self.config.TENSORBOARD_DIR, + flush_secs=self.flush_secs, + purge_step=0, + ) as writer: + for dagger_it in range(self.config.IL.DAGGER.iterations): + step_id = 0 + if not self.config.IL.DAGGER.preload_lmdb_features: + update_id = dagger_it + (1 if self.config.IL.load_from_ckpt else 0) + self._update_dataset( + update_id + ) + + if torch.cuda.is_available(): + with torch.cuda.device(self.device): + torch.cuda.empty_cache() + gc.collect() + + dataset = IWTrajectoryDataset( + self.lmdb_features_dir, + self.config.IL.use_iw, + inflection_weight_coef=self.config.IL.inflection_weight_coef, + lmdb_map_size=self.config.IL.DAGGER.lmdb_map_size, + batch_size=self.config.IL.batch_size, + preload_size=self.config.IL.preload_dataloader_size, + ) + diter = torch.utils.data.DataLoader( + dataset, + batch_size=self.config.IL.batch_size, + shuffle=False, + collate_fn=collate_fn, + pin_memory=False, + drop_last=True, # drop last batch if smaller + num_workers=workers, + ) + num_batch = dataset.length // dataset.batch_size + print("DAGGER Iteration", dagger_it, "dataset length: ", dataset.length) + if num_batch == 0: + num_batch = 1 + logger.info(f"Number of batches to process:{num_batch}") + # AuxLosses.activate() + for epoch in tqdm.trange( + self.config.IL.epochs, dynamic_ncols=True + ): + total_loss = 0.0 + for batch in tqdm.tqdm( + diter, + total=num_batch, + leave=False, + dynamic_ncols=True, + ): + epoch = start_epoch + epoch + ( + observations_batch, + prev_actions_batch, + not_done_masks, + corrected_actions_batch, + weights_batch, + ) = batch + + observations_batch = { + k: (v.to( + device=self.device, + non_blocking=True, + ) if v.dtype == torch.long else v.to(device=self.device, dtype=torch.float32, + non_blocking=True)) + for k, v in observations_batch.items() + } + + loss, action_loss, aux_loss = self._update_agent( + observations_batch, + prev_actions_batch.to( + device=self.device, non_blocking=True + ), + not_done_masks.to( + device=self.device, non_blocking=True + ), + corrected_actions_batch.to( + device=self.device, non_blocking=True + ), + weights_batch.to( + device=self.device, non_blocking=True + ), + ) + + writer.add_scalar( + f"train_loss_iter_{dagger_it}", loss, step_id + ) + writer.add_scalar( + f"train_action_loss_iter_{dagger_it}", + action_loss, + step_id, + ) + total_loss += loss + step_id += 1 # noqa: SIM113 + total_loss = total_loss / (num_batch) + writer.add_scalar( + f"train_total_loss{dagger_it}", + total_loss, + epoch, + ) + logger.info(f"Mean Loss for DAgger iter {dagger_it}, Epoch {epoch}: {total_loss}") + if total_loss <= self.config.IL.mean_loss_to_save_checkpoint and ( + epoch + 1) % self.config.IL.checkpoint_frequency == 0: + print("Save", f"ckpt.{dagger_it * self.config.IL.epochs + epoch}.pth") + self.save_checkpoint( + f"ckpt.{dagger_it * self.config.IL.epochs + epoch}.pth" + ) + if total_loss <= self.config.IL.mean_loss_to_stop_training: + logger.info(f"Stopping training early at epoch {epoch}") + break + # AuxLosses.deactivate() + + + + def _update_agent( + self, + observations, + prev_actions, + not_done_masks, + corrected_actions, + weights, + step_grad: bool = True, + loss_accumulation_scalar: int = 1, + ): + """ + Returns the agent loss in the training loop + :param observations: + :param prev_actions: + :param not_done_masks: + :param corrected_actions: + :param weights: + :param step_grad: + :param loss_accumulation_scalar: + :return: + """ + + hidden_states = None + + distribution = self.policy.build_distribution( + observations, hidden_states, prev_actions, not_done_masks + ) + + logits = distribution.logits + + # The permutation allows to keep the expected input shape (batch times classes) + # the third dimension gets interpreted as a sequence correctly, as the target actions + # have the correct shape + action_loss = F.cross_entropy( + logits.permute(0, 2, 1), corrected_actions, reduction="none" + ) + action_loss = ((weights * action_loss).sum(0) / weights.sum(0)).mean() + + aux_loss = 0.0 + loss = action_loss + loss = loss / loss_accumulation_scalar + loss.backward() + + if step_grad: + self.optimizer.step() + self.optimizer.zero_grad() + + if isinstance(aux_loss, torch.Tensor): + aux_loss = aux_loss.item() + return loss.item(), action_loss.item(), aux_loss diff --git a/vlnce_baselines/models/decision_transformer_policy.py b/vlnce_baselines/models/decision_transformer_policy.py new file mode 100644 index 00000000..e5c4bacc --- /dev/null +++ b/vlnce_baselines/models/decision_transformer_policy.py @@ -0,0 +1,634 @@ +import torch +import torch.nn as nn +from gym import Space +from habitat import Config +from habitat_baselines.common.baseline_registry import BaselineRegistry +from habitat_baselines.rl.ppo.policy import Net +from vlnce_baselines.models.encoders.min_gpt import GPT, NewGELU +from vlnce_baselines.models.encoders import resnet_encoders +from vlnce_baselines.models.encoders.instruction_encoder import ( + InstructionEncoder, Word2VecEmbeddings, InstructionEncoderWithTransformer +) +from vlnce_baselines.models.policy import ILPolicy +import numpy as np +from torch import Tensor + +from vlnce_baselines.models.utils import PositionalEncoding, VanillaMultiHeadAttention + + +@BaselineRegistry.register_policy +class DecisionTransformerPolicy(ILPolicy): + def __init__( + self, + observation_space: Space, + action_space: Space, + model_config: Config, + ): + net = "DecisionTransformerNet" + if hasattr(model_config.DECISION_TRANSFORMER, "net"): + net = model_config.DECISION_TRANSFORMER.net + assert net in model_config.DECISION_TRANSFORMER.allowed_models + print("Training with:", net) + super().__init__( + eval(net)( + observation_space=observation_space, + model_config=model_config, + num_actions=action_space.n, + ), + action_space.n, + ) + + def act( + self, + observations, + rnn_states, + prev_actions, + masks, + deterministic=False, + ): + actions, rnn_states = super().act(observations, rnn_states, prev_actions, masks, deterministic) + # We just want to return the last action of the transformer sequence... + return actions[:, -1, :], rnn_states + + @classmethod + def from_config( + cls, config: Config, observation_space: Space, action_space: Space + ): + config.defrost() + config.MODEL.TORCH_GPU_ID = config.TORCH_GPU_ID + config.freeze() + + return cls( + observation_space=observation_space, + action_space=action_space, + model_config=config.MODEL, + ) + + +class AbstractDecisionTransformerNet(Net): + # Decision Transformer where each time step is fed into a GPT backbone. + # Finally, a distribution over discrete actions (FWD, L, R, STOP) is produced. + def __init__( + self, observation_space: Space, model_config: Config, num_actions: int + ): + """ + + :param observation_space: Delivered by the Habitat Framework + :param model_config: General config + :param num_actions: 4 discrete actions (FWD, L, R, STOP) + """ + super().__init__() + self.model_config = model_config + assert model_config.DEPTH_ENCODER.cnn_type in ["VlnResnetDepthEncoder"] + assert model_config.RGB_ENCODER.cnn_type in [ + "TorchVisionResNet18", + "TorchVisionResNet50", + ] + + assert model_config.DECISION_TRANSFORMER.reward_type in model_config.DECISION_TRANSFORMER.allowed_rewards + + n = self.initialize_transformer_step_size() + self.set_transformer_step_size(n) + # Init the Depth visual encoder + self.depth_encoder = getattr( + resnet_encoders, model_config.DEPTH_ENCODER.cnn_type + )( + observation_space, + output_size=model_config.DEPTH_ENCODER.output_size, + checkpoint=model_config.DEPTH_ENCODER.ddppo_checkpoint, + backbone=model_config.DEPTH_ENCODER.backbone, + trainable=model_config.DEPTH_ENCODER.trainable, + spatial_output=model_config.DECISION_TRANSFORMER.spatial_output + ) + # Init the RGB visual encoder + self.rgb_encoder = getattr( + resnet_encoders, model_config.RGB_ENCODER.cnn_type + )( + model_config.RGB_ENCODER.output_size, + normalize_visual_inputs=model_config.DECISION_TRANSFORMER.normalize_rgb, + trainable=model_config.RGB_ENCODER.trainable, + spatial_output=model_config.DECISION_TRANSFORMER.spatial_output + ) + + if model_config.DECISION_TRANSFORMER.spatial_output: + self.rgb_linear = nn.Sequential( + nn.AdaptiveAvgPool1d(1), + nn.Flatten(), + nn.Linear( + self.rgb_encoder.output_shape[0], + model_config.RGB_ENCODER.output_size, + ), + nn.ReLU(True), + ) + self.depth_linear = nn.Sequential( + nn.Flatten(), + nn.Linear( + np.prod(self.depth_encoder.output_shape), + model_config.DEPTH_ENCODER.output_size, + ), + nn.ReLU(True), + ) + + self.dim_not_included_for_predictions = 2 # Action and reward + self.exclude_past_action_for_prediction = model_config.DECISION_TRANSFORMER.exclude_past_action_for_prediction + if not self.exclude_past_action_for_prediction: + self.dim_not_included_for_predictions = 1 + self.return_to_go_inference = model_config.DECISION_TRANSFORMER.return_to_go_inference + + self.initialize_instruction_encoder() + + self.reward_type = model_config.DECISION_TRANSFORMER.reward_type + self.action_activation = nn.Sequential( + nn.Dropout(p=self.model_config.DECISION_TRANSFORMER.activation_action_drop), NewGELU()) + self.instruction_activation = nn.Sequential( + nn.Dropout(p=self.model_config.DECISION_TRANSFORMER.activation_instruction_drop), NewGELU()) + self.rgb_activation = nn.Sequential( + nn.Dropout(p=self.model_config.DECISION_TRANSFORMER.activation_rgb_drop), NewGELU()) + self.depth_activation = nn.Sequential( + nn.Dropout(p=self.model_config.DECISION_TRANSFORMER.activation_depth_drop), NewGELU()) + self.gpt_encoder = GPT(self.model_config.DECISION_TRANSFORMER) + self.transformer_step_size = self.model_config.DECISION_TRANSFORMER.step_size + self.embed_timestep = nn.Embedding(model_config.DECISION_TRANSFORMER.episode_horizon, + model_config.DECISION_TRANSFORMER.hidden_dim) + self.embed_return = nn.Linear(1, model_config.DECISION_TRANSFORMER.hidden_dim) + self.embed_action = nn.Embedding(num_actions + 1, model_config.DECISION_TRANSFORMER.hidden_dim) + self.embed_ln = nn.LayerNorm(model_config.DECISION_TRANSFORMER.hidden_dim) + self.initialize_other_layers() + self.train() + + def _prepare_embeddings(self, observations): + """ + read the relevant features from observation and returns it + :param observations: + :return: instruction_embedding, depth_embedding, rgb_embedding + """ + # for all the following keys, we need tto merge the first 2 dimensions + # [batch, sequence length, all other dimensions] to [batch * sequence length, all other dimensions] + + original_batch_shape = observations["instruction"].shape[0:2] # excluding the embedding dimentions + batch_size, seq_length = original_batch_shape + # the observations were flattened for rnn processing + # the first dimension is actually equal to sequence length * original batch size. + # we also retrieve all other dimensions starting at index 1 + shape = lambda tensor: tuple([s for s in original_batch_shape] + [s for s in tensor.shape[1:]]) + + # Transpose dimension 0 and 1 and let the last one untouched + # resize_tensor = lambda tensor: tensor.reshape(shape(tensor)).permute(1,0,-1).contiguous() + resize_tensor = lambda tensor: tensor.reshape(shape(tensor)) + + self._flatten_batch(observations, "rgb") + self._flatten_batch(observations, "depth") + self._flatten_batch(observations, "rgb_features") + self._flatten_batch(observations, "depth_features") + if not self.model_config.DECISION_TRANSFORMER.use_transformer_encoded_instruction: + self._flatten_batch(observations, "instruction") + + depth_embedding = self.depth_encoder(observations) + rgb_embedding = self.rgb_encoder(observations) + if self.model_config.DECISION_TRANSFORMER.spatial_output: + depth_embedding = self.depth_linear( + torch.flatten(depth_embedding, 2)) + rgb_embedding = self.rgb_linear( + torch.flatten(rgb_embedding, 2)) + depth_embedding = self.depth_activation(depth_embedding) + rgb_embedding = self.rgb_activation(rgb_embedding) + # we just undo the permutation made in the original implementation + instruction_embedding = self.handle_instruction_embeddings(observations, resize_tensor, batch_size, seq_length) + + depth_embedding = resize_tensor(depth_embedding) + rgb_embedding = resize_tensor(rgb_embedding) + + if self.model_config.ablate_instruction: + instruction_embedding = instruction_embedding * 0 + if self.model_config.ablate_depth: + depth_embedding = depth_embedding * 0 + if self.model_config.ablate_rgb: + rgb_embedding = rgb_embedding * 0 + + return instruction_embedding, depth_embedding, rgb_embedding + + def handle_instruction_embeddings(self, observations, resize_tensor, batch_size, seq_length): + raise not NotImplementedError("Depending, if you get a sentence embedding or the whole word sequence!") + + def initialize_instruction_encoder(self): + raise not NotImplementedError("Should set instruction encoder used by your model!") + + def initialize_other_layers(self): + raise not NotImplementedError("Should set the layers used by your model!") + + def initialize_transformer_step_size(self): + raise not NotImplementedError("Should return the value needed for set_transformer_step_size(self, n) ") + + def create_tensors_for_gpt_as_tuple(self, prev_actions, returns_to_go, instruction_embedding, + depth_embedding, rgb_embedding, timesteps, batch_size, + seq_length): + raise not NotImplementedError( + "do the mo del specific work and return a tuple of tensors like (Action, S1, ... Sn, Reward)") + + def set_transformer_step_size(self, n): + self.model_config.defrost() + # a step has a size of 2 + n + # Actions, State 1, State 2... State n, Reward + self.model_config.DECISION_TRANSFORMER.step_size = n + self.model_config.freeze() + + def _flatten_batch(self, observations: Tensor, sensor_type: str): + + # quit silently + if not sensor_type in observations.keys(): + return + + dims = observations[sensor_type].size() + if len(dims) > 2: + observations[sensor_type] = observations[sensor_type].view(-1, *dims[2:]) + + @property + def output_size(self): + steps = max(1, (self.transformer_step_size - self.dim_not_included_for_predictions)) + + return self.model_config.DECISION_TRANSFORMER.hidden_dim * steps # - 2 because we exclude reward / actions for categorical layer + + def create_timesteps(self, sequence_length, batch_size): + + # TODO: use buffer? + timesteps = [torch.arange(0, sequence_length, dtype=torch.long) for _ in range(batch_size)] + timesteps = torch.stack(timesteps, dim=0).to(self.embed_ln.weight.device) + + return timesteps + + @property + def is_blind(self): + return self.rgb_encoder.is_blind or self.depth_encoder.is_blind + + @property + def num_recurrent_layers(self): + return self.state_encoder.num_recurrent_layers + + def forward(self, observations, rnn_states, prev_actions, masks): + original_batch_shape = observations["instruction"].shape[0:2] # excluding the embedding dimentions + batch_size, seq_length = original_batch_shape + + instruction_embedding, depth_embedding, rgb_embedding = self._prepare_embeddings(observations) + + if self.reward_type in observations.keys() and self.training: + returns_to_go = observations[self.reward_type] + else: + # If we don t have any rewards from the environment, just take one + # as mentioned in the paper during evaluation. + returns_to_go = torch.ones_like(prev_actions, dtype=torch.float).unsqueeze( + dim=-1) * self.return_to_go_inference + if "timesteps" in observations.keys(): + timesteps = observations["timesteps"] + else: + timesteps = self.create_timesteps(seq_length, batch_size) + + # squeeze to output the same shape as other embeddings + # after the operation with embedding layer + if len(timesteps.shape) > 2: + timesteps = timesteps.squeeze(-1) + + tensor_tuples = self.create_tensors_for_gpt_as_tuple(prev_actions, returns_to_go, instruction_embedding, + depth_embedding, rgb_embedding, timesteps, batch_size, + seq_length) + + stacked = ( + torch.stack(tensor_tuples, dim=1).permute(0, 2, 1, 3).reshape(batch_size, + self.transformer_step_size * seq_length, -1) + ) + + output = self.gpt_encoder(self.embed_ln(stacked)) + output = output.reshape(batch_size, seq_length, self.transformer_step_size, -1).permute(0, 2, 1, 3) + + start_dim = 1 + if not self.exclude_past_action_for_prediction: + start_dim = 0 + + + end_dim = max(start_dim + 1, self.transformer_step_size - 1) + # get predictions + action_preds = output[:, start_dim:end_dim].permute(0, 2, 1, 3).reshape(batch_size, + seq_length, + -1) + + return action_preds, None + + +class DecisionTransformerNet(AbstractDecisionTransformerNet): + """Decision Transformer, where RGB, DEPTH and Instructions are concatenated into one state. + """ + + def __init__( + self, observation_space: Space, model_config: Config, num_actions: int + ): + super().__init__(observation_space, model_config, num_actions) + + def handle_instruction_embeddings(self, observations, resize_tensor, batch_size, seq_length): + + instruction_embedding = self.instruction_activation(self.instruction_encoder(observations)) + + if not self.model_config.DECISION_TRANSFORMER.use_transformer_encoded_instruction: + instruction_embedding = resize_tensor(instruction_embedding) + else: + instruction_embedding = self.sentence_encoding(instruction_embedding.permute(0, 2, 1)).permute(0, 2, + 1).repeat( + (1, seq_length, 1)) + return instruction_embedding + + def initialize_other_layers(self): + # size due to concatenation of instruction, depth, and rgb features + input_state_size = self.instruction_encoder.output_size + self.model_config.DEPTH_ENCODER.output_size + self.model_config.RGB_ENCODER.output_size + self.embed_state = nn.Linear(input_state_size, self.model_config.DECISION_TRANSFORMER.hidden_dim) + + def create_tensors_for_gpt_as_tuple(self, prev_actions, returns_to_go, instruction_embedding, + depth_embedding, rgb_embedding, timesteps, batch_size, + seq_length): + states = torch.cat( + [instruction_embedding, depth_embedding, rgb_embedding], dim=2 + ) + # embed each modality with a different head + state_embeddings = self.embed_state(states) + action_embeddings = self.action_activation(self.embed_action(prev_actions)) + returns_embeddings = self.embed_return(returns_to_go) + time_embeddings = self.embed_timestep(timesteps) + + state_embeddings2 = state_embeddings + time_embeddings + action_embeddings2 = action_embeddings + time_embeddings + returns_embeddings2 = returns_embeddings + time_embeddings + + return returns_embeddings2, state_embeddings2, action_embeddings2 + + def initialize_instruction_encoder(self): + if not self.model_config.DECISION_TRANSFORMER.use_transformer_encoded_instruction: + # Init the instruction encoder + self.instruction_encoder = InstructionEncoder( + self.model_config.INSTRUCTION_ENCODER + ) + else: + self.instruction_encoder = InstructionEncoderWithTransformer(self.model_config) + if self.model_config.DECISION_TRANSFORMER.ENCODER.use_sentence_encoding: + self.sentence_encoding = nn.AdaptiveAvgPool1d(1) + + def initialize_transformer_step_size(self): + # reward, state, action + return 3 + + +class DecisionTransformerEnhancedNet(DecisionTransformerNet): + + def __init__( + self, observation_space: Space, model_config: Config, num_actions: int + ): + super().__init__(observation_space, model_config, num_actions) + + def initialize_other_layers(self): + out_dim = self.model_config.DECISION_TRANSFORMER.hidden_dim + self.instruction_embed_state = nn.Linear(self.instruction_encoder.output_size, + out_dim) + self.rgb_embed_state = nn.Linear(self.model_config.RGB_ENCODER.output_size, + out_dim) + self.depth_embed_state = nn.Linear(self.model_config.DEPTH_ENCODER.output_size, + out_dim) + + def initialize_transformer_step_size(self): + return 5 + + def create_tensors_for_gpt_as_tuple(self, prev_actions, returns_to_go, instruction_embedding, + depth_embedding, rgb_embedding, timesteps, batch_size, + seq_length): + instruction_state_embeddings = self.instruction_embed_state(instruction_embedding) + rgb_state_embeddings = self.rgb_embed_state(rgb_embedding) + depth_state_embeddings = self.depth_embed_state(depth_embedding) + + action_embeddings = self.embed_action(prev_actions) + returns_embeddings = self.embed_return(returns_to_go) + time_embeddings = self.embed_timestep(timesteps) + + # print(state_embeddings.shape, action_embeddings.shape, returns_embeddings.shape, time_embeddings.shape) + # time embeddings are treated similar to positional embeddings + instruction_state_embeddings2 = instruction_state_embeddings + time_embeddings + rgb_state_embeddings2 = rgb_state_embeddings + time_embeddings + depth_state_embeddings2 = depth_state_embeddings + time_embeddings + action_embeddings2 = action_embeddings + time_embeddings + returns_embeddings2 = returns_embeddings + time_embeddings + + return returns_embeddings2, instruction_state_embeddings2, rgb_state_embeddings2, depth_state_embeddings2, action_embeddings2 + + +class FullDecisionTransformerNet(AbstractDecisionTransformerNet): + def __init__( + self, observation_space: Space, model_config: Config, num_actions: int + ): + model_config.defrost() + # We do use Transformer encoding, but it is done to force to flatten the entry for instructions. + model_config.DECISION_TRANSFORMER.use_transformer_encoded_instruction = False + model_config.freeze() + super().__init__(observation_space, model_config, num_actions) + + def prepare_transformer_layer(self, model_config): + return nn.Transformer(d_model=model_config.DECISION_TRANSFORMER.hidden_dim + , nhead=model_config.DECISION_TRANSFORMER.n_head + , num_encoder_layers=model_config.DECISION_TRANSFORMER.ENCODER.n_layer + , num_decoder_layers=model_config.DECISION_TRANSFORMER.n_layer + , dim_feedforward=model_config.DECISION_TRANSFORMER.hidden_dim * 2 + , activation="gelu" + , batch_first=True) + + def handle_instruction_embeddings(self, observations, resize_tensor, batch_size, seq_length): + instruction_embedding = self.instruction_encoder(observations).permute(0, 2, 1) + instruction_embedding = resize_tensor(instruction_embedding) + return instruction_embedding + + def initialize_instruction_encoder(self): + self.instruction_encoder = Word2VecEmbeddings( + self.model_config.INSTRUCTION_ENCODER + ) + + def initialize_other_layers(self): + self.positional_encoding_for_instruction = PositionalEncoding(self.model_config.DECISION_TRANSFORMER.hidden_dim) + self.instruction_embed_state = nn.Linear(self.instruction_encoder.output_size, + self.model_config.DECISION_TRANSFORMER.hidden_dim) + self.rgb_embed_state = nn.Linear(self.model_config.RGB_ENCODER.output_size, + self.model_config.DECISION_TRANSFORMER.hidden_dim) + self.depth_embed_state = nn.Linear(self.model_config.DEPTH_ENCODER.output_size, + self.model_config.DECISION_TRANSFORMER.hidden_dim) + + self.encoder_instruction_to_rgb = self.prepare_transformer_layer(self.model_config) + self.encoder_instruction_to_depth = self.prepare_transformer_layer(self.model_config) + self.encoder_rgb_to_instruction = self.prepare_transformer_layer(self.model_config) + self.encoder_depth_to_instruction = self.prepare_transformer_layer(self.model_config) + self.visual_to_sentence_embed = nn.AdaptiveAvgPool1d(1) + + def initialize_transformer_step_size(self): + step_size = 2 + c = self.model_config.DECISION_TRANSFORMER.ENCODER + if c.use_rgb_state_embeddings is True: + step_size += 1 + if c.use_depth_state_embeddings is True: + step_size += 1 + if c.use_output_rgb_instructions is True: + step_size += 1 + if c.use_output_depth_instructions is True: + step_size += 1 + if c.use_output_rgb is True: + step_size += 1 + if c.use_output_depth is True: + step_size += 1 + return step_size + + def create_tensors_for_gpt_as_tuple(self, prev_actions, returns_to_go, instruction_embedding, + depth_embedding, rgb_embedding, timesteps, batch_size, + seq_length): + single_instruction_states = instruction_embedding[:, 0, :, :] + # embed each modality with a different head + instruction_state_embeddings = self.positional_encoding_for_instruction( + self.instruction_embed_state(single_instruction_states.permute(0, 2, 1))) + + # only 2D allowed in Pytorch Implementation + vision_causal_mask = VanillaMultiHeadAttention.create_causal_mask(seq_length, rgb_embedding.device)[0][0] + text_mask = VanillaMultiHeadAttention.create_padded_mask(instruction_state_embeddings) + + rgb_state_embeddings = self.rgb_activation(self.embed_ln(self.rgb_embed_state(rgb_embedding))) + depth_state_embeddings = self.depth_activation(self.embed_ln(self.depth_embed_state(depth_embedding))) + + action_embeddings = self.action_activation(self.embed_action(prev_actions)) + returns_embeddings = self.embed_return(returns_to_go) + time_embeddings = self.embed_timestep(timesteps) + + # print(state_embeddings.shape, action_embeddings.shape, returns_embeddings.shape, time_embeddings.shape) + # time embeddings are treated similar to positional embeddings + rgb_state_embeddings2 = self.embed_ln(rgb_state_embeddings) + time_embeddings + depth_state_embeddings2 = self.embed_ln(depth_state_embeddings) + time_embeddings + action_embeddings2 = action_embeddings + time_embeddings + returns_embeddings2 = returns_embeddings + time_embeddings + + causal_text_mask = VanillaMultiHeadAttention.create_causal_mask(instruction_state_embeddings.shape[1], + instruction_state_embeddings.device)[0][0] + rgb_mask = VanillaMultiHeadAttention.create_padded_mask(rgb_state_embeddings2) + depth_mask = VanillaMultiHeadAttention.create_padded_mask(depth_state_embeddings2) + + output_rgb_instructions = self.encoder_instruction_to_rgb(src=instruction_state_embeddings, + tgt=rgb_state_embeddings2, + src_key_padding_mask=text_mask, + tgt_mask=vision_causal_mask) + + output_depth_instructions = self.encoder_instruction_to_depth(src=instruction_state_embeddings, + tgt=depth_state_embeddings2, + src_key_padding_mask=text_mask, + tgt_mask=vision_causal_mask) + + output_rgb = self.encoder_rgb_to_instruction(src=rgb_state_embeddings2, tgt=instruction_state_embeddings, + src_key_padding_mask=rgb_mask, tgt_mask=causal_text_mask) + output_depth = self.encoder_depth_to_instruction(src=depth_state_embeddings2, tgt=instruction_state_embeddings, + src_key_padding_mask=depth_mask, tgt_mask=causal_text_mask) + + output_rgb = self.instruction_activation( + self.visual_to_sentence_embed(output_rgb.permute(0, 2, 1)).permute(0, 2, 1)) + time_embeddings + output_depth = self.instruction_activation( + self.visual_to_sentence_embed(output_depth.permute(0, 2, 1)).permute(0, 2, 1)) + time_embeddings + + c = self.model_config.DECISION_TRANSFORMER.ENCODER + t = [returns_embeddings2] + if c.use_rgb_state_embeddings is True: + t.append(rgb_state_embeddings2) + if c.use_depth_state_embeddings is True: + t.append(depth_state_embeddings2) + if c.use_output_rgb_instructions is True: + t.append(output_rgb_instructions) + if c.use_output_depth_instructions is True: + t.append(output_depth_instructions) + if c.use_output_rgb is True: + t.append(output_rgb) + if c.use_output_depth is True: + t.append(output_depth) + t.append(action_embeddings2) + assert len(t) >= 3 + + return tuple(t) + + +class FullDecisionTransformerSingleVisionStateNet(FullDecisionTransformerNet): + def __init__( + self, observation_space: Space, model_config: Config, num_actions: int + ): + super().__init__(observation_space, model_config, num_actions) + + def initialize_other_layers(self): + self.positional_encoding_for_instruction = PositionalEncoding(self.model_config.DECISION_TRANSFORMER.hidden_dim) + self.instruction_embed_state = nn.Linear(self.instruction_encoder.output_size, + self.model_config.DECISION_TRANSFORMER.hidden_dim) + self.rgb_embed_state = nn.Linear(self.model_config.RGB_ENCODER.output_size, + self.model_config.DECISION_TRANSFORMER.hidden_dim) + self.depth_embed_state = nn.Linear(self.model_config.DEPTH_ENCODER.output_size, + self.model_config.DECISION_TRANSFORMER.hidden_dim) + + self.embed_state = nn.Linear(self.rgb_embed_state.out_features + self.depth_embed_state.out_features, + self.model_config.DECISION_TRANSFORMER.hidden_dim) + + self.encoder_instruction_to_state = self.prepare_transformer_layer(self.model_config) + self.encoder_state_to_instruction = self.prepare_transformer_layer(self.model_config) + self.visual_to_sentence_embed = nn.AdaptiveAvgPool1d(1) + + def initialize_transformer_step_size(self): + step_size = 3 + c = self.model_config.DECISION_TRANSFORMER.ENCODER + if c.use_output_state_instructions is True: # like a different representation of instructions at each time steps + step_size += 1 + if c.use_output_state is True: # A single state representation of the whole sequence... + step_size += 1 + return step_size + + def create_tensors_for_gpt_as_tuple(self, prev_actions, returns_to_go, instruction_embedding, + depth_embedding, rgb_embedding, timesteps, batch_size, + seq_length): + single_instruction_states = instruction_embedding[:, 0, :, :] + # embed each modality with a different head + instruction_state_embeddings = self.positional_encoding_for_instruction( + self.instruction_embed_state(single_instruction_states.permute(0, 2, 1))) + + # only 2D allowed in Pytorch Implementation + vision_causal_mask = VanillaMultiHeadAttention.create_causal_mask(seq_length, rgb_embedding.device)[0][0] + text_mask = VanillaMultiHeadAttention.create_padded_mask(instruction_state_embeddings) + + rgb_state_embeddings = self.rgb_activation(self.embed_ln(self.rgb_embed_state(rgb_embedding))) + depth_state_embeddings = self.depth_activation(self.embed_ln(self.depth_embed_state(depth_embedding))) + + states = torch.cat( + [rgb_state_embeddings, depth_state_embeddings], dim=2 + ) + # embed each modality with a different head + state_embeddings = self.embed_ln(self.embed_state(states)) + + action_embeddings = self.action_activation(self.embed_action(prev_actions)) + returns_embeddings = self.embed_return(returns_to_go) + time_embeddings = self.embed_timestep(timesteps) + + # print(state_embeddings.shape, action_embeddings.shape, returns_embeddings.shape, time_embeddings.shape) + # time embeddings are treated similar to positional embeddings + state_embeddings2 = state_embeddings + time_embeddings + action_embeddings2 = action_embeddings + time_embeddings + returns_embeddings2 = returns_embeddings + time_embeddings + + causal_text_mask = VanillaMultiHeadAttention.create_causal_mask(instruction_state_embeddings.shape[1], + instruction_state_embeddings.device)[0][0] + state_mask = VanillaMultiHeadAttention.create_padded_mask(state_embeddings2) + + output_state_instructions = self.encoder_instruction_to_state(src=instruction_state_embeddings, + tgt=state_embeddings2, + src_key_padding_mask=text_mask, + tgt_mask=vision_causal_mask) + + output_state = self.encoder_state_to_instruction(src=state_embeddings2, tgt=instruction_state_embeddings, + src_key_padding_mask=state_mask, tgt_mask=causal_text_mask) + + output_state = self.instruction_activation( + self.visual_to_sentence_embed(output_state.permute(0, 2, 1)).permute(0, 2, 1)) + time_embeddings + + c = self.model_config.DECISION_TRANSFORMER.ENCODER + t = [returns_embeddings2, state_embeddings2] + if c.use_output_state_instructions is True: # like a different representation of instructions at each time steps + t.append(output_state_instructions) + if c.use_output_state is True: # A single state representation of the whole sequence... + t.append(output_state) + t.append(action_embeddings2) + + return tuple(t) diff --git a/vlnce_baselines/models/encoders/instruction_encoder.py b/vlnce_baselines/models/encoders/instruction_encoder.py index 4e28dbdd..a79d7419 100644 --- a/vlnce_baselines/models/encoders/instruction_encoder.py +++ b/vlnce_baselines/models/encoders/instruction_encoder.py @@ -6,8 +6,105 @@ from habitat import Config from habitat.core.simulator import Observations from torch import Tensor +from vlnce_baselines.models.utils import VanillaMultiHeadAttention +from vlnce_baselines.models.utils import PositionalEncoding +class Word2VecEmbeddings(nn.Module): + def __init__(self, config: Config) -> None: + """ + A layer that allows to us the pretrained Word2Vec embeddings + without further handling. + :param config: + """ + super().__init__() + + self.config = config + self.padding_idx = 0 + + if config.sensor_uuid == "instruction": + if self.config.use_pretrained_embeddings: + self.embedding_layer = nn.Embedding.from_pretrained( + embeddings=self._load_embeddings(), + freeze=not self.config.fine_tune_embeddings, + ) + else: # each embedding initialized to sampled Gaussian + self.embedding_layer = nn.Embedding( + num_embeddings=config.vocab_size, + embedding_dim=config.embedding_size, + padding_idx=self.padding_idx, + ) + + @property + def output_size(self): + return self.config.embedding_size + + def _load_embeddings(self) -> Tensor: + """Loads word embeddings from a pretrained embeddings file. + PAD: index 0. [0.0, ... 0.0] + UNK: index 1. mean of all R2R word embeddings: [mean_0, ..., mean_n] + why UNK is averaged: https://bit.ly/3u3hkYg + Returns: + embeddings tensor of size [num_words x embedding_dim] + """ + with gzip.open(self.config.embedding_file, "rt") as f: + embeddings = torch.tensor(json.load(f)) + return embeddings + + def forward(self, observations: Observations) -> Tensor: + """ + Tensor sizes after computation: + instruction: [batch_size x seq_length] + lengths: [batch_size] + hidden_state: [batch_size x hidden_size] + """ + assert self.config.sensor_uuid == "instruction" + instruction = observations["instruction"].long() + instruction = self.embedding_layer(instruction) + + return instruction + +class InstructionEncoderWithTransformer(nn.Module): + def __init__(self, config: Config) -> None: + """An encoder that uses RNN to encode an instruction. Returns + the final hidden state after processing the instruction sequence. + + Args: + config: must have + embedding_size: The dimension of each embedding vector + hidden_size: The hidden (output) size + rnn_type: The RNN cell type. Must be GRU or LSTM + final_state_only: If True, return just the final state + """ + super().__init__() + + self.config = config + self.word2vec = Word2VecEmbeddings(config.INSTRUCTION_ENCODER) + encoder_layer = nn.TransformerEncoderLayer(d_model=config.DECISION_TRANSFORMER.hidden_dim, nhead=config.DECISION_TRANSFORMER.n_head, batch_first=True) + self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=config.DECISION_TRANSFORMER.n_layer) + self.instruction_embed_state = nn.Linear(self.word2vec.output_size, + config.DECISION_TRANSFORMER.hidden_dim) + self.positional_encoding = PositionalEncoding(config.DECISION_TRANSFORMER.hidden_dim) + + def forward(self, observations: Observations) -> Tensor: + """ + Tensor sizes after computation: + instruction: [batch_size x seq_length] + lengths: [batch_size] + hidden_state: [batch_size x hidden_size] + """ + pretrained_embeddings = self.word2vec(observations) + # Instructions are repeated at each time step, we want only one instructions + embeddings = self.instruction_embed_state(pretrained_embeddings)[:,0,:,:] + embeddings = self.positional_encoding(embeddings) + padded_mask = VanillaMultiHeadAttention.create_padded_mask(embeddings) + encoded_instructions = self.transformer_encoder(embeddings, src_key_padding_mask=padded_mask) + return encoded_instructions + + @property + def output_size(self): + return self.config.DECISION_TRANSFORMER.hidden_dim + class InstructionEncoder(nn.Module): def __init__(self, config: Config) -> None: """An encoder that uses RNN to encode an instruction. Returns diff --git a/vlnce_baselines/models/encoders/min_gpt.py b/vlnce_baselines/models/encoders/min_gpt.py new file mode 100644 index 00000000..c228e9c2 --- /dev/null +++ b/vlnce_baselines/models/encoders/min_gpt.py @@ -0,0 +1,296 @@ +# Origin: https://github.com/karpathy/minGPT/blob/master/mingpt/model.py +# https://github.com/karpathy/minGPT/commit/90420ee978fed95e6eb7c9add728d33bb890fe39 + + +import math + +import torch +import torch.nn as nn +from torch.nn import functional as F + +import yacs.config + +# Default Habitat config node +class Config(yacs.config.CfgNode): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs, new_allowed=True) + + +CN = Config + +# ----------------------------------------------------------------------------- + +class NewGELU(nn.Module): + """ + Implementation of the GELU activation function currently in Google BERT repo (identical to OpenAI GPT). + Reference: Gaussian Error Linear Units (GELU) paper: https://arxiv.org/abs/1606.08415 + """ + def forward(self, x): + return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) + +class CausalSelfAttention(nn.Module): + """ + A vanilla multi-head masked self-attention layer with a projection at the end. + It is possible to use torch.nn.MultiheadAttention here but I am including an + explicit implementation here to show that there is nothing too scary here. + """ + + def __init__(self, config): + super().__init__() + assert config.n_embd % config.n_head == 0 + # key, query, value projections for all heads, but in a batch + self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd) + # output projection + self.c_proj = nn.Linear(config.n_embd, config.n_embd) + # regularization + self.attn_dropout = nn.Dropout(config.attn_pdrop) + self.resid_dropout = nn.Dropout(config.resid_pdrop) + # causal mask to ensure that attention is only applied to the left in the input sequence + self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size)) + .view(1, 1, config.block_size, config.block_size)) + self.n_head = config.n_head + self.n_embd = config.n_embd + + def forward(self, x): + B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd) + + # calculate query, key, values for all heads in batch and move head forward to be the batch dim + q, k ,v = self.c_attn(x).split(self.n_embd, dim=2) + k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) + q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) + v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) + + # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T) + att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) + att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf')) + att = F.softmax(att, dim=-1) + att = self.attn_dropout(att) + y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs) + y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side + + # output projection + y = self.resid_dropout(self.c_proj(y)) + return y + +class Block(nn.Module): + """ an unassuming Transformer block """ + + def __init__(self, config): + super().__init__() + self.ln_1 = nn.LayerNorm(config.n_embd) + self.attn = CausalSelfAttention(config) + self.ln_2 = nn.LayerNorm(config.n_embd) + self.mlp = nn.ModuleDict(dict( + c_fc = nn.Linear(config.n_embd, 4 * config.n_embd), + c_proj = nn.Linear(4 * config.n_embd, config.n_embd), + act = NewGELU(), + dropout = nn.Dropout(config.resid_pdrop), + )) + m = self.mlp + self.mlpf = lambda x: m.dropout(m.c_proj(m.act(m.c_fc(x)))) # MLP forward + # Add ReZero here + self.use_re_zero = config.use_re_zero + if self.use_re_zero: + self.re_zero_weights_1 = nn.Linear(config.n_embd, config.n_embd) + torch.nn.init.normal_(self.re_zero_weights_1.weight, 0, 0.25/config.n_embd) + self.re_zero_weights_2 = nn.Linear(config.n_embd, config.n_embd) + torch.nn.init.normal_(self.re_zero_weights_2.weight, 0, 0.25/config.n_embd) + + def forward(self, x): + if self.use_re_zero: + x = x + self.re_zero_weights_1(self.attn(x)) + x = x + self.re_zero_weights_2(self.mlpf(x)) + else: + x = x + self.attn(self.ln_1(x)) + x = x + self.mlpf(self.ln_2(x)) + return x + +class GPT(nn.Module): + """ GPT Language Model """ + + @staticmethod + def get_default_config(): + C = CN() + # either model_type or (n_layer, n_head, n_embd) must be given in the config + C.model_type = None + C.n_layer = 2 + C.n_head = 1 + C.n_embd = 128 + # these options must be filled in externally + C.vocab_size = 4 + C.block_size = 156 + C.episode_horizon = 128 + C.step_size = 1 + # dropout hyperparameters + C.embd_pdrop = 0.1 + C.resid_pdrop = 0.1 + C.attn_pdrop = 0.1 + return C + + def __init__(self, config): + super().__init__() + assert config.vocab_size is not None + assert config.episode_horizon is not None + assert config.step_size is not None + self.block_size = config.episode_horizon * config.step_size + config.defrost() + config.block_size = self.block_size + delete_n_embd = False + if not hasattr(config, "n_embd"): + delete_n_embd = True + + if hasattr(config, "hidden_dim"): + config.n_embd = config.hidden_dim + type_given = config.model_type is not None and len(config.model_type) > 0 + params_given = all([config.n_layer is not None, config.n_head is not None, config.n_embd is not None]) + assert type_given ^ params_given # exactly one of these (XOR) + if type_given: + # translate from model_type to detailed configuration + config.merge_from_dict({ + # names follow the huggingface naming conventions + # GPT-1 + 'openai-gpt': dict(n_layer=12, n_head=12, n_embd=768), # 117M params + # GPT-2 configs + 'gpt2': dict(n_layer=12, n_head=12, n_embd=768), # 124M params + 'gpt2-medium': dict(n_layer=24, n_head=16, n_embd=1024), # 350M params + 'gpt2-large': dict(n_layer=36, n_head=20, n_embd=1280), # 774M params + 'gpt2-xl': dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params + # Gophers + 'gopher-44m': dict(n_layer=8, n_head=16, n_embd=512), + # (there are a number more...) + # I made these tiny models up + 'gpt-mini': dict(n_layer=6, n_head=6, n_embd=192), + 'gpt-micro': dict(n_layer=4, n_head=4, n_embd=128), + 'gpt-nano': dict(n_layer=3, n_head=3, n_embd=48), + }[config.model_type]) + + self.transformer = nn.ModuleDict(dict( + # wte = nn.Embedding(config.vocab_size, config.n_embd), + # wpe = nn.Embedding(config.block_size, config.n_embd), + drop = nn.Dropout(config.embd_pdrop), + h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), + ln_f = nn.LayerNorm(config.n_embd), + )) + self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) + + # init all weights, and apply a special scaled init to the residual projections, per GPT-2 paper + self.apply(self._init_weights) + for pn, p in self.named_parameters(): + if pn.endswith('c_proj.weight'): + torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer)) + + # report number of parameters (note we don't count the decoder parameters in lm_head) + n_params = sum(p.numel() for p in self.transformer.parameters()) + print("number of parameters for GPT: %.2fM" % (n_params/1e6,)) + if delete_n_embd: + del config["n_embd"] + config.freeze() + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + if module.bias is not None: + torch.nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + elif isinstance(module, nn.LayerNorm): + torch.nn.init.zeros_(module.bias) + torch.nn.init.ones_(module.weight) + + + def configure_optimizers(self, train_config): + """ + This long function is unfortunately doing something very simple and is being very defensive: + We are separating out all parameters of the model into two buckets: those that will experience + weight decay for regularization and those that won't (biases, and layernorm/embedding weights). + We are then returning the PyTorch optimizer object. + """ + + # separate out all parameters to those that will and won't experience regularizing weight decay + decay = set() + no_decay = set() + whitelist_weight_modules = (torch.nn.Linear, ) + blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding) + for mn, m in self.named_modules(): + for pn, p in m.named_parameters(): + fpn = '%s.%s' % (mn, pn) if mn else pn # full param name + # random note: because named_modules and named_parameters are recursive + # we will see the same tensors p many many times. but doing it this way + # allows us to know which parent module any tensor p belongs to... + if pn.endswith('bias'): + # all biases will not be decayed + no_decay.add(fpn) + elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules): + # weights of whitelist modules will be weight decayed + decay.add(fpn) + elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules): + # weights of blacklist modules will NOT be weight decayed + no_decay.add(fpn) + + # validate that we considered every parameter + param_dict = {pn: p for pn, p in self.named_parameters()} + inter_params = decay & no_decay + union_params = decay | no_decay + assert len(inter_params) == 0, "parameters %s made it into both decay/no_decay sets!" % (str(inter_params), ) + assert len(param_dict.keys() - union_params) == 0, "parameters %s were not separated into either decay/no_decay set!" \ + % (str(param_dict.keys() - union_params), ) + + # create the pytorch optimizer object + optim_groups = [ + {"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": train_config.weight_decay}, + {"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0}, + ] + optimizer = torch.optim.AdamW(optim_groups, lr=train_config.learning_rate, betas=train_config.betas) + return optimizer + + def forward(self, x, targets=None): + # device = idx.device + # b, t = idx.size() + # assert t <= self.block_size, f"Cannot forward sequence of length {t}, block size is only {self.block_size}" + # pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0) # shape (1, t) + # + # # forward the GPT model itself + # tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd) + # pos_emb = self.transformer.wpe(pos) # position embeddings of shape (1, t, n_embd) + x = self.transformer.drop(x) + for block in self.transformer.h: + x = block(x) + x = self.transformer.ln_f(x) + # logits = self.lm_head(x) + # + # # if we are given some desired targets also calculate the loss + # loss = None + # if targets is not None: + # loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1) + + return x + + @torch.no_grad() + def generate(self, idx, max_new_tokens, temperature=1.0, do_sample=False, top_k=None): + """ + Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete + the sequence max_new_tokens times, feeding the predictions back into the model each time. + Most likely you'll want to make sure to be in model.eval() mode of operation for this. + """ + for _ in range(max_new_tokens): + # if the sequence context is growing too long we must crop it at block_size + idx_cond = idx if idx.size(1) <= self.block_size else idx[:, -self.block_size:] + # forward the model to get the logits for the index in the sequence + logits, _ = self(idx_cond) + # pluck the logits at the final step and scale by desired temperature + logits = logits[:, -1, :] / temperature + # optionally crop the logits to only the top k options + if top_k is not None: + v, _ = torch.topk(logits, top_k) + logits[logits < v[:, [-1]]] = -float('Inf') + # apply softmax to convert logits to (normalized) probabilities + probs = F.softmax(logits, dim=-1) + # either sample from the distribution or take the most likely element + if do_sample: + idx_next = torch.multinomial(probs, num_samples=1) + else: + _, idx_next = torch.topk(probs, k=1, dim=-1) + # append sampled index to the running sequence and continue + idx = torch.cat((idx, idx_next), dim=1) + + return diff --git a/vlnce_baselines/models/utils.py b/vlnce_baselines/models/utils.py index dd29b3e0..60388d3f 100644 --- a/vlnce_baselines/models/utils.py +++ b/vlnce_baselines/models/utils.py @@ -1,12 +1,14 @@ +import math from numbers import Number from typing import Any, Optional, Union import numpy as np import torch import torch.nn as nn -from torch import Size, Tensor +from torch import Size, Tensor, nn as nn from torch.distributions import constraints from torch.distributions.normal import Normal +from torch.nn import functional as F class TemperatureTanh(nn.Module): @@ -315,3 +317,96 @@ def batched_index_select( expanse[dim] = -1 index = index.view(views).expand(expanse) return torch.gather(x, dim, index).squeeze(dim) + + +class PositionalEncoding(nn.Module): + + def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 200): + super().__init__() + self.dropout = nn.Dropout(p=dropout) + + position = torch.arange(max_len).unsqueeze(1) + div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) + pe = torch.zeros(max_len, 1, d_model) + pe[:, 0, 0::2] = torch.sin(position * div_term) + pe[:, 0, 1::2] = torch.cos(position * div_term) + self.register_buffer('pe', pe.permute(1,0,2)) + + def forward(self, x: Tensor) -> Tensor: + """ + Args: + x: Tensor, shape [batch_size, seq_length, embedding_dim] + """ + x = x + self.pe[:x.size(1)] + return self.dropout(x) + + +class VanillaMultiHeadAttention(nn.Module): + """ + A vanilla multi-head masked attention layer with a projection at the end. + Depending on the mask and entries provided, it can serve as + Causal Multihead Self Attention. + """ + + def __init__(self, config): + super().__init__() + assert config.n_embd % config.n_head == 0 + # key, query, value projections for all heads, but in a batch + self.k_attn = nn.Linear(config.n_embd, config.n_embd) + self.q_attn = nn.Linear(config.n_embd, config.n_embd) + self.v_attn = nn.Linear(config.n_embd, config.n_embd) + # output projection + self.c_proj = nn.Linear(config.n_embd, config.n_embd) + # regularization + self.attn_dropout = nn.Dropout(config.attn_pdrop) + self.resid_dropout = nn.Dropout(config.resid_pdrop) + + self.n_head = config.n_head + self.n_embd = config.n_embd + + @staticmethod + def create_causal_mask(size: int, device="cpu"): + return (torch.tril(torch.ones(size, size)).view(1, 1, size, size) == 0).to(device) + + @staticmethod + def create_padded_mask(t: Tensor): + return (t == 0.0).all(dim=-1).to(t.device) + + def forward(self, q, k, v, mask=None): + q_B, q_T, q_C = q.size() # batch size, sequence length, embedding dimensionality (n_embd) + k_B, k_T, k_C = k.size() + v_B, v_T, v_C = v.size() + + # all entries must have the same batch size. Keys and values must have the same sequence length + assert q_B == k_B + assert k_B == v_B + assert k_C == v_C + + q = self.k_attn(q) + k = self.k_attn(k) + v = self.k_attn(v) + + # calculate query, key, values for all heads in batch and move head forward to be the batch dim + q = q.view(q_B, q_T, self.n_head, q_C // self.n_head).transpose(1, 2) # (B, nh, T, hs) + k = k.view(q_B, k_T, self.n_head, k_C // self.n_head).transpose(1, 2) # (B, nh, T, hs) + v = v.view(q_B, v_T, self.n_head, v_C // self.n_head).transpose(1, 2) # (B, nh, T, hs) + + # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T) + att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) + if mask is not None: + num_dim_att = len(att.shape) + num_dim_mask = len(mask.shape) + dim_to_add = num_dim_att - num_dim_mask + assert dim_to_add >= 0 + if dim_to_add > 0: + for i in range(1, dim_to_add+1): + mask = mask.unsqueeze(i) + att = att.masked_fill(mask, float('-inf')) + att = F.softmax(att, dim=-1) + att = self.attn_dropout(att) + y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs) + y = y.transpose(1, 2).contiguous().view(q_B, q_T, q_C) # re-assemble all head outputs side by side + + # output projection + y = self.resid_dropout(self.c_proj(y)) + return y