diff --git a/aloha-devel/requirements.txt b/aloha-devel/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..99b4b68bedb70af755d6e3c38f704e67bd961f8b --- /dev/null +++ b/aloha-devel/requirements.txt @@ -0,0 +1,23 @@ +# python=3.9 +# pip==23.0.1 +# pytorch==2.0.0 +# torchvision==0.15.0 +# pytorch-cuda==11.8 +pyquaternion==0.9.9 +pyyaml==6.0 +pexpect==4.8.0 +mujoco==2.3.7 +dm_control==1.0.14 +matplotlib==3.7.5 +einops==0.7.0 +packaging==23.0 +h5py==3.8.0 +ipython==8.12.3 +opencv-python==4.9.0.80 +rospkg==1.5.0 +empy==3.3.4 +catkin-pkg +diffusers==0.26.3 +termcolor==2.4.0 +imageio==2.34.0 +# cd act/detr && pip install -v -e . \ No newline at end of file diff --git a/aloha-devel/robomimic/__init__.py b/aloha-devel/robomimic/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b4d0fc53092a1d824d2d6493c4f97664024949d --- /dev/null +++ b/aloha-devel/robomimic/__init__.py @@ -0,0 +1,159 @@ +__version__ = "0.3.0" + + +# stores released dataset links and rollout horizons in global dictionary. +# Structure is given below for each type of dataset: + +# robosuite / real +# { +# task: +# dataset_type: +# hdf5_type: +# url: link +# horizon: value +# ... +# ... +# ... +# } +DATASET_REGISTRY = {} + +# momart +# { +# task: +# dataset_type: +# url: link +# size: value +# ... +# ... +# } +MOMART_DATASET_REGISTRY = {} + + +def register_dataset_link(task, dataset_type, hdf5_type, link, horizon): + """ + Helper function to register dataset link in global dictionary. + Also takes a @horizon parameter - this corresponds to the evaluation + rollout horizon that should be used during training. + + Args: + task (str): name of task for this dataset + dataset_type (str): type of dataset (usually identifies the dataset source) + hdf5_type (str): type of hdf5 - usually one of "raw", "low_dim", or "image", + to identify the kind of observations in the dataset + link (str): download link for the dataset + horizon (int): evaluation rollout horizon that should be used with this dataset + """ + if task not in DATASET_REGISTRY: + DATASET_REGISTRY[task] = {} + if dataset_type not in DATASET_REGISTRY[task]: + DATASET_REGISTRY[task][dataset_type] = {} + DATASET_REGISTRY[task][dataset_type][hdf5_type] = dict(url=link, horizon=horizon) + + +def register_all_links(): + """ + Record all dataset links in this function. + """ + + # all proficient human datasets + ph_tasks = ["lift", "can", "square", "transport", "tool_hang", "lift_real", "can_real", "tool_hang_real"] + ph_horizons = [400, 400, 400, 700, 700, 1000, 1000, 1000] + for task, horizon in zip(ph_tasks, ph_horizons): + register_dataset_link(task=task, dataset_type="ph", hdf5_type="raw", horizon=horizon, + link="http://downloads.cs.stanford.edu/downloads/rt_benchmark/{}/ph/demo{}.hdf5".format( + task, "" if "real" in task else "_v141" + ) + ) + # real world datasets only have demo.hdf5 files which already contain all observation modalities + # while sim datasets store raw low-dim mujoco states in the demo.hdf5 + if "real" not in task: + register_dataset_link(task=task, dataset_type="ph", hdf5_type="low_dim", horizon=horizon, + link="http://downloads.cs.stanford.edu/downloads/rt_benchmark/{}/ph/low_dim_v141.hdf5".format(task)) + register_dataset_link(task=task, dataset_type="ph", hdf5_type="image", horizon=horizon, + link=None) + + # all multi human datasets + mh_tasks = ["lift", "can", "square", "transport"] + mh_horizons = [500, 500, 500, 1100] + for task, horizon in zip(mh_tasks, mh_horizons): + register_dataset_link(task=task, dataset_type="mh", hdf5_type="raw", horizon=horizon, + link="http://downloads.cs.stanford.edu/downloads/rt_benchmark/{}/mh/demo_v141.hdf5".format(task)) + register_dataset_link(task=task, dataset_type="mh", hdf5_type="low_dim", horizon=horizon, + link="http://downloads.cs.stanford.edu/downloads/rt_benchmark/{}/mh/low_dim_v141.hdf5".format(task)) + register_dataset_link(task=task, dataset_type="mh", hdf5_type="image", horizon=horizon, + link=None) + + # all machine generated datasets + for task, horizon in zip(["lift", "can"], [400, 400]): + register_dataset_link(task=task, dataset_type="mg", hdf5_type="raw", horizon=horizon, + link="http://downloads.cs.stanford.edu/downloads/rt_benchmark/{}/mg/demo_v141.hdf5".format(task)) + register_dataset_link(task=task, dataset_type="mg", hdf5_type="low_dim_sparse", horizon=horizon, + link="http://downloads.cs.stanford.edu/downloads/rt_benchmark/{}/mg/low_dim_sparse_v141.hdf5".format(task)) + register_dataset_link(task=task, dataset_type="mg", hdf5_type="image_sparse", horizon=horizon, + link="http://downloads.cs.stanford.edu/downloads/rt_benchmark/{}/mg/image_sparse_v141.hdf5".format(task)) + register_dataset_link(task=task, dataset_type="mg", hdf5_type="low_dim_dense", horizon=horizon, + link="http://downloads.cs.stanford.edu/downloads/rt_benchmark/{}/mg/low_dim_dense_v141.hdf5".format(task)) + register_dataset_link(task=task, dataset_type="mg", hdf5_type="image_dense", horizon=horizon, + link="http://downloads.cs.stanford.edu/downloads/rt_benchmark/{}/mg/image_dense_v141.hdf5".format(task)) + + # can-paired dataset + register_dataset_link(task="can", dataset_type="paired", hdf5_type="raw", horizon=400, + link="http://downloads.cs.stanford.edu/downloads/rt_benchmark/can/paired/demo_v141.hdf5") + register_dataset_link(task="can", dataset_type="paired", hdf5_type="low_dim", horizon=400, + link="http://downloads.cs.stanford.edu/downloads/rt_benchmark/can/paired/low_dim_v141.hdf5") + register_dataset_link(task="can", dataset_type="paired", hdf5_type="image", horizon=400, + link=None) + + +def register_momart_dataset_link(task, dataset_type, link, dataset_size): + """ + Helper function to register dataset link in global dictionary. + Also takes a @horizon parameter - this corresponds to the evaluation + rollout horizon that should be used during training. + + Args: + task (str): name of task for this dataset + dataset_type (str): type of dataset (usually identifies the dataset source) + link (str): download link for the dataset + dataset_size (float): size of the dataset, in GB + """ + if task not in MOMART_DATASET_REGISTRY: + MOMART_DATASET_REGISTRY[task] = {} + if dataset_type not in MOMART_DATASET_REGISTRY[task]: + MOMART_DATASET_REGISTRY[task][dataset_type] = {} + MOMART_DATASET_REGISTRY[task][dataset_type] = dict(url=link, size=dataset_size) + + +def register_all_momart_links(): + """ + Record all dataset links in this function. + """ + # all tasks, mapped to their [exp, sub, gen, sam] sizes + momart_tasks = { + "table_setup_from_dishwasher": [14, 14, 3.3, 0.6], + "table_setup_from_dresser": [16, 17, 3.1, 0.7], + "table_cleanup_to_dishwasher": [23, 36, 5.3, 1.1], + "table_cleanup_to_sink": [17, 28, 2.9, 0.8], + "unload_dishwasher": [21, 27, 5.4, 1.0], + } + + momart_dataset_types = [ + "expert", + "suboptimal", + "generalize", + "sample", + ] + + # Iterate over all combos and register the link + for task, dataset_sizes in momart_tasks.items(): + for dataset_type, dataset_size in zip(momart_dataset_types, dataset_sizes): + register_momart_dataset_link( + task=task, + dataset_type=dataset_type, + link=f"http://downloads.cs.stanford.edu/downloads/rt_mm/{dataset_type}/{task}_{dataset_type}.hdf5", + dataset_size=dataset_size, + ) + + +register_all_links() +register_all_momart_links() diff --git a/aloha-devel/robomimic/__pycache__/__init__.cpython-38.pyc b/aloha-devel/robomimic/__pycache__/__init__.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e3f5b0fe1db2e1296d28c015b1c2161d8304118d Binary files /dev/null and b/aloha-devel/robomimic/__pycache__/__init__.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/algo/__pycache__/act.cpython-38.pyc b/aloha-devel/robomimic/algo/__pycache__/act.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..59b000aa4b83e8d99e2a1a271ac619a818b12354 Binary files /dev/null and b/aloha-devel/robomimic/algo/__pycache__/act.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/algo/__pycache__/gl.cpython-38.pyc b/aloha-devel/robomimic/algo/__pycache__/gl.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2de4df7b3349017435348b57444726dd4cf5a54a Binary files /dev/null and b/aloha-devel/robomimic/algo/__pycache__/gl.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/algo/__pycache__/iql.cpython-38.pyc b/aloha-devel/robomimic/algo/__pycache__/iql.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..de41d01bbb244c50ac6ac4d6cc4620d9ab103de9 Binary files /dev/null and b/aloha-devel/robomimic/algo/__pycache__/iql.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/algo/__pycache__/iris.cpython-38.pyc b/aloha-devel/robomimic/algo/__pycache__/iris.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d597c071932c226e7cfc3d96ade59b46d27e34ee Binary files /dev/null and b/aloha-devel/robomimic/algo/__pycache__/iris.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/config/__init__.py b/aloha-devel/robomimic/config/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b8b9da73ba059e7b0b77f9bd48cdca03fcb1f222 --- /dev/null +++ b/aloha-devel/robomimic/config/__init__.py @@ -0,0 +1,14 @@ +from robomimic.config.config import Config +from robomimic.config.base_config import config_factory, get_all_registered_configs + +# note: these imports are needed to register these classes in the global config registry +from robomimic.config.bc_config import BCConfig +from robomimic.config.bcq_config import BCQConfig +from robomimic.config.cql_config import CQLConfig +from robomimic.config.iql_config import IQLConfig +from robomimic.config.gl_config import GLConfig +from robomimic.config.hbc_config import HBCConfig +from robomimic.config.iris_config import IRISConfig +from robomimic.config.td3_bc_config import TD3_BCConfig +from robomimic.config.diffusion_policy_config import DiffusionPolicyConfig +from robomimic.config.act_config import ACTConfig diff --git a/aloha-devel/robomimic/config/__pycache__/act_config.cpython-38.pyc b/aloha-devel/robomimic/config/__pycache__/act_config.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..97945ebe49af15ed43862dcedccc918d26418200 Binary files /dev/null and b/aloha-devel/robomimic/config/__pycache__/act_config.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/config/__pycache__/base_config.cpython-38.pyc b/aloha-devel/robomimic/config/__pycache__/base_config.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e7cb9aa39fe54b961bdd95a12431d1ac1253ef94 Binary files /dev/null and b/aloha-devel/robomimic/config/__pycache__/base_config.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/config/__pycache__/config.cpython-38.pyc b/aloha-devel/robomimic/config/__pycache__/config.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ba2aeb7f1fcd6915436b48c287a3bfcdfe732eb4 Binary files /dev/null and b/aloha-devel/robomimic/config/__pycache__/config.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/config/__pycache__/cql_config.cpython-38.pyc b/aloha-devel/robomimic/config/__pycache__/cql_config.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1580b72a500c33e3151f3d58040f98c3be78273d Binary files /dev/null and b/aloha-devel/robomimic/config/__pycache__/cql_config.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/config/__pycache__/td3_bc_config.cpython-38.pyc b/aloha-devel/robomimic/config/__pycache__/td3_bc_config.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb75b2a42c375efa879b99eb687183a77c75ba7a Binary files /dev/null and b/aloha-devel/robomimic/config/__pycache__/td3_bc_config.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/config/bcq_config.py b/aloha-devel/robomimic/config/bcq_config.py new file mode 100644 index 0000000000000000000000000000000000000000..e28f5ba5668aa2e5e6d9ca2187953a4e05b56a7d --- /dev/null +++ b/aloha-devel/robomimic/config/bcq_config.py @@ -0,0 +1,83 @@ +""" +Config for BCQ algorithm. +""" + +from robomimic.config.base_config import BaseConfig +from robomimic.config.bc_config import BCConfig + + +class BCQConfig(BaseConfig): + ALGO_NAME = "bcq" + + def algo_config(self): + """ + This function populates the `config.algo` attribute of the config, and is given to the + `Algo` subclass (see `algo/algo.py`) for each algorithm through the `algo_config` + argument to the constructor. Any parameter that an algorithm needs to determine its + training and test-time behavior should be populated here. + """ + + # optimization parameters + self.algo.optim_params.critic.learning_rate.initial = 1e-3 # critic learning rate + self.algo.optim_params.critic.learning_rate.decay_factor = 0.1 # factor to decay LR by (if epoch schedule non-empty) + self.algo.optim_params.critic.learning_rate.epoch_schedule = [] # epochs where LR decay occurs + self.algo.optim_params.critic.regularization.L2 = 0.00 # L2 regularization strength + self.algo.optim_params.critic.start_epoch = -1 # number of epochs before starting critic training (-1 means start right away) + self.algo.optim_params.critic.end_epoch = -1 # number of epochs before ending critic training (-1 means start right away) + + self.algo.optim_params.action_sampler.learning_rate.initial = 1e-3 # action sampler learning rate + self.algo.optim_params.action_sampler.learning_rate.decay_factor = 0.1 # factor to decay LR by (if epoch schedule non-empty) + self.algo.optim_params.action_sampler.learning_rate.epoch_schedule = [] # epochs where LR decay occurs + self.algo.optim_params.action_sampler.regularization.L2 = 0.00 # L2 regularization strength + self.algo.optim_params.action_sampler.start_epoch = -1 # number of epochs before starting action sampler training (-1 means start right away) + self.algo.optim_params.action_sampler.end_epoch = -1 # number of epochs before ending action sampler training (-1 means start right away) + + self.algo.optim_params.actor.learning_rate.initial = 1e-3 # actor learning rate + self.algo.optim_params.actor.learning_rate.decay_factor = 0.1 # factor to decay LR by (if epoch schedule non-empty) + self.algo.optim_params.actor.learning_rate.epoch_schedule = [] # epochs where LR decay occurs + self.algo.optim_params.actor.regularization.L2 = 0.00 # L2 regularization strength + self.algo.optim_params.actor.start_epoch = -1 # number of epochs before starting actor training (-1 means start right away) + self.algo.optim_params.actor.end_epoch = -1 # number of epochs before ending actor training (-1 means start right away) + + # target network related parameters + self.algo.discount = 0.99 # discount factor to use + self.algo.n_step = 1 # for using n-step returns in TD-updates + self.algo.target_tau = 0.005 # update rate for target networks + self.algo.infinite_horizon = False # if True, scale terminal rewards by 1 / (1 - discount) to treat as infinite horizon + + # ================== Critic Network Config =================== + self.algo.critic.use_huber = False # Huber Loss instead of L2 for critic + self.algo.critic.max_gradient_norm = None # L2 gradient clipping for critic (None to use no clipping) + self.algo.critic.value_bounds = None # optional 2-tuple to ensure lower and upper bound on value estimates + self.algo.critic.num_action_samples = 10 # number of actions to sample per training batch to get target critic value + self.algo.critic.num_action_samples_rollout = 100 # number of actions to sample per environment step + + # critic ensemble parameters (TD3 trick) + self.algo.critic.ensemble.n = 2 # number of Q networks in the ensemble + self.algo.critic.ensemble.weight = 0.75 # weighting for mixing min and max for target Q value + + # distributional critic + self.algo.critic.distributional.enabled = False # train distributional critic (C51) + self.algo.critic.distributional.num_atoms = 51 # number of values in categorical distribution + + self.algo.critic.layer_dims = (300, 400) # size of critic MLP + + # ================== Action Sampler Config =================== + self.algo.action_sampler = BCConfig().algo + # use VAE by default + self.algo.action_sampler.vae.enabled = True + # remove unused parts of BCConfig algo config + del self.algo.action_sampler.optim_params # since action sampler optim params specified at top-level + del self.algo.action_sampler.loss + del self.algo.action_sampler.gaussian + del self.algo.action_sampler.rnn + del self.algo.action_sampler.transformer + + # Number of epochs before freezing encoder (-1 for no freezing). Only applies to cVAE-based action samplers. + with self.algo.action_sampler.unlocked(): + self.algo.action_sampler.freeze_encoder_epoch = -1 + + # ================== Actor Network Config =================== + self.algo.actor.enabled = False # whether to use the actor perturbation network + self.algo.actor.perturbation_scale = 0.05 # size of learned action perturbations + self.algo.actor.layer_dims = (300, 400) # size of actor MLP diff --git a/aloha-devel/robomimic/config/config.py b/aloha-devel/robomimic/config/config.py new file mode 100644 index 0000000000000000000000000000000000000000..74da6535b385f91aa5c34e20af731ba2e3d06ecb --- /dev/null +++ b/aloha-devel/robomimic/config/config.py @@ -0,0 +1,322 @@ +""" +Basic config class - provides a convenient way to work with nested +dictionaries (by exposing keys as attributes) and to save / load from jsons. + +Based on addict: https://github.com/mewwts/addict +""" + +import json +import copy +import contextlib +from copy import deepcopy + + +class Config(dict): + + def __init__(__self, *args, **kwargs): + object.__setattr__(__self, '__key_locked', False) # disallow adding new keys + object.__setattr__(__self, '__all_locked', False) # disallow both key and value update + object.__setattr__(__self, '__do_not_lock_keys', False) # cannot be key-locked + object.__setattr__(__self, '__parent', kwargs.pop('__parent', None)) + object.__setattr__(__self, '__key', kwargs.pop('__key', None)) + for arg in args: + if not arg: + continue + elif isinstance(arg, dict): + for key, val in arg.items(): + __self[key] = __self._hook(val) + elif isinstance(arg, tuple) and (not isinstance(arg[0], tuple)): + __self[arg[0]] = __self._hook(arg[1]) + else: + for key, val in iter(arg): + __self[key] = __self._hook(val) + + for key, val in kwargs.items(): + __self[key] = __self._hook(val) + + def lock(self): + """ + Lock the config. Afterwards, new keys cannot be added to the + config, and the values of existing keys cannot be modified. + """ + object.__setattr__(self, '__all_locked', True) + if self.key_lockable: + object.__setattr__(self, '__key_locked', True) + + for k in self: + if isinstance(self[k], Config): + self[k].lock() + + def unlock(self): + """ + Unlock the config. Afterwards, new keys can be added to the + config, and the values of existing keys can be modified. + """ + object.__setattr__(self, '__all_locked', False) + object.__setattr__(self, '__key_locked', False) + + for k in self: + if isinstance(self[k], Config): + self[k].unlock() + + def _get_lock_state_recursive(self): + """ + Internal helper function to get the lock state of all sub-configs recursively. + """ + lock_state = {"__all_locked": self.is_locked, "__key_locked": self.is_key_locked} + for k in self: + if isinstance(self[k], Config): + assert k not in ["__all_locked", "__key_locked"] + lock_state[k] = self[k]._get_lock_state_recursive() + return lock_state + + def _set_lock_state_recursive(self, lock_state): + """ + Internal helper function to set the lock state of all sub-configs recursively. + """ + lock_state = deepcopy(lock_state) + object.__setattr__(self, '__all_locked', lock_state.pop("__all_locked")) + object.__setattr__(self, '__key_locked', lock_state.pop("__key_locked")) + for k in lock_state: + if isinstance(self[k], Config): + self[k]._set_lock_state_recursive(lock_state[k]) + + def _get_lock_state(self): + """ + Retrieves the lock state of this config. + + Returns: + lock_state (dict): a dictionary with an "all_locked" key that is True + if both key and value updates are locked and False otherwise, and + a "key_locked" key that is True if only key updates are locked (value + updates still allowed) and False otherwise + """ + return { + "all_locked": self.is_locked, + "key_locked": self.is_key_locked + } + + def _set_lock_state(self, lock_state): + """ + Sets the lock state for this config. + + Args: + lock_state (dict): a dictionary with an "all_locked" key that is True + if both key and value updates should be locked and False otherwise, and + a "key_locked" key that is True if only key updates should be locked (value + updates still allowed) and False otherwise + """ + if lock_state["all_locked"]: + self.lock() + if lock_state["key_locked"]: + self.lock_keys() + + @contextlib.contextmanager + def unlocked(self): + """ + A context scope for modifying a Config object. Within the scope, + both keys and values can be updated. Upon leaving the scope, + the initial level of locking is restored. + """ + lock_state = self._get_lock_state() + self.unlock() + yield + self._set_lock_state(lock_state) + + @contextlib.contextmanager + def values_unlocked(self): + """ + A context scope for modifying a Config object. Within the scope, + only values can be updated (new keys cannot be created). Upon + leaving the scope, the initial level of locking is restored. + """ + lock_state = self._get_lock_state() + self.unlock() + self.lock_keys() + yield + self._set_lock_state(lock_state) + + def lock_keys(self): + """ + Lock this config so that new keys cannot be added. + """ + if not self.key_lockable: + return + object.__setattr__(self, '__key_locked', True) + for k in self: + if isinstance(self[k], Config): + self[k].lock_keys() + + def unlock_keys(self): + """ + Unlock this config so that new keys can be added. + """ + object.__setattr__(self, '__key_locked', False) + for k in self: + if isinstance(self[k], Config): + self[k].unlock_keys() + + @property + def is_locked(self): + """ + Returns True if the config is locked (no key or value updates allowed). + """ + return object.__getattribute__(self, '__all_locked') + + @property + def is_key_locked(self): + """ + Returns True if the config is key-locked (no key updates allowed). + """ + return object.__getattribute__(self, '__key_locked') + + def do_not_lock_keys(self): + """ + Calling this function on this config indicates that key updates should be + allowed even when this config is key-locked (but not when it is completely + locked). This is convenient for attributes that contain kwargs, where there + might be a variable type and number of arguments contained in the sub-config. + """ + object.__setattr__(self, '__do_not_lock_keys', True) + + @property + def key_lockable(self): + """ + Returns true if this config is key-lockable (new keys cannot be inserted in a + key-locked lock level). + """ + return not object.__getattribute__(self, '__do_not_lock_keys') + + def __setattr__(self, name, value): + if self.is_locked: + raise RuntimeError("This config has been locked - cannot set attribute '{}' to {}".format(name, value)) + + if hasattr(Config, name): + raise AttributeError("'Dict' object attribute " + "'{0}' is read-only".format(name)) + elif not hasattr(self, name) and self.is_key_locked: + raise RuntimeError("This config is key-locked - cannot add key '{}'".format(name)) + else: + self[name] = value + + def __setitem__(self, name, value): + super(Config, self).__setitem__(name, value) + p = object.__getattribute__(self, '__parent') + key = object.__getattribute__(self, '__key') + if p is not None: + p[key] = self + + def __add__(self, other): + if not self.keys(): + return other + else: + self_type = type(self).__name__ + other_type = type(other).__name__ + msg = "unsupported operand type(s) for +: '{}' and '{}'" + raise TypeError(msg.format(self_type, other_type)) + + @classmethod + def _hook(cls, item): + if isinstance(item, dict): + # We return Config instance instead of cls instance to ensure all sub-configs are not a top-level class + return Config(item) + elif isinstance(item, (list, tuple)): + return type(item)(Config._hook(elem) for elem in item) + return item + + def __getattr__(self, item): + return self.__getitem__(item) + + def __repr__(self): + json_string = json.dumps(self.to_dict(), indent=4) + return json_string + + def __getitem__(self, name): + if name not in self: + if object.__getattribute__(self, '__all_locked') or object.__getattribute__(self, '__key_locked'): + raise RuntimeError("This config has been locked and '{}' is not in this config".format(name)) + return Config(__parent=self, __key=name) + return super(Config, self).__getitem__(name) + + def __delattr__(self, name): + del self[name] + + def to_dict(self): + base = {} + for key, value in self.items(): + if isinstance(value, type(self)): + base[key] = value.to_dict() + elif isinstance(value, (list, tuple)): + base[key] = type(value)( + item.to_dict() if isinstance(item, type(self)) else + item for item in value) + else: + base[key] = value + return base + + def copy(self): + return copy.copy(self) + + def deepcopy(self): + return copy.deepcopy(self) + + def __deepcopy__(self, memo): + other = self.__class__() + memo[id(self)] = other + for key, value in self.items(): + other[copy.deepcopy(key, memo)] = copy.deepcopy(value, memo) + return other + + def update(self, *args, **kwargs): + """ + Update this config using another config or nested dictionary. + """ + if self.is_locked: + raise RuntimeError('Cannot update - this config has been locked') + other = {} + if args: + if len(args) > 1: + raise TypeError() + other.update(args[0]) + other.update(kwargs) + for k, v in other.items(): + if self.is_key_locked and k not in self: + raise RuntimeError("Cannot update - this config has been key-locked and key '{}' does not exist".format(k)) + if (not isinstance(self[k], dict)) or (not isinstance(v, dict)): + self[k] = v + else: + self[k].update(v) + + def __getnewargs__(self): + return tuple(self.items()) + + def __getstate__(self): + return self + + def __setstate__(self, state): + self.update(state) + + def setdefault(self, key, default=None): + if key in self: + return self[key] + else: + self[key] = default + return default + + def dump(self, filename=None): + """ + Dumps the config to a json. + + Args: + filename (str): if not None, save to json file. + + Returns: + json_string (str): json string representation of + this config + """ + json_string = json.dumps(self.to_dict(), indent=4) + if filename is not None: + f = open(filename, "w") + f.write(json_string) + f.close() + return json_string \ No newline at end of file diff --git a/aloha-devel/robomimic/config/hbc_config.py b/aloha-devel/robomimic/config/hbc_config.py new file mode 100644 index 0000000000000000000000000000000000000000..ae65c9b85fc168dc65666392fd334810b622ab04 --- /dev/null +++ b/aloha-devel/robomimic/config/hbc_config.py @@ -0,0 +1,96 @@ +""" +Config for HBC algorithm. +""" + +from robomimic.config.base_config import BaseConfig +from robomimic.config.gl_config import GLConfig +from robomimic.config.bc_config import BCConfig + + +class HBCConfig(BaseConfig): + ALGO_NAME = "hbc" + + def train_config(self): + """ + Update from superclass to change default sequence length to load from dataset. + """ + super(HBCConfig, self).train_config() + self.train.seq_length = 10 # length of experience sequence to fetch from the buffer + + def algo_config(self): + """ + This function populates the `config.algo` attribute of the config, and is given to the + `Algo` subclass (see `algo/algo.py`) for each algorithm through the `algo_config` + argument to the constructor. Any parameter that an algorithm needs to determine its + training and test-time behavior should be populated here. + """ + + # One of ["separate", "cascade"]. In "separate" mode (default), + # the planner and actor are trained independently and then the planner subgoal predictions are + # used to condition the actor at test-time. In "cascade" mode, the actor is trained directly + # on planner subgoal predictions. In "actor_only" mode, only the actor is trained, and in + # "planner_only" mode, only the planner is trained. + self.algo.mode = "separate" + self.algo.actor_use_random_subgoals = False # whether to sample subgoal index from [1, subgoal_horizon] + self.algo.subgoal_update_interval = 10 # how frequently the subgoal should be updated at test-time + + + # ================== Latent Subgoal Config ================== + self.algo.latent_subgoal.enabled = False # if True, use VAE latent space as subgoals for actor, instead of reconstructions + + # prior correction trick for actor and value training: instead of using encoder for + # transforming subgoals to latent subgoals, generate prior samples and choose + # the closest one to the encoder output + self.algo.latent_subgoal.prior_correction.enabled = False + self.algo.latent_subgoal.prior_correction.num_samples = 100 + + # ================== Planner Config ================== + self.algo.planner = GLConfig().algo # config for goal learning + # set subgoal horizon explicitly + self.algo.planner.subgoal_horizon = 10 + # ensure VAE is used + self.algo.planner.vae.enabled = True + + # ================== Actor Config =================== + self.algo.actor = BCConfig().algo + # use RNN + self.algo.actor.rnn.enabled = True + self.algo.actor.rnn.horizon = 10 + # remove unused parts of BCConfig algo config + del self.algo.actor.gaussian + del self.algo.actor.gmm + del self.algo.actor.vae + + def observation_config(self): + """ + Update from superclass so that planner and actor each get their own observation config. + """ + self.observation.planner = GLConfig().observation + self.observation.actor = BCConfig().observation + + @property + def use_goals(self): + """ + Update from superclass - planner goal modalities determine goal-conditioning + """ + return len( + self.observation.planner.modalities.goal.low_dim + + self.observation.planner.modalities.goal.rgb) > 0 + + @property + def all_obs_keys(self): + """ + Update from superclass to include modalities from planner and actor. + """ + # pool all modalities + return sorted(tuple(set([ + obs_key for group in [ + self.observation.planner.modalities.obs.values(), + self.observation.planner.modalities.goal.values(), + self.observation.planner.modalities.subgoal.values(), + self.observation.actor.modalities.obs.values(), + self.observation.actor.modalities.goal.values(), + ] + for modality in group + for obs_key in modality + ]))) diff --git a/aloha-devel/robomimic/config/iql_config.py b/aloha-devel/robomimic/config/iql_config.py new file mode 100644 index 0000000000000000000000000000000000000000..bd603d1aa0183639971b16747c5020afa6d04fe3 --- /dev/null +++ b/aloha-devel/robomimic/config/iql_config.py @@ -0,0 +1,73 @@ +""" +Config for IQL algorithm. +""" + +from robomimic.config.base_config import BaseConfig + + +class IQLConfig(BaseConfig): + ALGO_NAME = "iql" + + def algo_config(self): + """ + This function populates the `config.algo` attribute of the config, and is given to the + `Algo` subclass (see `algo/algo.py`) for each algorithm through the `algo_config` + argument to the constructor. Any parameter that an algorithm needs to determine its + training and test-time behavior should be populated here. + """ + super(IQLConfig, self).algo_config() + + # optimization parameters + self.algo.optim_params.critic.learning_rate.initial = 1e-4 # critic learning rate + self.algo.optim_params.critic.learning_rate.decay_factor = 0.0 # factor to decay LR by (if epoch schedule non-empty) + self.algo.optim_params.critic.learning_rate.epoch_schedule = [] # epochs where LR decay occurs + self.algo.optim_params.critic.regularization.L2 = 0.00 # L2 regularization strength + + self.algo.optim_params.vf.learning_rate.initial = 1e-4 # vf learning rate + self.algo.optim_params.vf.learning_rate.decay_factor = 0.0 # factor to decay LR by (if epoch schedule non-empty) + self.algo.optim_params.vf.learning_rate.epoch_schedule = [] # epochs where LR decay occurs + self.algo.optim_params.vf.regularization.L2 = 0.00 # L2 regularization strength + + self.algo.optim_params.actor.learning_rate.initial = 1e-4 # actor learning rate + self.algo.optim_params.actor.learning_rate.decay_factor = 0.0 # factor to decay LR by (if epoch schedule non-empty) + self.algo.optim_params.actor.learning_rate.epoch_schedule = [] # epochs where LR decay occurs + self.algo.optim_params.actor.regularization.L2 = 0.00 # L2 regularization strength + + # target network related parameters + self.algo.discount = 0.99 # discount factor to use + self.algo.target_tau = 0.01 # update rate for target networks + + # ================== Actor Network Config =================== + # Actor network settings + self.algo.actor.net.type = "gaussian" # Options are currently ["gaussian", "gmm"] + + # Actor network settings - shared + self.algo.actor.net.common.std_activation = "softplus" # Activation to use for std output from policy net + self.algo.actor.net.common.low_noise_eval = True # Whether to use deterministic action sampling at eval stage + self.algo.actor.net.common.use_tanh = False # Whether to use tanh at output of actor network + + # Actor network settings - gaussian + self.algo.actor.net.gaussian.init_last_fc_weight = 0.001 # If set, will override the initialization of the final fc layer to be uniformly sampled limited by this value + self.algo.actor.net.gaussian.init_std = 0.3 # Relative scaling factor for std from policy net + self.algo.actor.net.gaussian.fixed_std = False # Whether to learn std dev or not + + self.algo.actor.net.gmm.num_modes = 5 # number of GMM modes + self.algo.actor.net.gmm.min_std = 0.0001 # minimum std output from network + + self.algo.actor.layer_dims = (300, 400) # actor MLP layer dimensions + + self.algo.actor.max_gradient_norm = None # L2 gradient clipping for actor + + # ================== Critic Network Config =================== + # critic ensemble parameters + self.algo.critic.ensemble.n = 2 # number of Q networks in the ensemble + self.algo.critic.layer_dims = (300, 400) # critic MLP layer dimensions + self.algo.critic.use_huber = False # Huber Loss instead of L2 for critic + self.algo.critic.max_gradient_norm = None # L2 gradient clipping for actor + + # ================== Adv Config ============================== + self.algo.adv.clip_adv_value = None # whether to clip raw advantage estimates + self.algo.adv.beta = 1.0 # temperature for operator + self.algo.adv.use_final_clip = True # whether to clip final weight calculations + + self.algo.vf_quantile = 0.9 # quantile factor in quantile regression diff --git a/aloha-devel/robomimic/envs/__init__.py b/aloha-devel/robomimic/envs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/aloha-devel/robomimic/envs/env_base.py b/aloha-devel/robomimic/envs/env_base.py new file mode 100644 index 0000000000000000000000000000000000000000..ee3184c2aeff786eb0386e0e686de24e83610a9b --- /dev/null +++ b/aloha-devel/robomimic/envs/env_base.py @@ -0,0 +1,212 @@ +""" +This file contains the base class for environment wrappers that are used +to provide a standardized environment API for training policies and interacting +with metadata present in datasets. +""" +import abc + + +class EnvType: + """ + Holds environment types - one per environment class. + These act as identifiers for different environments. + """ + ROBOSUITE_TYPE = 1 + GYM_TYPE = 2 + IG_MOMART_TYPE = 3 + + +class EnvBase(abc.ABC): + """A base class method for environments used by this repo.""" + @abc.abstractmethod + def __init__( + self, + env_name, + render=False, + render_offscreen=False, + use_image_obs=False, + postprocess_visual_obs=True, + **kwargs, + ): + """ + Args: + env_name (str): name of environment. Only needs to be provided if making a different + environment from the one in @env_meta. + + render (bool): if True, environment supports on-screen rendering + + render_offscreen (bool): if True, environment supports off-screen rendering. This + is forced to be True if @env_meta["use_images"] is True. + + use_image_obs (bool): if True, environment is expected to render rgb image observations + on every env.step call. Set this to False for efficiency reasons, if image + observations are not required. + + postprocess_visual_obs (bool): if True, postprocess image observations + to prepare for learning. This should only be False when extracting observations + for saving to a dataset (to save space on RGB images for example). + """ + return + + @abc.abstractmethod + def step(self, action): + """ + Step in the environment with an action. + + Args: + action (np.array): action to take + + Returns: + observation (dict): new observation dictionary + reward (float): reward for this step + done (bool): whether the task is done + info (dict): extra information + """ + return + + @abc.abstractmethod + def reset(self): + """ + Reset environment. + + Returns: + observation (dict): initial observation dictionary. + """ + return + + @abc.abstractmethod + def reset_to(self, state): + """ + Reset to a specific simulator state. + + Args: + state (dict): current simulator state + + Returns: + observation (dict): observation dictionary after setting the simulator state + """ + return + + @abc.abstractmethod + def render(self, mode="human", height=None, width=None, camera_name=None): + """Render""" + return + + @abc.abstractmethod + def get_observation(self): + """Get environment observation""" + return + + @abc.abstractmethod + def get_state(self): + """Get environment simulator state, compatible with @reset_to""" + return + + @abc.abstractmethod + def get_reward(self): + """ + Get current reward. + """ + return + + @abc.abstractmethod + def get_goal(self): + """ + Get goal observation. Not all environments support this. + """ + return + + @abc.abstractmethod + def set_goal(self, **kwargs): + """ + Set goal observation with external specification. Not all environments support this. + """ + return + + @abc.abstractmethod + def is_done(self): + """ + Check if the task is done (not necessarily successful). + """ + return + + @abc.abstractmethod + def is_success(self): + """ + Check if the task condition(s) is reached. Should return a dictionary + { str: bool } with at least a "task" key for the overall task success, + and additional optional keys corresponding to other task criteria. + """ + return + + @property + @abc.abstractmethod + def action_dimension(self): + """ + Returns dimension of actions (int). + """ + return + + @property + @abc.abstractmethod + def name(self): + """ + Returns name of environment name (str). + """ + return + + @property + @abc.abstractmethod + def type(self): + """ + Returns environment type (int) for this kind of environment. + This helps identify this env class. + """ + return + + @property + def version(self): + """ + Returns version of environment (str). + This is not an abstract method, some subclasses do not implement it + """ + return None + + @abc.abstractmethod + def serialize(self): + """ + Save all information needed to re-instantiate this environment in a dictionary. + This is the same as @env_meta - environment metadata stored in hdf5 datasets, + and used in utils/env_utils.py. + """ + return + + @classmethod + @abc.abstractmethod + def create_for_data_processing(cls, camera_names, camera_height, camera_width, reward_shaping, **kwargs): + """ + Create environment for processing datasets, which includes extracting + observations, labeling dense / sparse rewards, and annotating dones in + transitions. + + Args: + camera_names ([str]): list of camera names that correspond to image observations + camera_height (int): camera height for all cameras + camera_width (int): camera width for all cameras + reward_shaping (bool): if True, use shaped environment rewards, else use sparse task completion rewards + + Returns: + env (EnvBase instance) + """ + return + + @property + @abc.abstractmethod + def rollout_exceptions(self): + """ + Return tuple of exceptions to except when doing rollouts. This is useful to ensure + that the entire training run doesn't crash because of a bad policy that causes unstable + simulation computations. + """ + return + diff --git a/aloha-devel/robomimic/envs/env_gym.py b/aloha-devel/robomimic/envs/env_gym.py new file mode 100644 index 0000000000000000000000000000000000000000..6cb1fa56586ad3698db2ee362f56282cb934256e --- /dev/null +++ b/aloha-devel/robomimic/envs/env_gym.py @@ -0,0 +1,247 @@ +""" +This file contains the gym environment wrapper that is used +to provide a standardized environment API for training policies and interacting +with metadata present in datasets. +""" +import json +import numpy as np +from copy import deepcopy + +import gym +try: + import d4rl +except: + print("WARNING: could not load d4rl environments!") + +import robomimic.envs.env_base as EB +import robomimic.utils.obs_utils as ObsUtils + + +class EnvGym(EB.EnvBase): + """Wrapper class for gym""" + def __init__( + self, + env_name, + render=False, + render_offscreen=False, + use_image_obs=False, + postprocess_visual_obs=True, + **kwargs, + ): + """ + Args: + env_name (str): name of environment. Only needs to be provided if making a different + environment from the one in @env_meta. + + render (bool): ignored - gym envs always support on-screen rendering + + render_offscreen (bool): ignored - gym envs always support off-screen rendering + + use_image_obs (bool): ignored - gym envs don't typically use images + + postprocess_visual_obs (bool): ignored - gym envs don't typically use images + """ + self._init_kwargs = deepcopy(kwargs) + self._env_name = env_name + self._current_obs = None + self._current_reward = None + self._current_done = None + self._done = None + self.env = gym.make(env_name, **kwargs) + + def step(self, action): + """ + Step in the environment with an action. + + Args: + action (np.array): action to take + + Returns: + observation (dict): new observation dictionary + reward (float): reward for this step + done (bool): whether the task is done + info (dict): extra information + """ + obs, reward, done, info = self.env.step(action) + self._current_obs = obs + self._current_reward = reward + self._current_done = done + return self.get_observation(obs), reward, self.is_done(), info + + def reset(self): + """ + Reset environment. + + Returns: + observation (dict): initial observation dictionary. + """ + self._current_obs = self.env.reset() + self._current_reward = None + self._current_done = None + return self.get_observation(self._current_obs) + + def reset_to(self, state): + """ + Reset to a specific simulator state. + + Args: + state (dict): current simulator state that contains: + - states (np.ndarray): initial state of the mujoco environment + + Returns: + observation (dict): observation dictionary after setting the simulator state + """ + if hasattr(self.env.unwrapped.sim, "set_state_from_flattened"): + self.env.unwrapped.sim.set_state_from_flattened(state["states"]) + self.env.unwrapped.sim.forward() + return { "flat" : self.env.unwrapped._get_obs() } + else: + raise NotImplementedError + + def render(self, mode="human", height=None, width=None, camera_name=None, **kwargs): + """ + Render from simulation to either an on-screen window or off-screen to RGB array. + + Args: + mode (str): pass "human" for on-screen rendering or "rgb_array" for off-screen rendering + height (int): height of image to render - only used if mode is "rgb_array" + width (int): width of image to render - only used if mode is "rgb_array" + """ + if mode =="human": + return self.env.render(mode=mode, **kwargs) + if mode == "rgb_array": + return self.env.render(mode="rgb_array", height=height, width=width) + else: + raise NotImplementedError("mode={} is not implemented".format(mode)) + + def get_observation(self, obs=None): + """ + Get current environment observation dictionary. + + Args: + ob (np.array): current flat observation vector to wrap and provide as a dictionary. + If not provided, uses self._current_obs. + """ + if obs is None: + assert self._current_obs is not None + obs = self._current_obs + return { "flat" : np.copy(obs) } + + def get_state(self): + """ + Get current environment simulator state as a dictionary. Should be compatible with @reset_to. + """ + # NOTE: assumes MuJoCo gym task! + xml = self.env.sim.model.get_xml() # model xml file + state = np.array(self.env.sim.get_state().flatten()) # simulator state + return dict(model=xml, states=state) + + def get_reward(self): + """ + Get current reward. + """ + assert self._current_reward is not None + return self._current_reward + + def get_goal(self): + """ + Get goal observation. Not all environments support this. + """ + raise NotImplementedError + + def set_goal(self, **kwargs): + """ + Set goal observation with external specification. Not all environments support this. + """ + raise NotImplementedError + + def is_done(self): + """ + Check if the task is done (not necessarily successful). + """ + assert self._current_done is not None + return self._current_done + + def is_success(self): + """ + Check if the task condition(s) is reached. Should return a dictionary + { str: bool } with at least a "task" key for the overall task success, + and additional optional keys corresponding to other task criteria. + """ + if hasattr(self.env.unwrapped, "_check_success"): + return self.env.unwrapped._check_success() + + # gym envs generally don't check task success - we only compare returns + return { "task" : False } + + @property + def action_dimension(self): + """ + Returns dimension of actions (int). + """ + return self.env.action_space.shape[0] + + @property + def name(self): + """ + Returns name of environment name (str). + """ + return self._env_name + + @property + def type(self): + """ + Returns environment type (int) for this kind of environment. + This helps identify this env class. + """ + return EB.EnvType.GYM_TYPE + + def serialize(self): + """ + Save all information needed to re-instantiate this environment in a dictionary. + This is the same as @env_meta - environment metadata stored in hdf5 datasets, + and used in utils/env_utils.py. + """ + return dict(env_name=self.name, type=self.type, env_kwargs=deepcopy(self._init_kwargs)) + + @classmethod + def create_for_data_processing(cls, env_name, camera_names, camera_height, camera_width, reward_shaping, **kwargs): + """ + Create environment for processing datasets, which includes extracting + observations, labeling dense / sparse rewards, and annotating dones in + transitions. For gym environments, input arguments (other than @env_name) + are ignored, since environments are mostly pre-configured. + + Args: + env_name (str): name of gym environment to create + + Returns: + env (EnvGym instance) + """ + + # make sure to initialize obs utils so it knows which modalities are image modalities. + # For currently supported gym tasks, there are no image observations. + obs_modality_specs = { + "obs": { + "low_dim": ["flat"], + "rgb": [], + } + } + ObsUtils.initialize_obs_utils_with_obs_specs(obs_modality_specs) + + return cls(env_name=env_name, **kwargs) + + @property + def rollout_exceptions(self): + """ + Return tuple of exceptions to except when doing rollouts. This is useful to ensure + that the entire training run doesn't crash because of a bad policy that causes unstable + simulation computations. + """ + return () + + def __repr__(self): + """ + Pretty-print env description. + """ + return self.name + "\n" + json.dumps(self._init_kwargs, sort_keys=True, indent=4) diff --git a/aloha-devel/robomimic/envs/env_ig_momart.py b/aloha-devel/robomimic/envs/env_ig_momart.py new file mode 100644 index 0000000000000000000000000000000000000000..81dd312fe784e31ef41d68783113ee39fb7b0209 --- /dev/null +++ b/aloha-devel/robomimic/envs/env_ig_momart.py @@ -0,0 +1,395 @@ +""" +Wrapper environment class to enable using iGibson-based environments used in the MOMART paper +""" + +from copy import deepcopy +import numpy as np +import json + +import pybullet as p +import gibson2 +from gibson2.envs.semantic_organize_and_fetch import SemanticOrganizeAndFetch +from gibson2.utils.custom_utils import ObjectConfig +import gibson2.external.pybullet_tools.utils as PBU +import tempfile +import os +import yaml +import cv2 + +import robomimic.utils.obs_utils as ObsUtils +import robomimic.envs.env_base as EB + + +# TODO: Once iG 2.0 is more stable, automate available environments, similar to robosuite +ENV_MAPPING = { + "SemanticOrganizeAndFetch": SemanticOrganizeAndFetch, +} + + +class EnvGibsonMOMART(EB.EnvBase): + """ + Wrapper class for gibson environments (https://github.com/StanfordVL/iGibson) specifically compatible with + MoMaRT datasets + """ + def __init__( + self, + env_name, + ig_config, + postprocess_visual_obs=True, + render=False, + render_offscreen=False, + use_image_obs=False, + image_height=None, + image_width=None, + physics_timestep=1./240., + action_timestep=1./20., + **kwargs, + ): + """ + Args: + ig_config (dict): YAML configuration to use for iGibson, as a dict + + postprocess_visual_obs (bool): if True, postprocess image observations + to prepare for learning + + render (bool): if True, environment supports on-screen rendering + + render_offscreen (bool): if True, environment supports off-screen rendering. This + is forced to be True if @use_image_obs is True. + + use_image_obs (bool): if True, environment is expected to render rgb image observations + on every env.step call. Set this to False for efficiency reasons, if image + observations are not required. + + render_mode (str): How to run simulation rendering. Options are {"pbgui", "iggui", or "headless"} + + image_height (int): If specified, overrides internal iG image height when rendering + + image_width (int): If specified, overrides internal iG image width when rendering + + physics_timestep (float): Pybullet physics timestep to use + + action_timestep (float): Action timestep to use for robot in simulation + + kwargs (unrolled dict): Any args to substitute in the ig_configuration + """ + self._env_name = env_name + self.ig_config = deepcopy(ig_config) + self.postprocess_visual_obs = postprocess_visual_obs + self._init_kwargs = kwargs + + # Determine rendering mode + self.render_mode = "iggui" if render else "headless" + self.render_onscreen = render + + # Make sure rgb is part of obs in ig config + self.ig_config["output"] = list(set(self.ig_config["output"] + ["rgb"])) + + # Warn user that iG always uses a renderer + if (not render) and (not render_offscreen): + print("WARNING: iGibson always uses a renderer -- using headless by default.") + + # Update ig config + for k, v in kwargs.items(): + assert k in self.ig_config, f"Got unknown ig configuration key {k}!" + self.ig_config[k] = v + + # Set rendering values + self.obs_img_height = image_height if image_height is not None else self.ig_config.get("obs_image_height", 120) + self.obs_img_width = image_width if image_width is not None else self.ig_config.get("obs_image_width", 120) + + # Get class to create + envClass = ENV_MAPPING.get(self._env_name, None) + + # Make sure we have a valid environment class + assert envClass is not None, "No valid environment for the requested task was found!" + + # Set device idx for rendering + # ensure that we select the correct GPU device for rendering by testing for EGL rendering + # NOTE: this package should be installed from this link (https://github.com/StanfordVL/egl_probe) + import egl_probe + device_idx = 0 + valid_gpu_devices = egl_probe.get_available_devices() + if len(valid_gpu_devices) > 0: + device_idx = valid_gpu_devices[0] + + # Create environment + self.env = envClass( + config_file=deepcopy(self.ig_config), + mode=self.render_mode, + physics_timestep=physics_timestep, + action_timestep=action_timestep, + device_idx=device_idx, + ) + + # If we have a viewer, make sure to remove all bodies belonging to the visual markers + self.exclude_body_ids = [] # Bodies to exclude when saving state + if self.env.simulator.viewer is not None: + self.exclude_body_ids.append(self.env.simulator.viewer.constraint_marker.body_id) + self.exclude_body_ids.append(self.env.simulator.viewer.constraint_marker2.body_id) + + def step(self, action): + """ + Step in the environment with an action + + Args: + action: action to take + + Returns: + observation: new observation + reward: step reward + done: whether the task is done + info: extra information + """ + obs, r, done, info = self.env.step(action) + obs = self.get_observation(obs) + return obs, r, self.is_done(), info + + def reset(self): + """Reset environment""" + di = self.env.reset() + return self.get_observation(di) + + def reset_to(self, state): + """ + Reset to a specific state + Args: + state (dict): contains: + - states (np.ndarray): initial state of the mujoco environment + - goal (dict): goal components to reset + Returns: + new observation + """ + if "states" in state: + self.env.reset_to(state["states"], exclude=self.exclude_body_ids) + + if "goal" in state: + self.set_goal(**state["goal"]) + + # Return obs + return self.get_observation() + + def render(self, mode="human", camera_name="rgb", height=None, width=None): + """ + Render + + Args: + mode (str): Mode(s) to render. Options are either 'human' (rendering onscreen) or 'rgb' (rendering to + frames offscreen) + camera_name (str): Name of the camera to use -- valid options are "rgb" or "rgb_wrist" + height (int): If specified with width, resizes the rendered image to this height + width (int): If specified with height, resizes the rendered image to this width + + Returns: + array or None: If rendering to frame, returns the rendered frame. Otherwise, returns None + """ + # Only robotview camera is currently supported + assert camera_name in {"rgb", "rgb_wrist"}, \ + f"Only rgb, rgb_wrist cameras currently supported, got {camera_name}." + + if mode == "human": + assert self.render_onscreen, "Rendering has not been enabled for onscreen!" + self.env.simulator.sync() + else: + assert self.env.simulator.renderer is not None, "No renderer enabled for this env!" + + frame = self.env.sensors["vision"].get_obs(self.env)[camera_name] + + # Reshape all frames + if height is not None and width is not None: + frame = cv2.resize(frame, dsize=(height, width), interpolation=cv2.INTER_CUBIC) + return frame + + def resize_obs_frame(self, frame): + """ + Resizes frame to be internal height and width values + """ + return cv2.resize(frame, dsize=(self.obs_img_width, self.obs_img_height), interpolation=cv2.INTER_CUBIC) + + def get_observation(self, di=None): + """Get environment observation""" + if di is None: + di = self.env.get_state() + ret = {} + for k in di: + # RGB Images + if "rgb" in k: + ret[k] = di[k] + # ret[k] = np.transpose(di[k], (2, 0, 1)) + if self.postprocess_visual_obs: + ret[k] = ObsUtils.process_obs(obs=self.resize_obs_frame(ret[k]), obs_key=k) + + # Depth images + elif "depth" in k: + # ret[k] = np.transpose(di[k], (2, 0, 1)) + # Values can be corrupted (negative or > 1.0, so we clip values) + ret[k] = np.clip(di[k], 0.0, 1.0) + if self.postprocess_visual_obs: + ret[k] = ObsUtils.process_obs(obs=self.resize_obs_frame(ret[k])[..., None], obs_key=k) + + # Segmentation Images + elif "seg" in k: + ret[k] = di[k][..., None] + if self.postprocess_visual_obs: + ret[k] = ObsUtils.process_obs(obs=self.resize_obs_frame(ret[k]), obs_key=k) + + # Scans + elif "scan" in k: + ret[k] = np.transpose(np.array(di[k]), axes=(1, 0)) + + # Compose proprio obs + proprio_obs = di["proprio"] + + # Compute intermediate values + lin_vel = np.linalg.norm(proprio_obs["base_lin_vel"][:2]) + ang_vel = proprio_obs["base_ang_vel"][2] + + ret["proprio"] = np.concatenate([ + proprio_obs["head_joint_pos"], + proprio_obs["grasped"], + proprio_obs["eef_pos"], + proprio_obs["eef_quat"], + ]) + + # Proprio info that's only relevant for navigation + ret["proprio_nav"] = np.concatenate([ + [lin_vel], + [ang_vel], + ]) + + # Compose task obs + ret["object"] = np.concatenate([ + np.array(di["task_obs"]["object-state"]), + ]) + + # Add ground truth navigational state + ret["gt_nav"] = np.concatenate([ + proprio_obs["base_pos"][:2], + [np.sin(proprio_obs["base_rpy"][2])], + [np.cos(proprio_obs["base_rpy"][2])], + ]) + + return ret + + def sync_task(self): + """ + Method to synchronize iG task, since we're not actually resetting the env but instead setting states directly. + Should only be called after resetting the initial state of an episode + """ + self.env.task.update_target_object_init_pos() + self.env.task.update_location_info() + + def set_task_conditions(self, task_conditions): + """ + Method to override task conditions (e.g.: target object), useful in cases such as playing back + from demonstrations + + Args: + task_conditions (dict): Keyword-mapped arguments to pass to task instance to set internally + """ + self.env.set_task_conditions(task_conditions) + + def get_state(self): + """Get iG flattened state""" + return {"states": PBU.WorldSaver(exclude_body_ids=self.exclude_body_ids).serialize()} + + def get_reward(self): + return self.env.task.get_reward(self.env)[0] + # return float(self.is_success()["task"]) + + def get_goal(self): + """Get goal specification""" + # No support yet in iG + raise NotImplementedError + + def set_goal(self, **kwargs): + """Set env target with external specification""" + # No support yet in iG + raise NotImplementedError + + def is_done(self): + """Check if the agent is done (not necessarily successful).""" + return False + + def is_success(self): + """ + Check if the task condition(s) is reached. Should return a dictionary + { str: bool } with at least a "task" key for the overall task success, + and additional optional keys corresponding to other task criteria. + """ + succ = self.env.check_success() + if isinstance(succ, dict): + assert "task" in succ + return succ + return { "task" : succ } + + @classmethod + def create_for_data_processing( + cls, + env_name, + camera_names, + camera_height, + camera_width, + reward_shaping, + **kwargs, + ): + """ + Create environment for processing datasets, which includes extracting + observations, labeling dense / sparse rewards, and annotating dones in + transitions. + + Args: + env_name (str): name of environment + camera_names (list of str): list of camera names that correspond to image observations + camera_height (int): camera height for all cameras + camera_width (int): camera width for all cameras + reward_shaping (bool): if True, use shaped environment rewards, else use sparse task completion rewards + """ + has_camera = (len(camera_names) > 0) + + # note that @postprocess_visual_obs is False since this env's images will be written to a dataset + return cls( + env_name=env_name, + render=False, + render_offscreen=has_camera, + use_image_obs=has_camera, + postprocess_visual_obs=False, + image_height=camera_height, + image_width=camera_width, + **kwargs, + ) + + @property + def action_dimension(self): + """Action dimension""" + return self.env.robots[0].action_dim + + @property + def name(self): + """Environment name""" + return self._env_name + + @property + def type(self): + """Environment type""" + return EB.EnvType.IG_MOMART_TYPE + + def serialize(self): + """Serialize to dictionary""" + return dict(env_name=self.name, type=self.type, + ig_config=self.ig_config, + env_kwargs=deepcopy(self._init_kwargs)) + + @classmethod + def deserialize(cls, info, postprocess_visual_obs=True): + """Create environment with external info""" + return cls(env_name=info["env_name"], ig_config=info["ig_config"], postprocess_visual_obs=postprocess_visual_obs, **info["env_kwargs"]) + + @property + def rollout_exceptions(self): + """Return tuple of exceptions to except when doing rollouts""" + return (RuntimeError) + + def __repr__(self): + return self.name + "\n" + json.dumps(self._init_kwargs, sort_keys=True, indent=4) + \ + "\niGibson Config: \n" + json.dumps(self.ig_config, sort_keys=True, indent=4) diff --git a/aloha-devel/robomimic/envs/env_robosuite.py b/aloha-devel/robomimic/envs/env_robosuite.py new file mode 100644 index 0000000000000000000000000000000000000000..675e7dd28ea740e4bbab62fbd355e2c5ac190fce --- /dev/null +++ b/aloha-devel/robomimic/envs/env_robosuite.py @@ -0,0 +1,415 @@ +""" +This file contains the robosuite environment wrapper that is used +to provide a standardized environment API for training policies and interacting +with metadata present in datasets. +""" +import json +import numpy as np +from copy import deepcopy + +import robosuite + +import robomimic.utils.obs_utils as ObsUtils +import robomimic.utils.lang_utils as LangUtils +import robomimic.envs.env_base as EB + + +class EnvRobosuite(EB.EnvBase): + """Wrapper class for robosuite environments (https://github.com/ARISE-Initiative/robosuite)""" + def __init__( + self, + env_name, + render=False, + render_offscreen=False, + use_image_obs=False, + postprocess_visual_obs=True, + lang=None, + **kwargs, + ): + """ + Args: + env_name (str): name of environment. Only needs to be provided if making a different + environment from the one in @env_meta. + + render (bool): if True, environment supports on-screen rendering + + render_offscreen (bool): if True, environment supports off-screen rendering. This + is forced to be True if @env_meta["use_images"] is True. + + use_image_obs (bool): if True, environment is expected to render rgb image observations + on every env.step call. Set this to False for efficiency reasons, if image + observations are not required. + + postprocess_visual_obs (bool): if True, postprocess image observations + to prepare for learning. This should only be False when extracting observations + for saving to a dataset (to save space on RGB images for example). + + lang: TODO add documentation + """ + self.postprocess_visual_obs = postprocess_visual_obs + + # robosuite version check + self._is_v1 = (robosuite.__version__.split(".")[0] == "1") + if self._is_v1: + assert (int(robosuite.__version__.split(".")[1]) >= 2), "only support robosuite v0.3 and v1.2+" + + kwargs = deepcopy(kwargs) + + # update kwargs based on passed arguments + update_kwargs = dict( + has_renderer=render, + has_offscreen_renderer=(render_offscreen or use_image_obs), + ignore_done=True, + use_object_obs=True, + use_camera_obs=use_image_obs, + camera_depths=False, + ) + kwargs.update(update_kwargs) + + if self._is_v1: + if kwargs["has_offscreen_renderer"]: + # ensure that we select the correct GPU device for rendering by testing for EGL rendering + # NOTE: this package should be installed from this link (https://github.com/StanfordVL/egl_probe) + import egl_probe + valid_gpu_devices = egl_probe.get_available_devices() + if len(valid_gpu_devices) > 0: + kwargs["render_gpu_device_id"] = valid_gpu_devices[0] + else: + # make sure gripper visualization is turned off (we almost always want this for learning) + kwargs["gripper_visualization"] = False + del kwargs["camera_depths"] + kwargs["camera_depth"] = False # rename kwarg + + self._env_name = env_name + self._init_kwargs = deepcopy(kwargs) + self.env = robosuite.make(self._env_name, **kwargs) + self.base_env = self.env # for mimicgen + self.lang = lang + self._lang_emb = LangUtils.get_lang_emb(self.lang) + + if self._is_v1: + # Make sure joint position observations and eef vel observations are active + for ob_name in self.env.observation_names: + if ("joint_pos" in ob_name) or ("eef_vel" in ob_name): + self.env.modify_observable(observable_name=ob_name, attribute="active", modifier=True) + + def step(self, action): + """ + Step in the environment with an action. + + Args: + action (np.array): action to take + + Returns: + observation (dict): new observation dictionary + reward (float): reward for this step + done (bool): whether the task is done + info (dict): extra information + """ + obs, r, done, info = self.env.step(action) + obs = self.get_observation(obs) + info["is_success"] = self.is_success() + return obs, r, self.is_done(), info + + def reset(self): + """ + Reset environment. + + Returns: + observation (dict): initial observation dictionary. + """ + di = self.env.reset() + return self.get_observation(di) + + #notifies the environment whether or not the next environemnt testing object should update its category + def update_env(self, attr, value): + self.env.attr = value + + + def reset_to(self, state): + """ + Reset to a specific simulator state. + + Args: + state (dict): current simulator state that contains one or more of: + - states (np.ndarray): initial state of the mujoco environment + - model (str): mujoco scene xml + + Returns: + observation (dict): observation dictionary after setting the simulator state (only + if "states" is in @state) + """ + should_ret = False + if "model" in state: + if state.get("ep_meta", None) is not None: + # set relevant episode information + ep_meta = json.loads(state["ep_meta"]) + self.env.set_attrs_from_ep_meta(ep_meta) + + # this reset is necessary. + # while the call to env.reset_from_xml_string does call reset, + # that is only a "soft" reset that doesn't actually reload the model. + self.reset() + robosuite_version_id = int(robosuite.__version__.split(".")[1]) + if robosuite_version_id <= 3: + from robosuite.utils.mjcf_utils import postprocess_model_xml + xml = postprocess_model_xml(state["model"]) + else: + # v1.4 and above use the class-based edit_model_xml function + xml = self.env.edit_model_xml(state["model"]) + self.env.reset_from_xml_string(xml) + self.env.sim.reset() + if not self._is_v1: + # hide teleop visualization after restoring from model + self.env.sim.model.site_rgba[self.env.eef_site_id] = np.array([0., 0., 0., 0.]) + self.env.sim.model.site_rgba[self.env.eef_cylinder_id] = np.array([0., 0., 0., 0.]) + if "states" in state: + self.env.sim.set_state_from_flattened(state["states"]) + self.env.sim.forward() + should_ret = True + + if "goal" in state: + self.set_goal(**state["goal"]) + if should_ret: + # only return obs if we've done a forward call - otherwise the observations will be garbage + return self.get_observation() + return None + + def render(self, mode="human", height=None, width=None, camera_name=None): + """ + Render from simulation to either an on-screen window or off-screen to RGB array. + + Args: + mode (str): pass "human" for on-screen rendering or "rgb_array" for off-screen rendering + height (int): height of image to render - only used if mode is "rgb_array" + width (int): width of image to render - only used if mode is "rgb_array" + camera_name (str): camera name to use for rendering + """ + # if camera_name is None, infer from initial env kwargs + if camera_name is None: + camera_name = self._init_kwargs.get("camera_names", ["agentview"])[0] + + if mode == "human": + cam_id = self.env.sim.model.camera_name2id(camera_name) + self.env.viewer.set_camera(cam_id) + return self.env.render() + elif mode == "rgb_array": + return self.env.sim.render(height=height, width=width, camera_name=camera_name)[::-1] + else: + raise NotImplementedError("mode={} is not implemented".format(mode)) + + def get_observation(self, di=None): + """ + Get current environment observation dictionary. + + Args: + di (dict): current raw observation dictionary from robosuite to wrap and provide + as a dictionary. If not provided, will be queried from robosuite. + """ + if di is None: + di = self.env._get_observations(force_update=True) if self._is_v1 else self.env._get_observation() + ret = {} + for k in di: + if (k in ObsUtils.OBS_KEYS_TO_MODALITIES) and ObsUtils.key_is_obs_modality(key=k, obs_modality="rgb"): + ret[k] = di[k][::-1] + if self.postprocess_visual_obs: + ret[k] = ObsUtils.process_obs(obs=ret[k], obs_key=k) + + # "object" key contains object information + ret["object"] = np.array(di["object-state"]) + + if self._is_v1: + for robot in self.env.robots: + # add all robot-arm-specific observations. Note the (k not in ret) check + # ensures that we don't accidentally add robot wrist images a second time + pf = robot.robot_model.naming_prefix + for k in di: + if k.startswith(pf) and (k not in ret) and \ + (not k.endswith("proprio-state")): + ret[k] = np.array(di[k]) + else: + # minimal proprioception for older versions of robosuite + ret["proprio"] = np.array(di["robot-state"]) + ret["eef_pos"] = np.array(di["eef_pos"]) + ret["eef_quat"] = np.array(di["eef_quat"]) + ret["gripper_qpos"] = np.array(di["gripper_qpos"]) + + if self._lang_emb is not None: + ret["lang_emb"] = np.array(self._lang_emb) + return ret + + def get_state(self): + """ + Get current environment simulator state as a dictionary. Should be compatible with @reset_to. + """ + xml = self.env.sim.model.get_xml() # model xml file + state = np.array(self.env.sim.get_state().flatten()) # simulator state + info = dict(model=xml, states=state) + if hasattr(self.env, "get_ep_meta"): + # get ep_meta if applicable + info["ep_meta"] = json.dumps(self.env.get_ep_meta(), indent=4) + return info + + def get_reward(self): + """ + Get current reward. + """ + return self.env.reward() + + def get_goal(self): + """ + Get goal observation. Not all environments support this. + """ + return self.get_observation(self.env._get_goal()) + + def set_goal(self, **kwargs): + """ + Set goal observation with external specification. Not all environments support this. + """ + return self.env.set_goal(**kwargs) + + def is_done(self): + """ + Check if the task is done (not necessarily successful). + """ + + # Robosuite envs always rollout to fixed horizon. + return False + + def is_success(self): + """ + Check if the task condition(s) is reached. Should return a dictionary + { str: bool } with at least a "task" key for the overall task success, + and additional optional keys corresponding to other task criteria. + """ + succ = self.env._check_success() + if isinstance(succ, dict): + assert "task" in succ + return succ + return { "task" : succ } + + @property + def action_dimension(self): + """ + Returns dimension of actions (int). + """ + return self.env.action_spec[0].shape[0] + + @property + def name(self): + """ + Returns name of environment name (str). + """ + return self._env_name + + @property + def type(self): + """ + Returns environment type (int) for this kind of environment. + This helps identify this env class. + """ + return EB.EnvType.ROBOSUITE_TYPE + + @property + def version(self): + """ + Returns version of robosuite used for this environment, eg. 1.2.0 + """ + return robosuite.__version__ + + def serialize(self): + """ + Save all information needed to re-instantiate this environment in a dictionary. + This is the same as @env_meta - environment metadata stored in hdf5 datasets, + and used in utils/env_utils.py. + """ + return dict( + env_name=self.name, + env_version=self.version, + type=self.type, + env_kwargs=deepcopy(self._init_kwargs) + ) + + @classmethod + def create_for_data_processing( + cls, + env_name, + camera_names, + camera_height, + camera_width, + reward_shaping, + **kwargs, + ): + """ + Create environment for processing datasets, which includes extracting + observations, labeling dense / sparse rewards, and annotating dones in + transitions. + + Args: + env_name (str): name of environment + camera_names (list of str): list of camera names that correspond to image observations + camera_height (int): camera height for all cameras + camera_width (int): camera width for all cameras + reward_shaping (bool): if True, use shaped environment rewards, else use sparse task completion rewards + """ + is_v1 = (robosuite.__version__.split(".")[0] == "1") + has_camera = (len(camera_names) > 0) + + new_kwargs = { + "reward_shaping": reward_shaping, + } + + if has_camera: + if is_v1: + new_kwargs["camera_names"] = list(camera_names) + new_kwargs["camera_heights"] = camera_height + new_kwargs["camera_widths"] = camera_width + else: + assert len(camera_names) == 1 + if has_camera: + new_kwargs["camera_name"] = camera_names[0] + new_kwargs["camera_height"] = camera_height + new_kwargs["camera_width"] = camera_width + + kwargs.update(new_kwargs) + + # also initialize obs utils so it knows which modalities are image modalities + image_modalities = list(camera_names) + if is_v1: + image_modalities = ["{}_image".format(cn) for cn in camera_names] + elif has_camera: + # v0.3 only had support for one image, and it was named "rgb" + assert len(image_modalities) == 1 + image_modalities = ["rgb"] + obs_modality_specs = { + "obs": { + "low_dim": [], # technically unused, so we don't have to specify all of them + "rgb": image_modalities, + } + } + ObsUtils.initialize_obs_utils_with_obs_specs(obs_modality_specs) + + # note that @postprocess_visual_obs is False since this env's images will be written to a dataset + return cls( + env_name=env_name, + render=False, + render_offscreen=has_camera, + use_image_obs=has_camera, + postprocess_visual_obs=False, + **kwargs, + ) + + @property + def rollout_exceptions(self): + """ + Return tuple of exceptions to except when doing rollouts. This is useful to ensure + that the entire training run doesn't crash because of a bad policy that causes unstable + simulation computations. + """ + return (Exception) + + def __repr__(self): + """ + Pretty-print env description. + """ + return self.name + "\n" + json.dumps(self._init_kwargs, sort_keys=True, indent=4) diff --git a/aloha-devel/robomimic/envs/wrappers.py b/aloha-devel/robomimic/envs/wrappers.py new file mode 100644 index 0000000000000000000000000000000000000000..fb46091ef33279ce9199f9d70e8add72818671a3 --- /dev/null +++ b/aloha-devel/robomimic/envs/wrappers.py @@ -0,0 +1,222 @@ +""" +A collection of useful environment wrappers. +""" +from copy import deepcopy +import textwrap +import numpy as np +from collections import deque + +import robomimic.envs.env_base as EB + + +class EnvWrapper(object): + """ + Base class for all environment wrappers in robomimic. + """ + def __init__(self, env): + """ + Args: + env (EnvBase instance): The environment to wrap. + """ + assert isinstance(env, EB.EnvBase) or isinstance(env, EnvWrapper) + self.env = env + + @classmethod + def class_name(cls): + return cls.__name__ + + def _warn_double_wrap(self): + """ + Utility function that checks if we're accidentally trying to double wrap an env + Raises: + Exception: [Double wrapping env] + """ + env = self.env + while True: + if isinstance(env, EnvWrapper): + if env.class_name() == self.class_name(): + raise Exception( + "Attempted to double wrap with Wrapper: {}".format( + self.__class__.__name__ + ) + ) + env = env.env + else: + break + + @property + def unwrapped(self): + """ + Grabs unwrapped environment + + Returns: + env (EnvBase instance): Unwrapped environment + """ + if hasattr(self.env, "unwrapped"): + return self.env.unwrapped + else: + return self.env + + def _to_string(self): + """ + Subclasses should override this method to print out info about the + wrapper (such as arguments passed to it). + """ + return '' + + def __repr__(self): + """Pretty print environment.""" + header = '{}'.format(str(self.__class__.__name__)) + msg = '' + indent = ' ' * 4 + if self._to_string() != '': + msg += textwrap.indent("\n" + self._to_string(), indent) + msg += textwrap.indent("\nenv={}".format(self.env), indent) + msg = header + '(' + msg + '\n)' + return msg + + # this method is a fallback option on any methods the original env might support + def __getattr__(self, attr): + # using getattr ensures that both __getattribute__ and __getattr__ (fallback) get called + # (see https://stackoverflow.com/questions/3278077/difference-between-getattr-vs-getattribute) + orig_attr = getattr(self.env, attr) + if callable(orig_attr): + + def hooked(*args, **kwargs): + result = orig_attr(*args, **kwargs) + # prevent wrapped_class from becoming unwrapped + if id(result) == id(self.env): + return self + return result + + return hooked + else: + return orig_attr + + +class FrameStackWrapper(EnvWrapper): + """ + Wrapper for frame stacking observations during rollouts. The agent + receives a sequence of past observations instead of a single observation + when it calls @env.reset, @env.reset_to, or @env.step in the rollout loop. + """ + def __init__(self, env, num_frames): + """ + Args: + env (EnvBase instance): The environment to wrap. + num_frames (int): number of past observations (including current observation) + to stack together. Must be greater than 1 (otherwise this wrapper would + be a no-op). + """ + assert num_frames > 1, "error: FrameStackWrapper must have num_frames > 1 but got num_frames of {}".format(num_frames) + + super(FrameStackWrapper, self).__init__(env=env) + self.num_frames = num_frames + + ### TODO: add action padding option + adding action to obs to include action history in obs ### + + # keep track of last @num_frames observations for each obs key + self.obs_history = None + + def _get_initial_obs_history(self, init_obs): + """ + Helper method to get observation history from the initial observation, by + repeating it. + + Returns: + obs_history (dict): a deque for each observation key, with an extra + leading dimension of 1 for each key (for easy concatenation later) + """ + obs_history = {} + for k in init_obs: + obs_history[k] = deque( + [init_obs[k][None] for _ in range(self.num_frames)], + maxlen=self.num_frames, + ) + return obs_history + + def _get_stacked_obs_from_history(self): + """ + Helper method to convert internal variable @self.obs_history to a + stacked observation where each key is a numpy array with leading dimension + @self.num_frames. + """ + # concatenate all frames per key so we return a numpy array per key + return { k : np.concatenate(self.obs_history[k], axis=0) for k in self.obs_history } + + def cache_obs_history(self): + self.obs_history_cache = deepcopy(self.obs_history) + + def uncache_obs_history(self): + self.obs_history = self.obs_history_cache + self.obs_history_cache = None + + def reset(self): + """ + Modify to return frame stacked observation which is @self.num_frames copies of + the initial observation. + + Returns: + obs_stacked (dict): each observation key in original observation now has + leading shape @self.num_frames and consists of the previous @self.num_frames + observations + """ + obs = self.env.reset() + self.timestep = 0 # always zero regardless of timestep type + self.update_obs(obs, reset=True) + self.obs_history = self._get_initial_obs_history(init_obs=obs) + return self._get_stacked_obs_from_history() + + def reset_to(self, state): + """ + Modify to return frame stacked observation which is @self.num_frames copies of + the initial observation. + + Returns: + obs_stacked (dict): each observation key in original observation now has + leading shape @self.num_frames and consists of the previous @self.num_frames + observations + """ + obs = self.env.reset_to(state) + self.timestep = 0 # always zero regardless of timestep type + self.update_obs(obs, reset=True) + self.obs_history = self._get_initial_obs_history(init_obs=obs) + return self._get_stacked_obs_from_history() + + def step(self, action): + """ + Modify to update the internal frame history and return frame stacked observation, + which will have leading dimension @self.num_frames for each key. + + Args: + action (np.array): action to take + + Returns: + obs_stacked (dict): each observation key in original observation now has + leading shape @self.num_frames and consists of the previous @self.num_frames + observations + reward (float): reward for this step + done (bool): whether the task is done + info (dict): extra information + """ + obs, r, done, info = self.env.step(action) + self.update_obs(obs, action=action, reset=False) + # update frame history + for k in obs: + # make sure to have leading dim of 1 for easy concatenation + self.obs_history[k].append(obs[k][None]) + obs_ret = self._get_stacked_obs_from_history() + return obs_ret, r, done, info + + def update_obs(self, obs, action=None, reset=False): + obs["timesteps"] = np.array([self.timestep]) + + if reset: + obs["actions"] = np.zeros(self.env.action_dimension) + else: + self.timestep += 1 + obs["actions"] = action[: self.env.action_dimension] + + def _to_string(self): + """Info to pretty print.""" + return "num_frames={}".format(self.num_frames) \ No newline at end of file diff --git a/aloha-devel/robomimic/exps/templates/act.json b/aloha-devel/robomimic/exps/templates/act.json new file mode 100644 index 0000000000000000000000000000000000000000..4512ecdfa0178c74f5621326889f85ca9559c9da --- /dev/null +++ b/aloha-devel/robomimic/exps/templates/act.json @@ -0,0 +1,160 @@ +{ + "algo_name": "act", + "experiment": { + "name": "test", + "validate": false, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "mse":{}, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 40, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 500, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": true, + "n": 50, + "horizon": 400, + "rate": 40, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir":"../act_trained_models", + "num_data_workers": 4, + "hdf5_cache_mode": "low_dim", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": false, + "hdf5_normalize_obs": false, + "hdf5_filter_key": null, + "seq_length": 10, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 128, + "num_epochs": 10000, + "seed": 1 + }, + "algo": { + "optim_params": { + "policy": { + "optimizer_type": "adamw", + "learning_rate": { + "initial": 0.00005, + "decay_factor": 1, + "epoch_schedule": [ + 100 + ], + "scheduler_type": "linear" + }, + "regularization": { + "L2": 0.0001 + } + } + }, + "loss": { + "l2_weight": 0.0, + "l1_weight": 1.0, + "cos_weight": 0.0 + }, + "act": { + "hidden_dim": 512, + "dim_feedforward": 3200, + "backbone": "resnet18", + "enc_layers": 4, + "dec_layers": 7, + "nheads": 8, + "latent_dim": 32, + "kl_weight": 20 + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": { + "feature_dimension": 64, + "backbone_class": "ResNet18Conv", + "backbone_kwargs": { + "pretrained": false, + "input_coord_conv": false + }, + "pool_class": "SpatialSoftmax", + "pool_kwargs": { + "num_kp": 32, + "learnable_temperature": false, + "temperature": 1.0, + "noise_std": 0.0 + } + }, + "obs_randomizer_class": "CropRandomizer", + "obs_randomizer_kwargs": { + "crop_height": 76, + "crop_width": 76, + "num_crops": 1, + "pos_enc": false + } + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + } +} \ No newline at end of file diff --git a/aloha-devel/robomimic/exps/templates/bc.json b/aloha-devel/robomimic/exps/templates/bc.json new file mode 100644 index 0000000000000000000000000000000000000000..a2c181057e32b1472e03d2a479ea48d1d1d234b3 --- /dev/null +++ b/aloha-devel/robomimic/exps/templates/bc.json @@ -0,0 +1,216 @@ +{ + "algo_name": "bc", + "experiment": { + "name": "test", + "validate": false, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "mse":{}, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": true, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../bc_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": false, + "hdf5_normalize_obs": false, + "hdf5_filter_key": null, + "hdf5_validation_filter_key": null, + "seq_length": 1, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "optim_params": { + "policy": { + "optimizer_type": "adam", + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.1, + "epoch_schedule": [], + "scheduler_type": "multistep" + }, + "regularization": { + "L2": 0.0 + } + } + }, + "loss": { + "l2_weight": 1.0, + "l1_weight": 0.0, + "cos_weight": 0.0 + }, + "actor_layer_dims": [ + 1024, + 1024 + ], + "gaussian": { + "enabled": false, + "fixed_std": false, + "init_std": 0.1, + "min_std": 0.01, + "std_activation": "softplus", + "low_noise_eval": true + }, + "gmm": { + "enabled": false, + "num_modes": 5, + "min_std": 0.0001, + "std_activation": "softplus", + "low_noise_eval": true + }, + "vae": { + "enabled": false, + "latent_dim": 14, + "latent_clip": null, + "kl_weight": 1.0, + "decoder": { + "is_conditioned": true, + "reconstruction_sum_across_elements": false + }, + "prior": { + "learn": false, + "is_conditioned": false, + "use_gmm": false, + "gmm_num_modes": 10, + "gmm_learn_weights": false, + "use_categorical": false, + "categorical_dim": 10, + "categorical_gumbel_softmax_hard": false, + "categorical_init_temp": 1.0, + "categorical_temp_anneal_step": 0.001, + "categorical_min_temp": 0.3 + }, + "encoder_layer_dims": [ + 300, + 400 + ], + "decoder_layer_dims": [ + 300, + 400 + ], + "prior_layer_dims": [ + 300, + 400 + ] + }, + "rnn": { + "enabled": false, + "horizon": 10, + "hidden_dim": 400, + "rnn_type": "LSTM", + "num_layers": 2, + "open_loop": false, + "kwargs": { + "bidirectional": false + } + }, + "transformer": { + "enabled": false, + "context_length": 10, + "embed_dim": 512, + "num_layers": 6, + "num_heads": 8, + "emb_dropout": 0.1, + "attn_dropout": 0.1, + "block_output_dropout": 0.1, + "sinusoidal_embedding": false, + "activation": "gelu", + "supervise_all_steps": false, + "nn_parameter_for_timesteps": true + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} \ No newline at end of file diff --git a/aloha-devel/robomimic/exps/templates/bc_transformer.json b/aloha-devel/robomimic/exps/templates/bc_transformer.json new file mode 100644 index 0000000000000000000000000000000000000000..71461ffa6da955a452a4df20c61d4e896fddef43 --- /dev/null +++ b/aloha-devel/robomimic/exps/templates/bc_transformer.json @@ -0,0 +1,172 @@ +{ + "algo_name": "bc", + "experiment": { + "name": "test", + "validate": true, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "mse":{}, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": true, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../bc_transformer_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "low_dim", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": false, + "hdf5_normalize_obs": false, + "hdf5_filter_key": null, + "seq_length": 1, + "pad_seq_length": true, + "frame_stack": 10, + "pad_frame_stack": true, + "dataset_keys": [ + "actions" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "optim_params": { + "policy": { + "optimizer_type": "adamw", + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.1, + "epoch_schedule": [100], + "scheduler_type": "linear" + }, + "regularization": { + "L2": 0.01 + } + } + }, + "loss": { + "l2_weight": 1.0, + "l1_weight": 0.0, + "cos_weight": 0.0 + }, + "actor_layer_dims": [], + "gaussian": { + "enabled": false + }, + "gmm": { + "enabled": true, + "num_modes": 5, + "min_std": 0.0001, + "std_activation": "softplus", + "low_noise_eval": true + }, + "vae": { + "enabled": false + }, + "rnn": { + "enabled": false + }, + "transformer": { + "enabled": true, + "supervise_all_steps": false, + "num_layers": 6, + "embed_dim": 512, + "num_heads": 8 + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": { + "feature_dimension": 64, + "backbone_class": "ResNet18Conv", + "backbone_kwargs": { + "pretrained": false, + "input_coord_conv": false + }, + "pool_class": "SpatialSoftmax", + "pool_kwargs": { + "num_kp": 32, + "learnable_temperature": false, + "temperature": 1.0, + "noise_std": 0.0 + } + }, + "obs_randomizer_class": "CropRandomizer", + "obs_randomizer_kwargs": { + "crop_height": 76, + "crop_width": 76, + "num_crops": 1, + "pos_enc": false + } + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + } +} \ No newline at end of file diff --git a/aloha-devel/robomimic/exps/templates/cql.json b/aloha-devel/robomimic/exps/templates/cql.json new file mode 100644 index 0000000000000000000000000000000000000000..a920efd6f01844971fba4881d73b762f7cf47ade --- /dev/null +++ b/aloha-devel/robomimic/exps/templates/cql.json @@ -0,0 +1,182 @@ +{ + "algo_name": "cql", + "experiment": { + "name": "test", + "validate": false, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": true, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../cql_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": true, + "hdf5_normalize_obs": false, + "hdf5_filter_key": null, + "hdf5_validation_filter_key": null, + "seq_length": 1, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 1024, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "optim_params": { + "critic": { + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.0, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + }, + "actor": { + "learning_rate": { + "initial": 0.0003, + "decay_factor": 0.0, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + } + }, + "discount": 0.99, + "n_step": 1, + "target_tau": 0.005, + "actor": { + "bc_start_steps": 0, + "target_entropy": "default", + "max_gradient_norm": null, + "net": { + "type": "gaussian", + "common": { + "std_activation": "exp", + "use_tanh": true, + "low_noise_eval": true + }, + "gaussian": { + "init_last_fc_weight": 0.001, + "init_std": 0.3, + "fixed_std": false + } + }, + "layer_dims": [ + 300, + 400 + ] + }, + "critic": { + "use_huber": false, + "max_gradient_norm": null, + "value_bounds": null, + "num_action_samples": 1, + "cql_weight": 1.0, + "deterministic_backup": true, + "min_q_weight": 1.0, + "target_q_gap": 5.0, + "num_random_actions": 10, + "ensemble": { + "n": 2 + }, + "layer_dims": [ + 300, + 400 + ] + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} \ No newline at end of file diff --git a/aloha-devel/robomimic/exps/templates/gl.json b/aloha-devel/robomimic/exps/templates/gl.json new file mode 100644 index 0000000000000000000000000000000000000000..39b4c2dbd65dad06afaaa1f88bd605a3477e3312 --- /dev/null +++ b/aloha-devel/robomimic/exps/templates/gl.json @@ -0,0 +1,182 @@ +{ + "algo_name": "gl", + "experiment": { + "name": "test", + "validate": false, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": true, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../gl_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": true, + "hdf5_normalize_obs": false, + "hdf5_filter_key": null, + "hdf5_validation_filter_key": null, + "seq_length": 1, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "optim_params": { + "goal_network": { + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + } + }, + "subgoal_horizon": 10, + "ae": { + "planner_layer_dims": [ + 300, + 400 + ] + }, + "vae": { + "enabled": true, + "latent_dim": 16, + "latent_clip": null, + "kl_weight": 1.0, + "decoder": { + "is_conditioned": true, + "reconstruction_sum_across_elements": false + }, + "prior": { + "learn": false, + "is_conditioned": false, + "use_gmm": false, + "gmm_num_modes": 10, + "gmm_learn_weights": false, + "use_categorical": false, + "categorical_dim": 10, + "categorical_gumbel_softmax_hard": false, + "categorical_init_temp": 1.0, + "categorical_temp_anneal_step": 0.001, + "categorical_min_temp": 0.3 + }, + "encoder_layer_dims": [ + 300, + 400 + ], + "decoder_layer_dims": [ + 300, + 400 + ], + "prior_layer_dims": [ + 300, + 400 + ] + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + }, + "subgoal": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} \ No newline at end of file diff --git a/aloha-devel/robomimic/exps/templates/hbc.json b/aloha-devel/robomimic/exps/templates/hbc.json new file mode 100644 index 0000000000000000000000000000000000000000..26eff76a8f40e3fd787c7a561a91155369101b7e --- /dev/null +++ b/aloha-devel/robomimic/exps/templates/hbc.json @@ -0,0 +1,293 @@ +{ + "algo_name": "hbc", + "experiment": { + "name": "test", + "validate": false, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": true, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../hbc_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": true, + "hdf5_normalize_obs": false, + "hdf5_filter_key": null, + "hdf5_validation_filter_key": null, + "seq_length": 10, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "mode": "separate", + "actor_use_random_subgoals": false, + "subgoal_update_interval": 10, + "latent_subgoal": { + "enabled": false, + "prior_correction": { + "enabled": false, + "num_samples": 100 + } + }, + "planner": { + "optim_params": { + "goal_network": { + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + } + }, + "subgoal_horizon": 10, + "ae": { + "planner_layer_dims": [ + 300, + 400 + ] + }, + "vae": { + "enabled": true, + "latent_dim": 16, + "latent_clip": null, + "kl_weight": 1.0, + "decoder": { + "is_conditioned": true, + "reconstruction_sum_across_elements": false + }, + "prior": { + "learn": false, + "is_conditioned": false, + "use_gmm": false, + "gmm_num_modes": 10, + "gmm_learn_weights": false, + "use_categorical": false, + "categorical_dim": 10, + "categorical_gumbel_softmax_hard": false, + "categorical_init_temp": 1.0, + "categorical_temp_anneal_step": 0.001, + "categorical_min_temp": 0.3 + }, + "encoder_layer_dims": [ + 300, + 400 + ], + "decoder_layer_dims": [ + 300, + 400 + ], + "prior_layer_dims": [ + 300, + 400 + ] + } + }, + "actor": { + "optim_params": { + "policy": { + "optimizer_type": "adam", + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.1, + "epoch_schedule": [], + "scheduler_type": "multistep" + }, + "regularization": { + "L2": 0.0 + } + } + }, + "loss": { + "l2_weight": 1.0, + "l1_weight": 0.0, + "cos_weight": 0.0 + }, + "actor_layer_dims": [ + 1024, + 1024 + ], + "rnn": { + "enabled": true, + "horizon": 10, + "hidden_dim": 400, + "rnn_type": "LSTM", + "num_layers": 2, + "open_loop": false, + "kwargs": { + "bidirectional": false + } + }, + "transformer": { + "enabled": false, + "context_length": 10, + "embed_dim": 512, + "num_layers": 6, + "num_heads": 8, + "emb_dropout": 0.1, + "attn_dropout": 0.1, + "block_output_dropout": 0.1, + "sinusoidal_embedding": false, + "activation": "gelu", + "supervise_all_steps": false, + "nn_parameter_for_timesteps": true + } + } + }, + "observation": { + "planner": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + }, + "subgoal": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "actor": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} \ No newline at end of file diff --git a/aloha-devel/robomimic/exps/templates/iql.json b/aloha-devel/robomimic/exps/templates/iql.json new file mode 100644 index 0000000000000000000000000000000000000000..4731788417924c649f1b92627fe6bf7f14668aac --- /dev/null +++ b/aloha-devel/robomimic/exps/templates/iql.json @@ -0,0 +1,192 @@ +{ + "algo_name": "iql", + "experiment": { + "name": "test", + "validate": false, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": true, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../iql_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": true, + "hdf5_normalize_obs": false, + "hdf5_filter_key": null, + "hdf5_validation_filter_key": null, + "seq_length": 1, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "optim_params": { + "critic": { + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.0, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + }, + "vf": { + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.0, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + }, + "actor": { + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.0, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + } + }, + "discount": 0.99, + "target_tau": 0.01, + "actor": { + "net": { + "type": "gaussian", + "common": { + "std_activation": "softplus", + "low_noise_eval": true, + "use_tanh": false + }, + "gaussian": { + "init_last_fc_weight": 0.001, + "init_std": 0.3, + "fixed_std": false + }, + "gmm": { + "num_modes": 5, + "min_std": 0.0001 + } + }, + "layer_dims": [ + 300, + 400 + ], + "max_gradient_norm": null + }, + "critic": { + "ensemble": { + "n": 2 + }, + "layer_dims": [ + 300, + 400 + ], + "use_huber": false, + "max_gradient_norm": null + }, + "adv": { + "clip_adv_value": null, + "beta": 1.0, + "use_final_clip": true + }, + "vf_quantile": 0.9 + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} \ No newline at end of file diff --git a/aloha-devel/robomimic/exps/templates/iris.json b/aloha-devel/robomimic/exps/templates/iris.json new file mode 100644 index 0000000000000000000000000000000000000000..6551663864a4d57d05d263de0069269ab115d8de --- /dev/null +++ b/aloha-devel/robomimic/exps/templates/iris.json @@ -0,0 +1,465 @@ +{ + "algo_name": "iris", + "experiment": { + "name": "test", + "validate": false, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 50, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": true, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../iris_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": true, + "hdf5_normalize_obs": false, + "hdf5_filter_key": null, + "hdf5_validation_filter_key": null, + "seq_length": 10, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 2000, + "seed": 1 + }, + "algo": { + "mode": "separate", + "actor_use_random_subgoals": false, + "subgoal_update_interval": 10, + "latent_subgoal": { + "enabled": false, + "prior_correction": { + "enabled": false, + "num_samples": 100 + } + }, + "value_planner": { + "planner": { + "optim_params": { + "goal_network": { + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + } + } + }, + "subgoal_horizon": 10, + "ae": { + "planner_layer_dims": [ + 300, + 400 + ] + }, + "vae": { + "enabled": true, + "latent_dim": 16, + "latent_clip": null, + "kl_weight": 1.0, + "decoder": { + "is_conditioned": true, + "reconstruction_sum_across_elements": false + }, + "prior": { + "learn": false, + "is_conditioned": false, + "use_gmm": false, + "gmm_num_modes": 10, + "gmm_learn_weights": false, + "use_categorical": false, + "categorical_dim": 10, + "categorical_gumbel_softmax_hard": false, + "categorical_init_temp": 1.0, + "categorical_temp_anneal_step": 0.001, + "categorical_min_temp": 0.3 + }, + "encoder_layer_dims": [ + 300, + 400 + ], + "decoder_layer_dims": [ + 300, + 400 + ], + "prior_layer_dims": [ + 300, + 400 + ] + } + }, + "value": { + "optim_params": { + "critic": { + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + }, + "action_sampler": { + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + }, + "actor": { + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + } + }, + "discount": 0.99, + "n_step": 1, + "target_tau": 0.005, + "infinite_horizon": false, + "critic": { + "use_huber": false, + "max_gradient_norm": null, + "value_bounds": null, + "num_action_samples": 10, + "num_action_samples_rollout": 100, + "ensemble": { + "n": 2, + "weight": 0.75 + }, + "distributional": { + "enabled": false, + "num_atoms": 51 + }, + "layer_dims": [ + 300, + 400 + ] + }, + "action_sampler": { + "actor_layer_dims": [ + 1024, + 1024 + ], + "gmm": { + "enabled": false, + "num_modes": 5, + "min_std": 0.0001, + "std_activation": "softplus", + "low_noise_eval": true + }, + "vae": { + "enabled": true, + "latent_dim": 14, + "latent_clip": null, + "kl_weight": 1.0, + "decoder": { + "is_conditioned": true, + "reconstruction_sum_across_elements": false + }, + "prior": { + "learn": false, + "is_conditioned": false, + "use_gmm": false, + "gmm_num_modes": 10, + "gmm_learn_weights": false, + "use_categorical": false, + "categorical_dim": 10, + "categorical_gumbel_softmax_hard": false, + "categorical_init_temp": 1.0, + "categorical_temp_anneal_step": 0.001, + "categorical_min_temp": 0.3 + }, + "encoder_layer_dims": [ + 300, + 400 + ], + "decoder_layer_dims": [ + 300, + 400 + ], + "prior_layer_dims": [ + 300, + 400 + ] + }, + "freeze_encoder_epoch": -1 + }, + "actor": { + "enabled": false, + "perturbation_scale": 0.05, + "layer_dims": [ + 300, + 400 + ] + } + }, + "num_samples": 100 + }, + "actor": { + "optim_params": { + "policy": { + "optimizer_type": "adam", + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.1, + "epoch_schedule": [], + "scheduler_type": "multistep" + }, + "regularization": { + "L2": 0.0 + } + } + }, + "loss": { + "l2_weight": 1.0, + "l1_weight": 0.0, + "cos_weight": 0.0 + }, + "actor_layer_dims": [ + 1024, + 1024 + ], + "rnn": { + "enabled": true, + "horizon": 10, + "hidden_dim": 400, + "rnn_type": "LSTM", + "num_layers": 2, + "open_loop": false, + "kwargs": { + "bidirectional": false + } + }, + "transformer": { + "enabled": false, + "context_length": 10, + "embed_dim": 512, + "num_layers": 6, + "num_heads": 8, + "emb_dropout": 0.1, + "attn_dropout": 0.1, + "block_output_dropout": 0.1, + "sinusoidal_embedding": false, + "activation": "gelu", + "supervise_all_steps": false, + "nn_parameter_for_timesteps": true + } + } + }, + "observation": { + "value_planner": { + "planner": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + }, + "subgoal": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "value": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + } + }, + "actor": { + "modalities": { + "obs": { + "low_dim": [ + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} \ No newline at end of file diff --git a/aloha-devel/robomimic/exps/templates/td3_bc.json b/aloha-devel/robomimic/exps/templates/td3_bc.json new file mode 100644 index 0000000000000000000000000000000000000000..414a8f04f0cce7c9857207b1b1269ff10c3ee38b --- /dev/null +++ b/aloha-devel/robomimic/exps/templates/td3_bc.json @@ -0,0 +1,167 @@ +{ + "algo_name": "td3_bc", + "experiment": { + "name": "test", + "validate": false, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": false, + "wandb_proj_name": "debug" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 20, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": true, + "on_best_rollout_success_rate": false + }, + "epoch_every_n_steps": 5000, + "validation_epoch_every_n_steps": 10, + "env": null, + "additional_envs": null, + "render": false, + "render_video": false, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": true, + "n": 50, + "horizon": 1000, + "rate": 1, + "warmstart": 0, + "terminate_on_success": true + } + }, + "train": { + "data": null, + "output_dir": "../td3_bc_trained_models", + "num_data_workers": 0, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": true, + "hdf5_normalize_obs": true, + "hdf5_filter_key": null, + "hdf5_validation_filter_key": null, + "seq_length": 1, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 256, + "num_epochs": 200, + "seed": 1 + }, + "algo": { + "optim_params": { + "critic": { + "learning_rate": { + "initial": 0.0003, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + }, + "actor": { + "learning_rate": { + "initial": 0.0003, + "decay_factor": 0.1, + "epoch_schedule": [] + }, + "regularization": { + "L2": 0.0 + }, + "start_epoch": -1, + "end_epoch": -1 + } + }, + "alpha": 2.5, + "discount": 0.99, + "n_step": 1, + "target_tau": 0.005, + "infinite_horizon": false, + "critic": { + "use_huber": false, + "max_gradient_norm": null, + "value_bounds": null, + "ensemble": { + "n": 2, + "weight": 1.0 + }, + "layer_dims": [ + 256, + 256 + ] + }, + "actor": { + "update_freq": 2, + "noise_std": 0.2, + "noise_clip": 0.5, + "layer_dims": [ + 256, + 256 + ] + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "flat" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} \ No newline at end of file diff --git a/aloha-devel/robomimic/macros.py b/aloha-devel/robomimic/macros.py new file mode 100644 index 0000000000000000000000000000000000000000..3b6c05038ca2021a3981908cea3888601a21351c --- /dev/null +++ b/aloha-devel/robomimic/macros.py @@ -0,0 +1,27 @@ +""" +Set of global variables shared across robomimic +""" +# Sets debugging mode. Should be set at top-level script so that internal +# debugging functionalities are made active +DEBUG = False + +# Whether to visualize the before & after of an observation randomizer +VISUALIZE_RANDOMIZER = False + +# wandb entity (eg. username or team name) +WANDB_ENTITY = None + +# wandb api key (obtain from https://wandb.ai/authorize) +# alternatively, set up wandb from terminal with `wandb login` +WANDB_API_KEY = None + +try: + from robomimic.macros_private import * +except ImportError: + from robomimic.utils.log_utils import log_warning + import robomimic + log_warning( + "No private macro file found!"\ + "\nIt is recommended to use a private macro file"\ + "\nTo setup, run: python {}/scripts/setup_macros.py".format(robomimic.__path__[0]) + ) diff --git a/aloha-devel/robomimic/models/__pycache__/__init__.cpython-38.pyc b/aloha-devel/robomimic/models/__pycache__/__init__.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bea7bc19bae5e0a6fdfaddcd4c692c202e0c89f9 Binary files /dev/null and b/aloha-devel/robomimic/models/__pycache__/__init__.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/models/__pycache__/base_nets.cpython-38.pyc b/aloha-devel/robomimic/models/__pycache__/base_nets.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0789e433cf8c2869f032b852db26e6c7809b4ec2 Binary files /dev/null and b/aloha-devel/robomimic/models/__pycache__/base_nets.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/models/__pycache__/distributions.cpython-38.pyc b/aloha-devel/robomimic/models/__pycache__/distributions.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..91c8d7f03933270d7243a4724ea9d5a8465067ed Binary files /dev/null and b/aloha-devel/robomimic/models/__pycache__/distributions.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/models/__pycache__/policy_nets.cpython-38.pyc b/aloha-devel/robomimic/models/__pycache__/policy_nets.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5991f363f95e99a5ad1cd17f79a8cf0239481a79 Binary files /dev/null and b/aloha-devel/robomimic/models/__pycache__/policy_nets.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/models/__pycache__/transformers.cpython-38.pyc b/aloha-devel/robomimic/models/__pycache__/transformers.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1b0cfa47d39e0d251f006c567a5e3d84ed71d04c Binary files /dev/null and b/aloha-devel/robomimic/models/__pycache__/transformers.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/models/__pycache__/value_nets.cpython-38.pyc b/aloha-devel/robomimic/models/__pycache__/value_nets.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..809334f439940de389bfa0a5282f548096109c7c Binary files /dev/null and b/aloha-devel/robomimic/models/__pycache__/value_nets.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/models/base_nets.py b/aloha-devel/robomimic/models/base_nets.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5d9b5c110c80e35c1ab128b9b8169c7f12ad89 --- /dev/null +++ b/aloha-devel/robomimic/models/base_nets.py @@ -0,0 +1,1156 @@ +""" +Contains torch Modules that correspond to basic network building blocks, like +MLP, RNN, and CNN backbones. +""" + +import math +import abc +import numpy as np +import textwrap +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision import models as vision_models +from torchvision import transforms + +import robomimic.utils.tensor_utils as TensorUtils + +CONV_ACTIVATIONS = { + "relu": nn.ReLU, + "None": None, + None: None, +} + + +def rnn_args_from_config(rnn_config): + """ + Takes a Config object corresponding to RNN settings + (for example `config.algo.rnn` in BCConfig) and extracts + rnn kwargs for instantiating rnn networks. + """ + return dict( + rnn_hidden_dim=rnn_config.hidden_dim, + rnn_num_layers=rnn_config.num_layers, + rnn_type=rnn_config.rnn_type, + rnn_kwargs=dict(rnn_config.kwargs), + ) + + +def transformer_args_from_config(transformer_config): + """ + Takes a Config object corresponding to Transformer settings + (for example `config.algo.transformer` in BCConfig) and extracts + transformer kwargs for instantiating transformer networks. + """ + transformer_args = dict( + transformer_context_length=transformer_config.context_length, + transformer_embed_dim=transformer_config.embed_dim, + transformer_num_heads=transformer_config.num_heads, + transformer_emb_dropout=transformer_config.emb_dropout, + transformer_attn_dropout=transformer_config.attn_dropout, + transformer_block_output_dropout=transformer_config.block_output_dropout, + transformer_sinusoidal_embedding=transformer_config.sinusoidal_embedding, + transformer_activation=transformer_config.activation, + transformer_nn_parameter_for_timesteps=transformer_config.nn_parameter_for_timesteps, + ) + + if "num_layers" in transformer_config: + transformer_args["transformer_num_layers"] = transformer_config.num_layers + + return transformer_args + + +class Module(torch.nn.Module): + """ + Base class for networks. The only difference from torch.nn.Module is that it + requires implementing @output_shape. + """ + @abc.abstractmethod + def output_shape(self, input_shape=None): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + raise NotImplementedError + + +class Sequential(torch.nn.Sequential, Module): + """ + Compose multiple Modules together (defined above). + """ + def __init__(self, *args): + for arg in args: + assert isinstance(arg, Module) + torch.nn.Sequential.__init__(self, *args) + self.fixed = False + + def output_shape(self, input_shape=None): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + out_shape = input_shape + for module in self: + out_shape = module.output_shape(out_shape) + return out_shape + + def freeze(self): + self.fixed = True + + def train(self, mode): + if self.fixed: + super().train(False) + else: + super().train(mode) + + +class Parameter(Module): + """ + A class that is a thin wrapper around a torch.nn.Parameter to make for easy saving + and optimization. + """ + def __init__(self, init_tensor): + """ + Args: + init_tensor (torch.Tensor): initial tensor + """ + super(Parameter, self).__init__() + self.param = torch.nn.Parameter(init_tensor) + + def output_shape(self, input_shape=None): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + return list(self.param.shape) + + def forward(self, inputs=None): + """ + Forward call just returns the parameter tensor. + """ + return self.param + + +class Unsqueeze(Module): + """ + Trivial class that unsqueezes the input. Useful for including in a nn.Sequential network + """ + def __init__(self, dim): + super(Unsqueeze, self).__init__() + self.dim = dim + + def output_shape(self, input_shape=None): + assert input_shape is not None + return input_shape + [1] if self.dim == -1 else input_shape[:self.dim + 1] + [1] + input_shape[self.dim + 1:] + + def forward(self, x): + return x.unsqueeze(dim=self.dim) + + +class Squeeze(Module): + """ + Trivial class that squeezes the input. Useful for including in a nn.Sequential network + """ + + def __init__(self, dim): + super(Squeeze, self).__init__() + self.dim = dim + + def output_shape(self, input_shape=None): + assert input_shape is not None + return input_shape[:self.dim] + input_shape[self.dim+1:] if input_shape[self.dim] == 1 else input_shape + + def forward(self, x): + return x.squeeze(dim=self.dim) + + +class MLP(Module): + """ + Base class for simple Multi-Layer Perceptrons. + """ + def __init__( + self, + input_dim, + output_dim, + layer_dims=(), + layer_func=nn.Linear, + layer_func_kwargs=None, + activation=nn.ReLU, + dropouts=None, + normalization=False, + output_activation=None, + ): + """ + Args: + input_dim (int): dimension of inputs + + output_dim (int): dimension of outputs + + layer_dims ([int]): sequence of integers for the hidden layers sizes + + layer_func: mapping per layer - defaults to Linear + + layer_func_kwargs (dict): kwargs for @layer_func + + activation: non-linearity per layer - defaults to ReLU + + dropouts ([float]): if not None, adds dropout layers with the corresponding probabilities + after every layer. Must be same size as @layer_dims. + + normalization (bool): if True, apply layer normalization after each layer + + output_activation: if provided, applies the provided non-linearity to the output layer + """ + super(MLP, self).__init__() + layers = [] + dim = input_dim + if layer_func_kwargs is None: + layer_func_kwargs = dict() + if dropouts is not None: + assert(len(dropouts) == len(layer_dims)) + for i, l in enumerate(layer_dims): + layers.append(layer_func(dim, l, **layer_func_kwargs)) + if normalization: + layers.append(nn.LayerNorm(l)) + layers.append(activation()) + if dropouts is not None and dropouts[i] > 0.: + layers.append(nn.Dropout(dropouts[i])) + dim = l + layers.append(layer_func(dim, output_dim)) + if output_activation is not None: + layers.append(output_activation()) + self._layer_func = layer_func + self.nets = layers + self._model = nn.Sequential(*layers) + + self._layer_dims = layer_dims + self._input_dim = input_dim + self._output_dim = output_dim + self._dropouts = dropouts + self._act = activation + self._output_act = output_activation + + def output_shape(self, input_shape=None): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + return [self._output_dim] + + def forward(self, inputs): + """ + Forward pass. + """ + return self._model(inputs) + + def __repr__(self): + """Pretty print network.""" + header = str(self.__class__.__name__) + act = None if self._act is None else self._act.__name__ + output_act = None if self._output_act is None else self._output_act.__name__ + + indent = ' ' * 4 + msg = "input_dim={}\noutput_dim={}\nlayer_dims={}\nlayer_func={}\ndropout={}\nact={}\noutput_act={}".format( + self._input_dim, self._output_dim, self._layer_dims, + self._layer_func.__name__, self._dropouts, act, output_act + ) + msg = textwrap.indent(msg, indent) + msg = header + '(\n' + msg + '\n)' + return msg + + +class RNN_Base(Module): + """ + A wrapper class for a multi-step RNN and a per-step network. + """ + def __init__( + self, + input_dim, + rnn_hidden_dim, + rnn_num_layers, + rnn_type="LSTM", # [LSTM, GRU] + rnn_kwargs=None, + per_step_net=None, + ): + """ + Args: + input_dim (int): dimension of inputs + + rnn_hidden_dim (int): RNN hidden dimension + + rnn_num_layers (int): number of RNN layers + + rnn_type (str): [LSTM, GRU] + + rnn_kwargs (dict): kwargs for the torch.nn.LSTM / GRU + + per_step_net: a network that runs per time step on top of the RNN output + """ + super(RNN_Base, self).__init__() + self.per_step_net = per_step_net + if per_step_net is not None: + assert isinstance(per_step_net, Module), "RNN_Base: per_step_net is not instance of Module" + + assert rnn_type in ["LSTM", "GRU"] + rnn_cls = nn.LSTM if rnn_type == "LSTM" else nn.GRU + rnn_kwargs = rnn_kwargs if rnn_kwargs is not None else {} + rnn_is_bidirectional = rnn_kwargs.get("bidirectional", False) + + self.nets = rnn_cls( + input_size=input_dim, + hidden_size=rnn_hidden_dim, + num_layers=rnn_num_layers, + batch_first=True, + **rnn_kwargs, + ) + + self._hidden_dim = rnn_hidden_dim + self._num_layers = rnn_num_layers + self._rnn_type = rnn_type + self._num_directions = int(rnn_is_bidirectional) + 1 # 2 if bidirectional, 1 otherwise + + @property + def rnn_type(self): + return self._rnn_type + + def get_rnn_init_state(self, batch_size, device): + """ + Get a default RNN state (zeros) + Args: + batch_size (int): batch size dimension + + device: device the hidden state should be sent to. + + Returns: + hidden_state (torch.Tensor or tuple): returns hidden state tensor or tuple of hidden state tensors + depending on the RNN type + """ + h_0 = torch.zeros(self._num_layers * self._num_directions, batch_size, self._hidden_dim).to(device) + if self._rnn_type == "LSTM": + c_0 = torch.zeros(self._num_layers * self._num_directions, batch_size, self._hidden_dim).to(device) + return h_0, c_0 + else: + return h_0 + + def output_shape(self, input_shape): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + + # infer time dimension from input shape and add to per_step_net output shape + if self.per_step_net is not None: + out = self.per_step_net.output_shape(input_shape[1:]) + if isinstance(out, dict): + out = {k: [input_shape[0]] + out[k] for k in out} + else: + out = [input_shape[0]] + out + else: + out = [input_shape[0], self._num_layers * self._hidden_dim] + return out + + def forward(self, inputs, rnn_init_state=None, return_state=False): + """ + Forward a sequence of inputs through the RNN and the per-step network. + + Args: + inputs (torch.Tensor): tensor input of shape [B, T, D], where D is the RNN input size + + rnn_init_state: rnn hidden state, initialize to zero state if set to None + + return_state (bool): whether to return hidden state + + Returns: + outputs: outputs of the per_step_net + + rnn_state: return rnn state at the end if return_state is set to True + """ + assert inputs.ndimension() == 3 # [B, T, D] + batch_size, seq_length, inp_dim = inputs.shape + if rnn_init_state is None: + rnn_init_state = self.get_rnn_init_state(batch_size, device=inputs.device) + + outputs, rnn_state = self.nets(inputs, rnn_init_state) + if self.per_step_net is not None: + outputs = TensorUtils.time_distributed(outputs, self.per_step_net) + + if return_state: + return outputs, rnn_state + else: + return outputs + + def forward_step(self, inputs, rnn_state): + """ + Forward a single step input through the RNN and per-step network, and return the new hidden state. + Args: + inputs (torch.Tensor): tensor input of shape [B, D], where D is the RNN input size + + rnn_state: rnn hidden state, initialize to zero state if set to None + + Returns: + outputs: outputs of the per_step_net + + rnn_state: return the new rnn state + """ + assert inputs.ndimension() == 2 + inputs = TensorUtils.to_sequence(inputs) + outputs, rnn_state = self.forward( + inputs, + rnn_init_state=rnn_state, + return_state=True, + ) + return outputs[:, 0], rnn_state + + +""" +================================================ +Visual Backbone Networks +================================================ +""" +class ConvBase(Module): + """ + Base class for ConvNets. + """ + def __init__(self): + super(ConvBase, self).__init__() + + # dirty hack - re-implement to pass the buck onto subclasses from ABC parent + def output_shape(self, input_shape): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + raise NotImplementedError + + def forward(self, inputs): + x = self.nets(inputs) + if list(self.output_shape(list(inputs.shape)[1:])) != list(x.shape)[1:]: + raise ValueError('Size mismatch: expect size %s, but got size %s' % ( + str(self.output_shape(list(inputs.shape)[1:])), str(list(x.shape)[1:])) + ) + return x + + +class ResNet18Conv(ConvBase): + """ + A ResNet18 block that can be used to process input images. + """ + def __init__( + self, + input_channel=3, + pretrained=False, + input_coord_conv=False, + ): + """ + Args: + input_channel (int): number of input channels for input images to the network. + If not equal to 3, modifies first conv layer in ResNet to handle the number + of input channels. + pretrained (bool): if True, load pretrained weights for all ResNet layers. + input_coord_conv (bool): if True, use a coordinate convolution for the first layer + (a convolution where input channels are modified to encode spatial pixel location) + """ + super(ResNet18Conv, self).__init__() + net = vision_models.resnet18(pretrained=pretrained) + + if input_coord_conv: + net.conv1 = CoordConv2d(input_channel, 64, kernel_size=7, stride=2, padding=3, bias=False) + elif input_channel != 3: + net.conv1 = nn.Conv2d(input_channel, 64, kernel_size=7, stride=2, padding=3, bias=False) + + # cut the last fc layer + self._input_coord_conv = input_coord_conv + self._input_channel = input_channel + self.nets = torch.nn.Sequential(*(list(net.children())[:-2])) + + def output_shape(self, input_shape): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + assert(len(input_shape) == 3) + out_h = int(math.ceil(input_shape[1] / 32.)) + out_w = int(math.ceil(input_shape[2] / 32.)) + return [512, out_h, out_w] + + def __repr__(self): + """Pretty print network.""" + header = '{}'.format(str(self.__class__.__name__)) + return header + '(input_channel={}, input_coord_conv={})'.format(self._input_channel, self._input_coord_conv) + + +class ResNet50Conv(ConvBase): + """ + A ResNet50 block that can be used to process input images. + """ + def __init__( + self, + input_channel=3, + pretrained=False, + input_coord_conv=False, + ): + """ + Args: + input_channel (int): number of input channels for input images to the network. + If not equal to 3, modifies first conv layer in ResNet to handle the number + of input channels. + pretrained (bool): if True, load pretrained weights for all ResNet layers. + input_coord_conv (bool): if True, use a coordinate convolution for the first layer + (a convolution where input channels are modified to encode spatial pixel location) + """ + super(ResNet50Conv, self).__init__() + net = vision_models.resnet50(pretrained=pretrained) + + if input_coord_conv: + # copied from ResNet18Conv. TODO: check if sizes need to be changed + net.conv1 = CoordConv2d(input_channel, 64, kernel_size=7, stride=2, padding=3, bias=False) + elif input_channel != 3: + # copied from ResNet18Conv. TODO: check if sizes need to be changed + net.conv1 = nn.Conv2d(input_channel, 64, kernel_size=7, stride=2, padding=3, bias=False) + + # cut the last fc layer + self._input_coord_conv = input_coord_conv + self._input_channel = input_channel + self.nets = torch.nn.Sequential(*(list(net.children())[:-2])) + + def output_shape(self, input_shape): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + assert(len(input_shape) == 3) + out_h = int(math.ceil(input_shape[1] / 32.)) + out_w = int(math.ceil(input_shape[2] / 32.)) + return [2048, out_h, out_w] + + def __repr__(self): + """Pretty print network.""" + header = '{}'.format(str(self.__class__.__name__)) + return header + '(input_channel={}, input_coord_conv={})'.format(self._input_channel, self._input_coord_conv) + + +class R3MConv(ConvBase): + """ + Base class for ConvNets pretrained with R3M (https://arxiv.org/abs/2203.12601) + """ + def __init__( + self, + input_channel=3, + r3m_model_class='resnet18', + freeze=True, + ): + """ + Using R3M pretrained observation encoder network proposed by https://arxiv.org/abs/2203.12601 + Args: + input_channel (int): number of input channels for input images to the network. + If not equal to 3, modifies first conv layer in ResNet to handle the number + of input channels. + r3m_model_class (str): select one of the r3m pretrained model "resnet18", "resnet34" or "resnet50" + freeze (bool): if True, use a frozen R3M pretrained model. + """ + super(R3MConv, self).__init__() + + try: + from r3m import load_r3m + except ImportError: + print("WARNING: could not load r3m library! Please follow https://github.com/facebookresearch/r3m to install R3M") + + net = load_r3m(r3m_model_class) + + assert input_channel == 3 # R3M only support input image with channel size 3 + assert r3m_model_class in ["resnet18", "resnet34", "resnet50"] # make sure the selected r3m model do exist + + # cut the last fc layer + self._input_channel = input_channel + self._r3m_model_class = r3m_model_class + self._freeze = freeze + self._input_coord_conv = False + self._pretrained = True + + preprocess = nn.Sequential( + transforms.Resize(256), + transforms.CenterCrop(224), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ) + self.nets = Sequential(*([preprocess] + list(net.module.convnet.children()))) + if freeze: + self.nets.freeze() + + self.weight_sum = np.sum([param.cpu().data.numpy().sum() for param in self.nets.parameters()]) + if freeze: + for param in self.nets.parameters(): + param.requires_grad = False + + self.nets.eval() + + def output_shape(self, input_shape): + """ + Function to compute output shape from inputs to this module. + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + assert(len(input_shape) == 3) + + if self._r3m_model_class == 'resnet50': + out_dim = 2048 + else: + out_dim = 512 + + return [out_dim, 1, 1] + + def __repr__(self): + """Pretty print network.""" + header = '{}'.format(str(self.__class__.__name__)) + return header + '(input_channel={}, input_coord_conv={}, pretrained={}, freeze={})'.format(self._input_channel, self._input_coord_conv, self._pretrained, self._freeze) + + +class MVPConv(ConvBase): + """ + Base class for ConvNets pretrained with MVP (https://arxiv.org/abs/2203.06173) + """ + def __init__( + self, + input_channel=3, + mvp_model_class='vitb-mae-egosoup', + freeze=True, + ): + """ + Using MVP pretrained observation encoder network proposed by https://arxiv.org/abs/2203.06173 + Args: + input_channel (int): number of input channels for input images to the network. + If not equal to 3, modifies first conv layer in ResNet to handle the number + of input channels. + mvp_model_class (str): select one of the mvp pretrained model "vits-mae-hoi", "vits-mae-in", "vits-sup-in", "vitb-mae-egosoup" or "vitl-256-mae-egosoup" + freeze (bool): if True, use a frozen MVP pretrained model. + """ + super(MVPConv, self).__init__() + + try: + import mvp + except ImportError: + print("WARNING: could not load mvp library! Please follow https://github.com/ir413/mvp to install MVP.") + + self.nets = mvp.load(mvp_model_class) + if freeze: + self.nets.freeze() + + assert input_channel == 3 # MVP only support input image with channel size 3 + assert mvp_model_class in ["vits-mae-hoi", "vits-mae-in", "vits-sup-in", "vitb-mae-egosoup", "vitl-256-mae-egosoup"] # make sure the selected r3m model do exist + + self._input_channel = input_channel + self._freeze = freeze + self._mvp_model_class = mvp_model_class + self._input_coord_conv = False + self._pretrained = True + + if '256' in mvp_model_class: + input_img_size = 256 + else: + input_img_size = 224 + self.preprocess = nn.Sequential( + transforms.Resize(input_img_size) + ) + + def forward(self, inputs): + x = self.preprocess(inputs) + x = self.nets(x) + if list(self.output_shape(list(inputs.shape)[1:])) != list(x.shape)[1:]: + raise ValueError('Size mismatch: expect size %s, but got size %s' % ( + str(self.output_shape(list(inputs.shape)[1:])), str(list(x.shape)[1:])) + ) + return x + + def output_shape(self, input_shape): + """ + Function to compute output shape from inputs to this module. + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + assert(len(input_shape) == 3) + if 'vitb' in self._mvp_model_class: + output_shape = [768] + elif 'vitl' in self._mvp_model_class: + output_shape = [1024] + else: + output_shape = [384] + return output_shape + + def __repr__(self): + """Pretty print network.""" + header = '{}'.format(str(self.__class__.__name__)) + return header + '(input_channel={}, input_coord_conv={}, pretrained={}, freeze={})'.format(self._input_channel, self._input_coord_conv, self._pretrained, self._freeze) + + +class CoordConv2d(nn.Conv2d, Module): + """ + 2D Coordinate Convolution + + Source: An Intriguing Failing of Convolutional Neural Networks and the CoordConv Solution + https://arxiv.org/abs/1807.03247 + (e.g. adds 2 channels per input feature map corresponding to (x, y) location on map) + """ + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + padding=0, + dilation=1, + groups=1, + bias=True, + padding_mode='zeros', + coord_encoding='position', + ): + """ + Args: + in_channels: number of channels of the input tensor [C, H, W] + out_channels: number of output channels of the layer + kernel_size: convolution kernel size + stride: conv stride + padding: conv padding + dilation: conv dilation + groups: conv groups + bias: conv bias + padding_mode: conv padding mode + coord_encoding: type of coordinate encoding. currently only 'position' is implemented + """ + + assert(coord_encoding in ['position']) + self.coord_encoding = coord_encoding + if coord_encoding == 'position': + in_channels += 2 # two extra channel for positional encoding + self._position_enc = None # position encoding + else: + raise Exception("CoordConv2d: coord encoding {} not implemented".format(self.coord_encoding)) + nn.Conv2d.__init__( + self, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=padding_mode + ) + + def output_shape(self, input_shape): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + + # adds 2 to channel dimension + return [input_shape[0] + 2] + input_shape[1:] + + def forward(self, input): + b, c, h, w = input.shape + if self.coord_encoding == 'position': + if self._position_enc is None: + pos_y, pos_x = torch.meshgrid(torch.arange(h), torch.arange(w)) + pos_y = pos_y.float().to(input.device) / float(h) + pos_x = pos_x.float().to(input.device) / float(w) + self._position_enc = torch.stack((pos_y, pos_x)).unsqueeze(0) + pos_enc = self._position_enc.expand(b, -1, -1, -1) + input = torch.cat((input, pos_enc), dim=1) + return super(CoordConv2d, self).forward(input) + + +class ShallowConv(ConvBase): + """ + A shallow convolutional encoder from https://rll.berkeley.edu/dsae/dsae.pdf + """ + def __init__(self, input_channel=3, output_channel=32): + super(ShallowConv, self).__init__() + self._input_channel = input_channel + self._output_channel = output_channel + self.nets = nn.Sequential( + torch.nn.Conv2d(input_channel, 64, kernel_size=7, stride=2, padding=3), + torch.nn.ReLU(), + torch.nn.Conv2d(64, 32, kernel_size=1, stride=1, padding=0), + torch.nn.ReLU(), + torch.nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1), + torch.nn.ReLU(), + torch.nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1), + ) + + def output_shape(self, input_shape): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + assert(len(input_shape) == 3) + assert(input_shape[0] == self._input_channel) + out_h = int(math.floor(input_shape[1] / 2.)) + out_w = int(math.floor(input_shape[2] / 2.)) + return [self._output_channel, out_h, out_w] + + +class Conv1dBase(Module): + """ + Base class for stacked Conv1d layers. + + Args: + input_channel (int): Number of channels for inputs to this network + activation (None or str): Per-layer activation to use. Defaults to "relu". Valid options are + currently {relu, None} for no activation + out_channels (list of int): Output channel size for each sequential Conv1d layer + kernel_size (list of int): Kernel sizes for each sequential Conv1d layer + stride (list of int): Stride sizes for each sequential Conv1d layer + conv_kwargs (dict): additional nn.Conv1D args to use, in list form, where the ith element corresponds to the + argument to be passed to the ith Conv1D layer. + See https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html for specific possible arguments. + """ + def __init__( + self, + input_channel=1, + activation="relu", + out_channels=(32, 64, 64), + kernel_size=(8, 4, 2), + stride=(4, 2, 1), + **conv_kwargs, + ): + super(Conv1dBase, self).__init__() + + # Get activation requested + activation = CONV_ACTIVATIONS[activation] + + # Generate network + self.n_layers = len(out_channels) + layers = OrderedDict() + for i in range(self.n_layers): + layer_kwargs = {k: v[i] for k, v in conv_kwargs.items()} + layers[f'conv{i}'] = nn.Conv1d( + in_channels=input_channel, + **layer_kwargs, + ) + if activation is not None: + layers[f'act{i}'] = activation() + input_channel = layer_kwargs["out_channels"] + + # Store network + self.nets = nn.Sequential(layers) + + def output_shape(self, input_shape): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + channels, length = input_shape + for i in range(self.n_layers): + net = getattr(self.nets, f"conv{i}") + channels = net.out_channels + length = int((length + 2 * net.padding[0] - net.dilation[0] * (net.kernel_size[0] - 1) - 1) / net.stride[0]) + 1 + return [channels, length] + + def forward(self, inputs): + x = self.nets(inputs) + if list(self.output_shape(list(inputs.shape)[1:])) != list(x.shape)[1:]: + raise ValueError('Size mismatch: expect size %s, but got size %s' % ( + str(self.output_shape(list(inputs.shape)[1:])), str(list(x.shape)[1:])) + ) + return x + + +""" +================================================ +Pooling Networks +================================================ +""" +class SpatialSoftmax(ConvBase): + """ + Spatial Softmax Layer. + + Based on Deep Spatial Autoencoders for Visuomotor Learning by Finn et al. + https://rll.berkeley.edu/dsae/dsae.pdf + """ + def __init__( + self, + input_shape, + num_kp=32, + temperature=1., + learnable_temperature=False, + output_variance=False, + noise_std=0.0, + ): + """ + Args: + input_shape (list): shape of the input feature (C, H, W) + num_kp (int): number of keypoints (None for not using spatialsoftmax) + temperature (float): temperature term for the softmax. + learnable_temperature (bool): whether to learn the temperature + output_variance (bool): treat attention as a distribution, and compute second-order statistics to return + noise_std (float): add random spatial noise to the predicted keypoints + """ + super(SpatialSoftmax, self).__init__() + assert len(input_shape) == 3 + self._in_c, self._in_h, self._in_w = input_shape # (C, H, W) + + if num_kp is not None: + self.nets = torch.nn.Conv2d(self._in_c, num_kp, kernel_size=1) + self._num_kp = num_kp + else: + self.nets = None + self._num_kp = self._in_c + self.learnable_temperature = learnable_temperature + self.output_variance = output_variance + self.noise_std = noise_std + + if self.learnable_temperature: + # temperature will be learned + temperature = torch.nn.Parameter(torch.ones(1) * temperature, requires_grad=True) + self.register_parameter('temperature', temperature) + else: + # temperature held constant after initialization + temperature = torch.nn.Parameter(torch.ones(1) * temperature, requires_grad=False) + self.register_buffer('temperature', temperature) + + pos_x, pos_y = np.meshgrid( + np.linspace(-1., 1., self._in_w), + np.linspace(-1., 1., self._in_h) + ) + pos_x = torch.from_numpy(pos_x.reshape(1, self._in_h * self._in_w)).float() + pos_y = torch.from_numpy(pos_y.reshape(1, self._in_h * self._in_w)).float() + self.register_buffer('pos_x', pos_x) + self.register_buffer('pos_y', pos_y) + + self.kps = None + + def __repr__(self): + """Pretty print network.""" + header = format(str(self.__class__.__name__)) + return header + '(num_kp={}, temperature={}, noise={})'.format( + self._num_kp, self.temperature.item(), self.noise_std) + + def output_shape(self, input_shape): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + assert(len(input_shape) == 3) + assert(input_shape[0] == self._in_c) + return [self._num_kp, 2] + + def forward(self, feature): + """ + Forward pass through spatial softmax layer. For each keypoint, a 2D spatial + probability distribution is created using a softmax, where the support is the + pixel locations. This distribution is used to compute the expected value of + the pixel location, which becomes a keypoint of dimension 2. K such keypoints + are created. + + Returns: + out (torch.Tensor or tuple): mean keypoints of shape [B, K, 2], and possibly + keypoint variance of shape [B, K, 2, 2] corresponding to the covariance + under the 2D spatial softmax distribution + """ + assert(feature.shape[1] == self._in_c) + assert(feature.shape[2] == self._in_h) + assert(feature.shape[3] == self._in_w) + if self.nets is not None: + feature = self.nets(feature) + + # [B, K, H, W] -> [B * K, H * W] where K is number of keypoints + feature = feature.reshape(-1, self._in_h * self._in_w) + # 2d softmax normalization + attention = F.softmax(feature / self.temperature, dim=-1) + # [1, H * W] x [B * K, H * W] -> [B * K, 1] for spatial coordinate mean in x and y dimensions + expected_x = torch.sum(self.pos_x * attention, dim=1, keepdim=True) + expected_y = torch.sum(self.pos_y * attention, dim=1, keepdim=True) + # stack to [B * K, 2] + expected_xy = torch.cat([expected_x, expected_y], 1) + # reshape to [B, K, 2] + feature_keypoints = expected_xy.view(-1, self._num_kp, 2) + + if self.training: + noise = torch.randn_like(feature_keypoints) * self.noise_std + feature_keypoints += noise + + if self.output_variance: + # treat attention as a distribution, and compute second-order statistics to return + expected_xx = torch.sum(self.pos_x * self.pos_x * attention, dim=1, keepdim=True) + expected_yy = torch.sum(self.pos_y * self.pos_y * attention, dim=1, keepdim=True) + expected_xy = torch.sum(self.pos_x * self.pos_y * attention, dim=1, keepdim=True) + var_x = expected_xx - expected_x * expected_x + var_y = expected_yy - expected_y * expected_y + var_xy = expected_xy - expected_x * expected_y + # stack to [B * K, 4] and then reshape to [B, K, 2, 2] where last 2 dims are covariance matrix + feature_covar = torch.cat([var_x, var_xy, var_xy, var_y], 1).reshape(-1, self._num_kp, 2, 2) + feature_keypoints = (feature_keypoints, feature_covar) + + if isinstance(feature_keypoints, tuple): + self.kps = (feature_keypoints[0].detach(), feature_keypoints[1].detach()) + else: + self.kps = feature_keypoints.detach() + return feature_keypoints + + +class SpatialMeanPool(Module): + """ + Module that averages inputs across all spatial dimensions (dimension 2 and after), + leaving only the batch and channel dimensions. + """ + def __init__(self, input_shape): + super(SpatialMeanPool, self).__init__() + assert len(input_shape) == 3 # [C, H, W] + self.in_shape = input_shape + + def output_shape(self, input_shape=None): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + return list(self.in_shape[:1]) # [C, H, W] -> [C] + + def forward(self, inputs): + """Forward pass - average across all dimensions except batch and channel.""" + return TensorUtils.flatten(inputs, begin_axis=2).mean(dim=2) + + +class FeatureAggregator(Module): + """ + Helpful class for aggregating features across a dimension. This is useful in + practice when training models that break an input image up into several patches + since features can be extraced per-patch using the same encoder and then + aggregated using this module. + """ + def __init__(self, dim=1, agg_type="avg"): + super(FeatureAggregator, self).__init__() + self.dim = dim + self.agg_type = agg_type + + def set_weight(self, w): + assert self.agg_type == "w_avg" + self.agg_weight = w + + def clear_weight(self): + assert self.agg_type == "w_avg" + self.agg_weight = None + + def output_shape(self, input_shape): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + # aggregates on @self.dim, so it is removed from the output shape + return list(input_shape[:self.dim]) + list(input_shape[self.dim+1:]) + + def forward(self, x): + """Forward pooling pass.""" + if self.agg_type == "avg": + # mean-pooling + return torch.mean(x, dim=1) + if self.agg_type == "w_avg": + # weighted mean-pooling + return torch.sum(x * self.agg_weight, dim=1) + raise Exception("unexpected agg type: {}".forward(self.agg_type)) diff --git a/aloha-devel/robomimic/models/obs_nets.py b/aloha-devel/robomimic/models/obs_nets.py new file mode 100644 index 0000000000000000000000000000000000000000..7aa6feceef4070f13ff1ecaf7da67051fb9010f4 --- /dev/null +++ b/aloha-devel/robomimic/models/obs_nets.py @@ -0,0 +1,1121 @@ +""" +Contains torch Modules that help deal with inputs consisting of multiple +modalities. This is extremely common when networks must deal with one or +more observation dictionaries, where each input dictionary can have +observation keys of a certain modality and shape. + +As an example, an observation could consist of a flat "robot0_eef_pos" observation key, +and a 3-channel RGB "agentview_image" observation key. +""" +import sys +import numpy as np +import textwrap +from copy import deepcopy +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.distributions as D + +from robomimic.utils.python_utils import extract_class_init_kwargs_from_dict +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.obs_utils as ObsUtils +from robomimic.models.base_nets import Module, Sequential, MLP, RNN_Base, ResNet18Conv, SpatialSoftmax, \ + FeatureAggregator +from robomimic.models.obs_core import VisualCore, Randomizer +from robomimic.models.transformers import PositionalEncoding, GPT_Backbone + + +def obs_encoder_factory( + obs_shapes, + feature_activation=nn.ReLU, + encoder_kwargs=None, + ): + """ + Utility function to create an @ObservationEncoder from kwargs specified in config. + + Args: + obs_shapes (OrderedDict): a dictionary that maps observation key to + expected shapes for observations. + + feature_activation: non-linearity to apply after each obs net - defaults to ReLU. Pass + None to apply no activation. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should be + nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + enc = ObservationEncoder(feature_activation=feature_activation) + for k, obs_shape in obs_shapes.items(): + obs_modality = ObsUtils.OBS_KEYS_TO_MODALITIES[k] + enc_kwargs = deepcopy(ObsUtils.DEFAULT_ENCODER_KWARGS[obs_modality]) if encoder_kwargs is None else \ + deepcopy(encoder_kwargs[obs_modality]) + + # Sanity check for kwargs in case they don't exist / are None + if enc_kwargs.get("core_kwargs", None) is None: + enc_kwargs["core_kwargs"] = {} + # Add in input shape info + enc_kwargs["core_kwargs"]["input_shape"] = obs_shape + # If group class is specified, then make sure corresponding kwargs only contain relevant kwargs + if enc_kwargs["core_class"] is not None: + enc_kwargs["core_kwargs"] = extract_class_init_kwargs_from_dict( + cls=ObsUtils.OBS_ENCODER_CORES[enc_kwargs["core_class"]], + dic=enc_kwargs["core_kwargs"], + copy=False, + ) + + # Add in input shape info + randomizers = [] + obs_randomizer_class_list = enc_kwargs["obs_randomizer_class"] + obs_randomizer_kwargs_list = enc_kwargs["obs_randomizer_kwargs"] + + if not isinstance(obs_randomizer_class_list, list): + obs_randomizer_class_list = [obs_randomizer_class_list] + + if not isinstance(obs_randomizer_kwargs_list, list): + obs_randomizer_kwargs_list = [obs_randomizer_kwargs_list] + + for rand_class, rand_kwargs in zip(obs_randomizer_class_list, obs_randomizer_kwargs_list): + rand = None + if rand_class is not None: + rand_kwargs["input_shape"] = obs_shape + rand_kwargs = extract_class_init_kwargs_from_dict( + cls=ObsUtils.OBS_RANDOMIZERS[rand_class], + dic=rand_kwargs, + copy=False, + ) + rand = ObsUtils.OBS_RANDOMIZERS[rand_class](**rand_kwargs) + randomizers.append(rand) + + enc.register_obs_key( + name=k, + shape=obs_shape, + net_class=enc_kwargs["core_class"], + net_kwargs=enc_kwargs["core_kwargs"], + randomizers=randomizers, + ) + + enc.make() + return enc + + +class ObservationEncoder(Module): + """ + Module that processes inputs by observation key and then concatenates the processed + observation keys together. Each key is processed with an encoder head network. + Call @register_obs_key to register observation keys with the encoder and then + finally call @make to create the encoder networks. + """ + def __init__(self, feature_activation=nn.ReLU): + """ + Args: + feature_activation: non-linearity to apply after each obs net - defaults to ReLU. Pass + None to apply no activation. + """ + super(ObservationEncoder, self).__init__() + self.obs_shapes = OrderedDict() + self.obs_nets_classes = OrderedDict() + self.obs_nets_kwargs = OrderedDict() + self.obs_share_mods = OrderedDict() + self.obs_nets = nn.ModuleDict() + self.obs_randomizers = nn.ModuleDict() + self.feature_activation = feature_activation + self._locked = False + + def register_obs_key( + self, + name, + shape, + net_class=None, + net_kwargs=None, + net=None, + randomizers=None, + share_net_from=None, + ): + """ + Register an observation key that this encoder should be responsible for. + + Args: + name (str): modality name + shape (int tuple): shape of modality + net_class (str): name of class in base_nets.py that should be used + to process this observation key before concatenation. Pass None to flatten + and concatenate the observation key directly. + net_kwargs (dict): arguments to pass to @net_class + net (Module instance): if provided, use this Module to process the observation key + instead of creating a different net + randomizer (Randomizer instance): if provided, use this Module to augment observation keys + coming in to the encoder, and possibly augment the processed output as well + share_net_from (str): if provided, use the same instance of @net_class + as another observation key. This observation key must already exist in this encoder. + Warning: Note that this does not share the observation key randomizer + """ + assert not self._locked, "ObservationEncoder: @register_obs_key called after @make" + assert name not in self.obs_shapes, "ObservationEncoder: modality {} already exists".format(name) + + if net is not None: + assert isinstance(net, Module), "ObservationEncoder: @net must be instance of Module class" + assert (net_class is None) and (net_kwargs is None) and (share_net_from is None), \ + "ObservationEncoder: @net provided - ignore other net creation options" + + if share_net_from is not None: + # share processing with another modality + assert (net_class is None) and (net_kwargs is None) + assert share_net_from in self.obs_shapes + + net_kwargs = deepcopy(net_kwargs) if net_kwargs is not None else {} + for rand in randomizers: + if rand is not None: + assert isinstance(rand, Randomizer) + if net_kwargs is not None: + # update input shape to visual core + net_kwargs["input_shape"] = rand.output_shape_in(shape) + + self.obs_shapes[name] = shape + self.obs_nets_classes[name] = net_class + self.obs_nets_kwargs[name] = net_kwargs + self.obs_nets[name] = net + self.obs_randomizers[name] = nn.ModuleList(randomizers) + self.obs_share_mods[name] = share_net_from + + def make(self): + """ + Creates the encoder networks and locks the encoder so that more modalities cannot be added. + """ + assert not self._locked, "ObservationEncoder: @make called more than once" + self._create_layers() + self._locked = True + + def _create_layers(self): + """ + Creates all networks and layers required by this encoder using the registered modalities. + """ + assert not self._locked, "ObservationEncoder: layers have already been created" + + for k in self.obs_shapes: + if self.obs_nets_classes[k] is not None: + # create net to process this modality + self.obs_nets[k] = ObsUtils.OBS_ENCODER_CORES[self.obs_nets_classes[k]](**self.obs_nets_kwargs[k]) + elif self.obs_share_mods[k] is not None: + # make sure net is shared with another modality + self.obs_nets[k] = self.obs_nets[self.obs_share_mods[k]] + + self.activation = None + if self.feature_activation is not None: + self.activation = self.feature_activation() + + def forward(self, obs_dict): + """ + Processes modalities according to the ordering in @self.obs_shapes. For each + modality, it is processed with a randomizer (if present), an encoder + network (if present), and again with the randomizer (if present), flattened, + and then concatenated with the other processed modalities. + + Args: + obs_dict (OrderedDict): dictionary that maps modalities to torch.Tensor + batches that agree with @self.obs_shapes. All modalities in + @self.obs_shapes must be present, but additional modalities + can also be present. + + Returns: + feats (torch.Tensor): flat features of shape [B, D] + """ + assert self._locked, "ObservationEncoder: @make has not been called yet" + + # ensure all modalities that the encoder handles are present + assert set(self.obs_shapes.keys()).issubset(obs_dict), "ObservationEncoder: {} does not contain all modalities {}".format( + list(obs_dict.keys()), list(self.obs_shapes.keys()) + ) + + # process modalities by order given by @self.obs_shapes + feats = [] + for k in self.obs_shapes: + x = obs_dict[k] + # maybe process encoder input with randomizer + for rand in self.obs_randomizers[k]: + if rand is not None: + x = rand.forward_in(x) + # maybe process with obs net + if self.obs_nets[k] is not None: + x = self.obs_nets[k](x) + if self.activation is not None: + x = self.activation(x) + # maybe process encoder output with randomizer + for rand in self.obs_randomizers[k]: + if rand is not None: + x = rand.forward_out(x) + # flatten to [B, D] + x = TensorUtils.flatten(x, begin_axis=1) + feats.append(x) + + # concatenate all features together + return torch.cat(feats, dim=-1) + + def output_shape(self, input_shape=None): + """ + Compute the output shape of the encoder. + """ + feat_dim = 0 + for k in self.obs_shapes: + feat_shape = self.obs_shapes[k] + for rand in self.obs_randomizers[k]: + if rand is not None: + feat_shape = rand.output_shape_in(feat_shape) + if self.obs_nets[k] is not None: + feat_shape = self.obs_nets[k].output_shape(feat_shape) + for rand in self.obs_randomizers[k]: + if rand is not None: + feat_shape = rand.output_shape_out(feat_shape) + feat_dim += int(np.prod(feat_shape)) + return [feat_dim] + + def __repr__(self): + """ + Pretty print the encoder. + """ + header = '{}'.format(str(self.__class__.__name__)) + msg = '' + for k in self.obs_shapes: + msg += textwrap.indent('\nKey(\n', ' ' * 4) + indent = ' ' * 8 + msg += textwrap.indent("name={}\nshape={}\n".format(k, self.obs_shapes[k]), indent) + msg += textwrap.indent("modality={}\n".format(ObsUtils.OBS_KEYS_TO_MODALITIES[k]), indent) + msg += textwrap.indent("randomizer={}\n".format(self.obs_randomizers[k]), indent) + msg += textwrap.indent("net={}\n".format(self.obs_nets[k]), indent) + msg += textwrap.indent("sharing_from={}\n".format(self.obs_share_mods[k]), indent) + msg += textwrap.indent(")", ' ' * 4) + msg += textwrap.indent("\noutput_shape={}".format(self.output_shape()), ' ' * 4) + msg = header + '(' + msg + '\n)' + return msg + + +class ObservationDecoder(Module): + """ + Module that can generate observation outputs by modality. Inputs are assumed + to be flat (usually outputs from some hidden layer). Each observation output + is generated with a linear layer from these flat inputs. Subclass this + module in order to implement more complex schemes for generating each + modality. + """ + def __init__( + self, + decode_shapes, + input_feat_dim, + ): + """ + Args: + decode_shapes (OrderedDict): a dictionary that maps observation key to + expected shape. This is used to generate output modalities from the + input features. + + input_feat_dim (int): flat input dimension size + """ + super(ObservationDecoder, self).__init__() + + # important: sort observation keys to ensure consistent ordering of modalities + assert isinstance(decode_shapes, OrderedDict) + self.obs_shapes = OrderedDict() + for k in decode_shapes: + self.obs_shapes[k] = decode_shapes[k] + + self.input_feat_dim = input_feat_dim + self._create_layers() + + def _create_layers(self): + """ + Create a linear layer to predict each modality. + """ + self.nets = nn.ModuleDict() + for k in self.obs_shapes: + layer_out_dim = int(np.prod(self.obs_shapes[k])) + self.nets[k] = nn.Linear(self.input_feat_dim, layer_out_dim) + + def output_shape(self, input_shape=None): + """ + Returns output shape for this module, which is a dictionary instead + of a list since outputs are dictionaries. + """ + return { k : list(self.obs_shapes[k]) for k in self.obs_shapes } + + def forward(self, feats): + """ + Predict each modality from input features, and reshape to each modality's shape. + """ + output = {} + for k in self.obs_shapes: + out = self.nets[k](feats) + output[k] = out.reshape(-1, *self.obs_shapes[k]) + return output + + def __repr__(self): + """Pretty print network.""" + header = '{}'.format(str(self.__class__.__name__)) + msg = '' + for k in self.obs_shapes: + msg += textwrap.indent('\nKey(\n', ' ' * 4) + indent = ' ' * 8 + msg += textwrap.indent("name={}\nshape={}\n".format(k, self.obs_shapes[k]), indent) + msg += textwrap.indent("modality={}\n".format(ObsUtils.OBS_KEYS_TO_MODALITIES[k]), indent) + msg += textwrap.indent("net=({})\n".format(self.nets[k]), indent) + msg += textwrap.indent(")", ' ' * 4) + msg = header + '(' + msg + '\n)' + return msg + + +class ObservationGroupEncoder(Module): + """ + This class allows networks to encode multiple observation dictionaries into a single + flat, concatenated vector representation. It does this by assigning each observation + dictionary (observation group) an @ObservationEncoder object. + + The class takes a dictionary of dictionaries, @observation_group_shapes. + Each key corresponds to a observation group (e.g. 'obs', 'subgoal', 'goal') + and each OrderedDict should be a map between modalities and + expected input shapes (e.g. { 'image' : (3, 120, 160) }). + """ + def __init__( + self, + observation_group_shapes, + feature_activation=nn.ReLU, + encoder_kwargs=None, + ): + """ + Args: + observation_group_shapes (OrderedDict): a dictionary of dictionaries. + Each key in this dictionary should specify an observation group, and + the value should be an OrderedDict that maps modalities to + expected shapes. + + feature_activation: non-linearity to apply after each obs net - defaults to ReLU. Pass + None to apply no activation. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + super(ObservationGroupEncoder, self).__init__() + + # type checking + assert isinstance(observation_group_shapes, OrderedDict) + assert np.all([isinstance(observation_group_shapes[k], OrderedDict) for k in observation_group_shapes]) + + self.observation_group_shapes = observation_group_shapes + + # create an observation encoder per observation group + self.nets = nn.ModuleDict() + for obs_group in self.observation_group_shapes: + self.nets[obs_group] = obs_encoder_factory( + obs_shapes=self.observation_group_shapes[obs_group], + feature_activation=feature_activation, + encoder_kwargs=encoder_kwargs, + ) + + def forward(self, **inputs): + """ + Process each set of inputs in its own observation group. + + Args: + inputs (dict): dictionary that maps observation groups to observation + dictionaries of torch.Tensor batches that agree with + @self.observation_group_shapes. All observation groups in + @self.observation_group_shapes must be present, but additional + observation groups can also be present. Note that these are specified + as kwargs for ease of use with networks that name each observation + stream in their forward calls. + + Returns: + outputs (torch.Tensor): flat outputs of shape [B, D] + """ + + # ensure all observation groups we need are present + assert set(self.observation_group_shapes.keys()).issubset(inputs), "{} does not contain all observation groups {}".format( + list(inputs.keys()), list(self.observation_group_shapes.keys()) + ) + + outputs = [] + # Deterministic order since self.observation_group_shapes is OrderedDict + for obs_group in self.observation_group_shapes: + # pass through encoder + outputs.append( + self.nets[obs_group].forward(inputs[obs_group]) + ) + + return torch.cat(outputs, dim=-1) + + def output_shape(self): + """ + Compute the output shape of this encoder. + """ + feat_dim = 0 + for obs_group in self.observation_group_shapes: + # get feature dimension of these keys + feat_dim += self.nets[obs_group].output_shape()[0] + return [feat_dim] + + def __repr__(self): + """Pretty print network.""" + header = '{}'.format(str(self.__class__.__name__)) + msg = '' + for k in self.observation_group_shapes: + msg += '\n' + indent = ' ' * 4 + msg += textwrap.indent("group={}\n{}".format(k, self.nets[k]), indent) + msg = header + '(' + msg + '\n)' + return msg + + +class MIMO_MLP(Module): + """ + Extension to MLP to accept multiple observation dictionaries as input and + to output dictionaries of tensors. Inputs are specified as a dictionary of + observation dictionaries, with each key corresponding to an observation group. + + This module utilizes @ObservationGroupEncoder to process the multiple input dictionaries and + @ObservationDecoder to generate tensor dictionaries. The default behavior + for encoding the inputs is to process visual inputs with a learned CNN and concatenating + the flat encodings with the other flat inputs. The default behavior for generating + outputs is to use a linear layer branch to produce each modality separately + (including visual outputs). + """ + def __init__( + self, + input_obs_group_shapes, + output_shapes, + layer_dims, + layer_func=nn.Linear, + activation=nn.ReLU, + encoder_kwargs=None, + ): + """ + Args: + input_obs_group_shapes (OrderedDict): a dictionary of dictionaries. + Each key in this dictionary should specify an observation group, and + the value should be an OrderedDict that maps modalities to + expected shapes. + + output_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for outputs. + + layer_dims ([int]): sequence of integers for the MLP hidden layer sizes + + layer_func: mapping per MLP layer - defaults to Linear + + activation: non-linearity per MLP layer - defaults to ReLU + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + super(MIMO_MLP, self).__init__() + + assert isinstance(input_obs_group_shapes, OrderedDict) + assert np.all([isinstance(input_obs_group_shapes[k], OrderedDict) for k in input_obs_group_shapes]) + assert isinstance(output_shapes, OrderedDict) + + self.input_obs_group_shapes = input_obs_group_shapes + self.output_shapes = output_shapes + + self.nets = nn.ModuleDict() + + # Encoder for all observation groups. + self.nets["encoder"] = ObservationGroupEncoder( + observation_group_shapes=input_obs_group_shapes, + encoder_kwargs=encoder_kwargs, + ) + + # flat encoder output dimension + mlp_input_dim = self.nets["encoder"].output_shape()[0] + + # intermediate MLP layers + self.nets["mlp"] = MLP( + input_dim=mlp_input_dim, + output_dim=layer_dims[-1], + layer_dims=layer_dims[:-1], + layer_func=layer_func, + activation=activation, + output_activation=activation, # make sure non-linearity is applied before decoder + ) + + # decoder for output modalities + self.nets["decoder"] = ObservationDecoder( + decode_shapes=self.output_shapes, + input_feat_dim=layer_dims[-1], + ) + + def output_shape(self, input_shape=None): + """ + Returns output shape for this module, which is a dictionary instead + of a list since outputs are dictionaries. + """ + return { k : list(self.output_shapes[k]) for k in self.output_shapes } + + def forward(self, **inputs): + """ + Process each set of inputs in its own observation group. + + Args: + inputs (dict): a dictionary of dictionaries with one dictionary per + observation group. Each observation group's dictionary should map + modality to torch.Tensor batches. Should be consistent with + @self.input_obs_group_shapes. + + Returns: + outputs (dict): dictionary of output torch.Tensors, that corresponds + to @self.output_shapes + """ + enc_outputs = self.nets["encoder"](**inputs) + mlp_out = self.nets["mlp"](enc_outputs) + return self.nets["decoder"](mlp_out) + + def _to_string(self): + """ + Subclasses should override this method to print out info about network / policy. + """ + return '' + + def __repr__(self): + """Pretty print network.""" + header = '{}'.format(str(self.__class__.__name__)) + msg = '' + indent = ' ' * 4 + if self._to_string() != '': + msg += textwrap.indent("\n" + self._to_string() + "\n", indent) + msg += textwrap.indent("\nencoder={}".format(self.nets["encoder"]), indent) + msg += textwrap.indent("\n\nmlp={}".format(self.nets["mlp"]), indent) + msg += textwrap.indent("\n\ndecoder={}".format(self.nets["decoder"]), indent) + msg = header + '(' + msg + '\n)' + return msg + + +class RNN_MIMO_MLP(Module): + """ + A wrapper class for a multi-step RNN and a per-step MLP and a decoder. + + Structure: [encoder -> rnn -> mlp -> decoder] + + All temporal inputs are processed by a shared @ObservationGroupEncoder, + followed by an RNN, and then a per-step multi-output MLP. + """ + def __init__( + self, + input_obs_group_shapes, + output_shapes, + mlp_layer_dims, + rnn_hidden_dim, + rnn_num_layers, + rnn_type="LSTM", # [LSTM, GRU] + rnn_kwargs=None, + mlp_activation=nn.ReLU, + mlp_layer_func=nn.Linear, + per_step=True, + encoder_kwargs=None, + ): + """ + Args: + input_obs_group_shapes (OrderedDict): a dictionary of dictionaries. + Each key in this dictionary should specify an observation group, and + the value should be an OrderedDict that maps modalities to + expected shapes. + + output_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for outputs. + + rnn_hidden_dim (int): RNN hidden dimension + + rnn_num_layers (int): number of RNN layers + + rnn_type (str): [LSTM, GRU] + + rnn_kwargs (dict): kwargs for the rnn model + + per_step (bool): if True, apply the MLP and observation decoder into @output_shapes + at every step of the RNN. Otherwise, apply them to the final hidden state of the + RNN. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + super(RNN_MIMO_MLP, self).__init__() + assert isinstance(input_obs_group_shapes, OrderedDict) + assert np.all([isinstance(input_obs_group_shapes[k], OrderedDict) for k in input_obs_group_shapes]) + assert isinstance(output_shapes, OrderedDict) + self.input_obs_group_shapes = input_obs_group_shapes + self.output_shapes = output_shapes + self.per_step = per_step + + self.nets = nn.ModuleDict() + + # Encoder for all observation groups. + self.nets["encoder"] = ObservationGroupEncoder( + observation_group_shapes=input_obs_group_shapes, + encoder_kwargs=encoder_kwargs, + ) + + # flat encoder output dimension + rnn_input_dim = self.nets["encoder"].output_shape()[0] + + # bidirectional RNNs mean that the output of RNN will be twice the hidden dimension + rnn_is_bidirectional = rnn_kwargs.get("bidirectional", False) + num_directions = int(rnn_is_bidirectional) + 1 # 2 if bidirectional, 1 otherwise + rnn_output_dim = num_directions * rnn_hidden_dim + + per_step_net = None + self._has_mlp = (len(mlp_layer_dims) > 0) + if self._has_mlp: + self.nets["mlp"] = MLP( + input_dim=rnn_output_dim, + output_dim=mlp_layer_dims[-1], + layer_dims=mlp_layer_dims[:-1], + output_activation=mlp_activation, + layer_func=mlp_layer_func + ) + self.nets["decoder"] = ObservationDecoder( + decode_shapes=self.output_shapes, + input_feat_dim=mlp_layer_dims[-1], + ) + if self.per_step: + per_step_net = Sequential(self.nets["mlp"], self.nets["decoder"]) + else: + self.nets["decoder"] = ObservationDecoder( + decode_shapes=self.output_shapes, + input_feat_dim=rnn_output_dim, + ) + if self.per_step: + per_step_net = self.nets["decoder"] + + # core network + self.nets["rnn"] = RNN_Base( + input_dim=rnn_input_dim, + rnn_hidden_dim=rnn_hidden_dim, + rnn_num_layers=rnn_num_layers, + rnn_type=rnn_type, + per_step_net=per_step_net, + rnn_kwargs=rnn_kwargs + ) + + def get_rnn_init_state(self, batch_size, device): + """ + Get a default RNN state (zeros) + + Args: + batch_size (int): batch size dimension + + device: device the hidden state should be sent to. + + Returns: + hidden_state (torch.Tensor or tuple): returns hidden state tensor or tuple of hidden state tensors + depending on the RNN type + """ + return self.nets["rnn"].get_rnn_init_state(batch_size, device=device) + + def output_shape(self, input_shape): + """ + Returns output shape for this module, which is a dictionary instead + of a list since outputs are dictionaries. + + Args: + input_shape (dict): dictionary of dictionaries, where each top-level key + corresponds to an observation group, and the low-level dictionaries + specify the shape for each modality in an observation dictionary + """ + + # infers temporal dimension from input shape + obs_group = list(self.input_obs_group_shapes.keys())[0] + mod = list(self.input_obs_group_shapes[obs_group].keys())[0] + T = input_shape[obs_group][mod][0] + TensorUtils.assert_size_at_dim(input_shape, size=T, dim=0, + msg="RNN_MIMO_MLP: input_shape inconsistent in temporal dimension") + # returns a dictionary instead of list since outputs are dictionaries + return { k : [T] + list(self.output_shapes[k]) for k in self.output_shapes } + + def forward(self, rnn_init_state=None, return_state=False, **inputs): + """ + Args: + inputs (dict): a dictionary of dictionaries with one dictionary per + observation group. Each observation group's dictionary should map + modality to torch.Tensor batches. Should be consistent with + @self.input_obs_group_shapes. First two leading dimensions should + be batch and time [B, T, ...] for each tensor. + + rnn_init_state: rnn hidden state, initialize to zero state if set to None + + return_state (bool): whether to return hidden state + + Returns: + outputs (dict): dictionary of output torch.Tensors, that corresponds + to @self.output_shapes. Leading dimensions will be batch and time [B, T, ...] + for each tensor. + + rnn_state (torch.Tensor or tuple): return the new rnn state (if @return_state) + """ + for obs_group in self.input_obs_group_shapes: + for k in self.input_obs_group_shapes[obs_group]: + # first two dimensions should be [B, T] for inputs + assert inputs[obs_group][k].ndim - 2 == len(self.input_obs_group_shapes[obs_group][k]) + + # use encoder to extract flat rnn inputs + rnn_inputs = TensorUtils.time_distributed(inputs, self.nets["encoder"], inputs_as_kwargs=True) + assert rnn_inputs.ndim == 3 # [B, T, D] + if self.per_step: + return self.nets["rnn"].forward(inputs=rnn_inputs, rnn_init_state=rnn_init_state, return_state=return_state) + + # apply MLP + decoder to last RNN output + outputs = self.nets["rnn"].forward(inputs=rnn_inputs, rnn_init_state=rnn_init_state, return_state=return_state) + if return_state: + outputs, rnn_state = outputs + + assert outputs.ndim == 3 # [B, T, D] + if self._has_mlp: + outputs = self.nets["decoder"](self.nets["mlp"](outputs[:, -1])) + else: + outputs = self.nets["decoder"](outputs[:, -1]) + + if return_state: + return outputs, rnn_state + return outputs + + def forward_step(self, rnn_state, **inputs): + """ + Unroll network over a single timestep. + + Args: + inputs (dict): expects same modalities as @self.input_shapes, with + additional batch dimension (but NOT time), since this is a + single time step. + + rnn_state (torch.Tensor): rnn hidden state + + Returns: + outputs (dict): dictionary of output torch.Tensors, that corresponds + to @self.output_shapes. Does not contain time dimension. + + rnn_state: return the new rnn state + """ + # ensure that the only extra dimension is batch dim, not temporal dim + assert np.all([inputs[k].ndim - 1 == len(self.input_shapes[k]) for k in self.input_shapes]) + + inputs = TensorUtils.to_sequence(inputs) + outputs, rnn_state = self.forward( + inputs, + rnn_init_state=rnn_state, + return_state=True, + ) + if self.per_step: + # if outputs are not per-step, the time dimension is already reduced + outputs = outputs[:, 0] + return outputs, rnn_state + + def _to_string(self): + """ + Subclasses should override this method to print out info about network / policy. + """ + return '' + + def __repr__(self): + """Pretty print network.""" + header = '{}'.format(str(self.__class__.__name__)) + msg = '' + indent = ' ' * 4 + msg += textwrap.indent("\n" + self._to_string(), indent) + msg += textwrap.indent("\n\nencoder={}".format(self.nets["encoder"]), indent) + msg += textwrap.indent("\n\nrnn={}".format(self.nets["rnn"]), indent) + msg = header + '(' + msg + '\n)' + return msg + + +class MIMO_Transformer(Module): + """ + Extension to Transformer (based on GPT architecture) to accept multiple observation + dictionaries as input and to output dictionaries of tensors. Inputs are specified as + a dictionary of observation dictionaries, with each key corresponding to an observation group. + This module utilizes @ObservationGroupEncoder to process the multiple input dictionaries and + @ObservationDecoder to generate tensor dictionaries. The default behavior + for encoding the inputs is to process visual inputs with a learned CNN and concatenating + the flat encodings with the other flat inputs. The default behavior for generating + outputs is to use a linear layer branch to produce each modality separately + (including visual outputs). + """ + def __init__( + self, + input_obs_group_shapes, + output_shapes, + transformer_embed_dim, + transformer_num_layers, + transformer_num_heads, + transformer_context_length, + transformer_emb_dropout=0.1, + transformer_attn_dropout=0.1, + transformer_block_output_dropout=0.1, + transformer_sinusoidal_embedding=False, + transformer_activation="gelu", + transformer_nn_parameter_for_timesteps=False, + encoder_kwargs=None, + ): + """ + Args: + input_obs_group_shapes (OrderedDict): a dictionary of dictionaries. + Each key in this dictionary should specify an observation group, and + the value should be an OrderedDict that maps modalities to + expected shapes. + output_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for outputs. + transformer_embed_dim (int): dimension for embeddings used by transformer + transformer_num_layers (int): number of transformer blocks to stack + transformer_num_heads (int): number of attention heads for each + transformer block - must divide @transformer_embed_dim evenly. Self-attention is + computed over this many partitions of the embedding dimension separately. + transformer_context_length (int): expected length of input sequences + transformer_activation: non-linearity for input and output layers used in transformer + transformer_emb_dropout (float): dropout probability for embedding inputs in transformer + transformer_attn_dropout (float): dropout probability for attention outputs for each transformer block + transformer_block_output_dropout (float): dropout probability for final outputs for each transformer block + encoder_kwargs (dict): observation encoder config + """ + super(MIMO_Transformer, self).__init__() + + assert isinstance(input_obs_group_shapes, OrderedDict) + assert np.all([isinstance(input_obs_group_shapes[k], OrderedDict) for k in input_obs_group_shapes]) + assert isinstance(output_shapes, OrderedDict) + + self.input_obs_group_shapes = input_obs_group_shapes + self.output_shapes = output_shapes + + self.nets = nn.ModuleDict() + self.params = nn.ParameterDict() + + # Encoder for all observation groups. + self.nets["encoder"] = ObservationGroupEncoder( + observation_group_shapes=input_obs_group_shapes, + encoder_kwargs=encoder_kwargs, + feature_activation=None, + ) + + # flat encoder output dimension + transformer_input_dim = self.nets["encoder"].output_shape()[0] + + self.nets["embed_encoder"] = nn.Linear( + transformer_input_dim, transformer_embed_dim + ) + + max_timestep = transformer_context_length + + if transformer_sinusoidal_embedding: + self.nets["embed_timestep"] = PositionalEncoding(transformer_embed_dim) + elif transformer_nn_parameter_for_timesteps: + assert ( + not transformer_sinusoidal_embedding + ), "nn.Parameter only works with learned embeddings" + self.params["embed_timestep"] = nn.Parameter( + torch.zeros(1, max_timestep, transformer_embed_dim) + ) + else: + self.nets["embed_timestep"] = nn.Embedding(max_timestep, transformer_embed_dim) + + # layer norm for embeddings + self.nets["embed_ln"] = nn.LayerNorm(transformer_embed_dim) + + # dropout for input embeddings + self.nets["embed_drop"] = nn.Dropout(transformer_emb_dropout) + + # GPT transformer + self.nets["transformer"] = GPT_Backbone( + embed_dim=transformer_embed_dim, + num_layers=transformer_num_layers, + num_heads=transformer_num_heads, + context_length=transformer_context_length, + attn_dropout=transformer_attn_dropout, + block_output_dropout=transformer_block_output_dropout, + activation=transformer_activation, + ) + + # decoder for output modalities + self.nets["decoder"] = ObservationDecoder( + decode_shapes=self.output_shapes, + input_feat_dim=transformer_embed_dim, + ) + + self.transformer_context_length = transformer_context_length + self.transformer_embed_dim = transformer_embed_dim + self.transformer_sinusoidal_embedding = transformer_sinusoidal_embedding + self.transformer_nn_parameter_for_timesteps = transformer_nn_parameter_for_timesteps + + def output_shape(self, input_shape=None): + """ + Returns output shape for this module, which is a dictionary instead + of a list since outputs are dictionaries. + """ + return { k : list(self.output_shapes[k]) for k in self.output_shapes } + + def embed_timesteps(self, embeddings): + """ + Computes timestep-based embeddings (aka positional embeddings) to add to embeddings. + Args: + embeddings (torch.Tensor): embeddings prior to positional embeddings are computed + Returns: + time_embeddings (torch.Tensor): positional embeddings to add to embeddings + """ + timesteps = ( + torch.arange( + 0, + embeddings.shape[1], + dtype=embeddings.dtype, + device=embeddings.device, + ) + .unsqueeze(0) + .repeat(embeddings.shape[0], 1) + ) + assert (timesteps >= 0.0).all(), "timesteps must be positive!" + if self.transformer_sinusoidal_embedding: + assert torch.is_floating_point(timesteps), timesteps.dtype + else: + timesteps = timesteps.long() + + if self.transformer_nn_parameter_for_timesteps: + time_embeddings = self.params["embed_timestep"] + else: + time_embeddings = self.nets["embed_timestep"]( + timesteps + ) # these are NOT fed into transformer, only added to the inputs. + # compute how many modalities were combined into embeddings, replicate time embeddings that many times + num_replicates = embeddings.shape[-1] // self.transformer_embed_dim + time_embeddings = torch.cat([time_embeddings for _ in range(num_replicates)], -1) + assert ( + embeddings.shape == time_embeddings.shape + ), f"{embeddings.shape}, {time_embeddings.shape}" + return time_embeddings + + def input_embedding( + self, + inputs, + ): + """ + Process encoded observations into embeddings to pass to transformer, + Adds timestep-based embeddings (aka positional embeddings) to inputs. + Args: + inputs (torch.Tensor): outputs from observation encoder + Returns: + embeddings (torch.Tensor): input embeddings to pass to transformer backbone. + """ + embeddings = self.nets["embed_encoder"](inputs) + time_embeddings = self.embed_timesteps(embeddings) + embeddings = embeddings + time_embeddings + embeddings = self.nets["embed_ln"](embeddings) + embeddings = self.nets["embed_drop"](embeddings) + + return embeddings + + + def forward(self, **inputs): + """ + Process each set of inputs in its own observation group. + Args: + inputs (dict): a dictionary of dictionaries with one dictionary per + observation group. Each observation group's dictionary should map + modality to torch.Tensor batches. Should be consistent with + @self.input_obs_group_shapes. First two leading dimensions should + be batch and time [B, T, ...] for each tensor. + Returns: + outputs (dict): dictionary of output torch.Tensors, that corresponds + to @self.output_shapes. Leading dimensions will be batch and time [B, T, ...] + for each tensor. + """ + for obs_group in self.input_obs_group_shapes: + for k in self.input_obs_group_shapes[obs_group]: + # first two dimensions should be [B, T] for inputs + if inputs[obs_group][k] is None: + continue + assert inputs[obs_group][k].ndim - 2 == len(self.input_obs_group_shapes[obs_group][k]) + + inputs = inputs.copy() + + transformer_encoder_outputs = None + transformer_inputs = TensorUtils.time_distributed( + inputs, self.nets["encoder"], inputs_as_kwargs=True + ) + assert transformer_inputs.ndim == 3 # [B, T, D] + + if transformer_encoder_outputs is None: + transformer_embeddings = self.input_embedding(transformer_inputs) + # pass encoded sequences through transformer + transformer_encoder_outputs = self.nets["transformer"].forward(transformer_embeddings) + + transformer_outputs = transformer_encoder_outputs + # apply decoder to each timestep of sequence to get a dictionary of outputs + transformer_outputs = TensorUtils.time_distributed( + transformer_outputs, self.nets["decoder"] + ) + transformer_outputs["transformer_encoder_outputs"] = transformer_encoder_outputs + return transformer_outputs + + def _to_string(self): + """ + Subclasses should override this method to print out info about network / policy. + """ + return '' + + def __repr__(self): + """Pretty print network.""" + header = '{}'.format(str(self.__class__.__name__)) + msg = '' + indent = ' ' * 4 + if self._to_string() != '': + msg += textwrap.indent("\n" + self._to_string() + "\n", indent) + msg += textwrap.indent("\nencoder={}".format(self.nets["encoder"]), indent) + msg += textwrap.indent("\n\ntransformer={}".format(self.nets["transformer"]), indent) + msg += textwrap.indent("\n\ndecoder={}".format(self.nets["decoder"]), indent) + msg = header + '(' + msg + '\n)' + return msg \ No newline at end of file diff --git a/aloha-devel/robomimic/models/value_nets.py b/aloha-devel/robomimic/models/value_nets.py new file mode 100644 index 0000000000000000000000000000000000000000..c98fa7e4e0f4185b2a11e5581158268aa9cda2cf --- /dev/null +++ b/aloha-devel/robomimic/models/value_nets.py @@ -0,0 +1,318 @@ +""" +Contains torch Modules for value networks. These networks take an +observation dictionary as input (and possibly additional conditioning, +such as subgoal or goal dictionaries) and produce value or +action-value estimates or distributions. +""" +import numpy as np +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.distributions as D + +import robomimic.utils.tensor_utils as TensorUtils +from robomimic.models.obs_nets import MIMO_MLP +from robomimic.models.distributions import DiscreteValueDistribution + + +class ValueNetwork(MIMO_MLP): + """ + A basic value network that predicts values from observations. + Can optionally be goal conditioned on future observations. + """ + def __init__( + self, + obs_shapes, + mlp_layer_dims, + value_bounds=None, + goal_shapes=None, + encoder_kwargs=None, + ): + """ + Args: + obs_shapes (OrderedDict): a dictionary that maps observation keys to + expected shapes for observations. + + mlp_layer_dims ([int]): sequence of integers for the MLP hidden layers sizes. + + value_bounds (tuple): a 2-tuple corresponding to the lowest and highest possible return + that the network should be possible of generating. The network will rescale outputs + using a tanh layer to lie within these bounds. If None, no tanh re-scaling is done. + + goal_shapes (OrderedDict): a dictionary that maps observation keys to + expected shapes for goal observations. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-observation key information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + self.value_bounds = value_bounds + if self.value_bounds is not None: + # convert [lb, ub] to a scale and offset for the tanh output, which is in [-1, 1] + self._value_scale = (float(self.value_bounds[1]) - float(self.value_bounds[0])) / 2. + self._value_offset = (float(self.value_bounds[1]) + float(self.value_bounds[0])) / 2. + + assert isinstance(obs_shapes, OrderedDict) + self.obs_shapes = obs_shapes + + # set up different observation groups for @MIMO_MLP + observation_group_shapes = OrderedDict() + observation_group_shapes["obs"] = OrderedDict(self.obs_shapes) + + self._is_goal_conditioned = False + if goal_shapes is not None and len(goal_shapes) > 0: + assert isinstance(goal_shapes, OrderedDict) + self._is_goal_conditioned = True + self.goal_shapes = OrderedDict(goal_shapes) + observation_group_shapes["goal"] = OrderedDict(self.goal_shapes) + else: + self.goal_shapes = OrderedDict() + + output_shapes = self._get_output_shapes() + super(ValueNetwork, self).__init__( + input_obs_group_shapes=observation_group_shapes, + output_shapes=output_shapes, + layer_dims=mlp_layer_dims, + encoder_kwargs=encoder_kwargs, + ) + + def _get_output_shapes(self): + """ + Allow subclasses to re-define outputs from @MIMO_MLP, since we won't + always directly predict values, but may instead predict the parameters + of a value distribution. + """ + return OrderedDict(value=(1,)) + + def output_shape(self, input_shape=None): + """ + Function to compute output shape from inputs to this module. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + return [1] + + def forward(self, obs_dict, goal_dict=None): + """ + Forward through value network, and then optionally use tanh scaling. + """ + values = super(ValueNetwork, self).forward(obs=obs_dict, goal=goal_dict)["value"] + if self.value_bounds is not None: + values = self._value_offset + self._value_scale * torch.tanh(values) + return values + + def _to_string(self): + return "value_bounds={}".format(self.value_bounds) + + +class ActionValueNetwork(ValueNetwork): + """ + A basic Q (action-value) network that predicts values from observations + and actions. Can optionally be goal conditioned on future observations. + """ + def __init__( + self, + obs_shapes, + ac_dim, + mlp_layer_dims, + value_bounds=None, + goal_shapes=None, + encoder_kwargs=None, + ): + """ + Args: + obs_shapes (OrderedDict): a dictionary that maps observation keys to + expected shapes for observations. + + ac_dim (int): dimension of action space. + + mlp_layer_dims ([int]): sequence of integers for the MLP hidden layers sizes. + + value_bounds (tuple): a 2-tuple corresponding to the lowest and highest possible return + that the network should be possible of generating. The network will rescale outputs + using a tanh layer to lie within these bounds. If None, no tanh re-scaling is done. + + goal_shapes (OrderedDict): a dictionary that maps observation keys to + expected shapes for goal observations. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-observation key information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + + # add in action as a modality + new_obs_shapes = OrderedDict(obs_shapes) + new_obs_shapes["action"] = (ac_dim,) + self.ac_dim = ac_dim + + # pass to super class to instantiate network + super(ActionValueNetwork, self).__init__( + obs_shapes=new_obs_shapes, + mlp_layer_dims=mlp_layer_dims, + value_bounds=value_bounds, + goal_shapes=goal_shapes, + encoder_kwargs=encoder_kwargs, + ) + + def forward(self, obs_dict, acts, goal_dict=None): + """ + Modify forward from super class to include actions in inputs. + """ + inputs = dict(obs_dict) + inputs["action"] = acts + return super(ActionValueNetwork, self).forward(inputs, goal_dict) + + def _to_string(self): + return "action_dim={}\nvalue_bounds={}".format(self.ac_dim, self.value_bounds) + + +class DistributionalActionValueNetwork(ActionValueNetwork): + """ + Distributional Q (action-value) network that outputs a categorical distribution over + a discrete grid of value atoms. See https://arxiv.org/pdf/1707.06887.pdf for + more details. + """ + def __init__( + self, + obs_shapes, + ac_dim, + mlp_layer_dims, + value_bounds, + num_atoms, + goal_shapes=None, + encoder_kwargs=None, + ): + """ + Args: + obs_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for observations. + + ac_dim (int): dimension of action space. + + mlp_layer_dims ([int]): sequence of integers for the MLP hidden layers sizes. + + value_bounds (tuple): a 2-tuple corresponding to the lowest and highest possible return + that the network should be possible of generating. This defines the support + of the value distribution. + + num_atoms (int): number of value atoms to use for the categorical distribution - which + is the representation of the value distribution. + + goal_shapes (OrderedDict): a dictionary that maps modality to + expected shapes for goal observations. + + encoder_kwargs (dict or None): If None, results in default encoder_kwargs being applied. Otherwise, should + be nested dictionary containing relevant per-modality information for encoder networks. + Should be of form: + + obs_modality1: dict + feature_dimension: int + core_class: str + core_kwargs: dict + ... + ... + obs_randomizer_class: str + obs_randomizer_kwargs: dict + ... + ... + obs_modality2: dict + ... + """ + + # parameters specific to DistributionalActionValueNetwork + self.num_atoms = num_atoms + self._atoms = np.linspace(value_bounds[0], value_bounds[1], num_atoms) + + # pass to super class to instantiate network + super(DistributionalActionValueNetwork, self).__init__( + obs_shapes=obs_shapes, + ac_dim=ac_dim, + mlp_layer_dims=mlp_layer_dims, + value_bounds=value_bounds, + goal_shapes=goal_shapes, + encoder_kwargs=encoder_kwargs, + ) + + def _get_output_shapes(self): + """ + Network outputs log probabilities for categorical distribution over discrete value grid. + """ + return OrderedDict(log_probs=(self.num_atoms,)) + + def forward_train(self, obs_dict, acts, goal_dict=None): + """ + Return full critic categorical distribution. + + Args: + obs_dict (dict): batch of observations + acts (torch.Tensor): batch of actions + goal_dict (dict): if not None, batch of goal observations + + Returns: + value_distribution (DiscreteValueDistribution instance) + """ + + # add in actions + inputs = dict(obs_dict) + inputs["action"] = acts + + # network returns unnormalized log probabilities (logits) for each of the value atoms + logits = MIMO_MLP.forward(self, obs=inputs, goal=goal_dict)["log_probs"] + + # turn these logits into a categorical distribution over the value atoms. + # (unsqueeze to make sure atoms are compatible with batch operations) + value_atoms = torch.Tensor(self._atoms).unsqueeze(0).to(logits.device) + return DiscreteValueDistribution(values=value_atoms, logits=logits) + + def forward(self, obs_dict, acts, goal_dict=None): + """ + Return mean of critic categorical distribution. Useful for obtaining + point estimates of critic values. + + Args: + obs_dict (dict): batch of observations + acts (torch.Tensor): batch of actions + goal_dict (dict): if not None, batch of goal observations + + Returns: + mean_value (torch.Tensor): expectation of value distribution + """ + vd = self.forward_train(obs_dict=obs_dict, acts=acts, goal_dict=goal_dict) + return vd.mean() + + def _to_string(self): + return "action_dim={}\nvalue_bounds={}\nnum_atoms={}".format(self.ac_dim, self.value_bounds, self.num_atoms) \ No newline at end of file diff --git a/aloha-devel/robomimic/utils/__init__.py b/aloha-devel/robomimic/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/aloha-devel/robomimic/utils/__pycache__/__init__.cpython-38.pyc b/aloha-devel/robomimic/utils/__pycache__/__init__.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5345e80043eb0dce3f2576a62f151ec8f7ca3545 Binary files /dev/null and b/aloha-devel/robomimic/utils/__pycache__/__init__.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/utils/__pycache__/tensor_utils.cpython-38.pyc b/aloha-devel/robomimic/utils/__pycache__/tensor_utils.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7aa13fe2bec07006bb2ea884d63ed3c668f7e182 Binary files /dev/null and b/aloha-devel/robomimic/utils/__pycache__/tensor_utils.cpython-38.pyc differ diff --git a/aloha-devel/robomimic/utils/action_utils.py b/aloha-devel/robomimic/utils/action_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ac974d50c5fab46b85cd8c3bb76d8e05a4f56aba --- /dev/null +++ b/aloha-devel/robomimic/utils/action_utils.py @@ -0,0 +1,35 @@ + +from typing import Union, Sequence, Dict, Optional, Tuple + +from copy import deepcopy +from collections import OrderedDict +import functools + +import numpy as np + + +def action_dict_to_vector( + action_dict: Dict[str, np.ndarray], + action_keys: Optional[Sequence[str]]=None) -> np.ndarray: + if action_keys is None: + action_keys = list(action_dict.keys()) + actions = [action_dict[k] for k in action_keys] + + action_vec = np.concatenate(actions, axis=-1) + return action_vec + + +def vector_to_action_dict( + action: np.ndarray, + action_shapes: Dict[str, Tuple[int]], + action_keys: Sequence[str]) -> Dict[str, np.ndarray]: + action_dict = dict() + start_idx = 0 + for key in action_keys: + this_act_shape = action_shapes[key] + this_act_dim = np.prod(this_act_shape) + end_idx = start_idx + this_act_dim + action_dict[key] = action[...,start_idx:end_idx].reshape( + action.shape[:-1]+this_act_shape) + start_idx = end_idx + return action_dict diff --git a/aloha-devel/robomimic/utils/dataset.py b/aloha-devel/robomimic/utils/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..15ab22641ad71d5dc00027bb395de27d46ab30ce --- /dev/null +++ b/aloha-devel/robomimic/utils/dataset.py @@ -0,0 +1,1158 @@ +""" +This file contains Dataset classes that are used by torch dataloaders +to fetch batches from hdf5 files. +""" +import os +import h5py +import numpy as np +import random +from copy import deepcopy +from contextlib import contextmanager +from collections import OrderedDict + +import torch.utils.data + +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.obs_utils as ObsUtils +import robomimic.utils.action_utils as AcUtils +import robomimic.utils.log_utils as LogUtils +import robomimic.utils.lang_utils as LangUtils + + +class SequenceDataset(torch.utils.data.Dataset): + def __init__( + self, + hdf5_path, + obs_keys, + action_keys, + dataset_keys, + action_config, + frame_stack=1, + seq_length=1, + pad_frame_stack=True, + pad_seq_length=True, + get_pad_mask=False, + goal_mode=None, + hdf5_cache_mode=None, + hdf5_use_swmr=True, + hdf5_normalize_obs=False, + filter_by_attribute=None, + load_next_obs=True, + shuffled_obs_key_groups=None, + lang=None, + ): + """ + Dataset class for fetching sequences of experience. + Length of the fetched sequence is equal to (@frame_stack - 1 + @seq_length) + + Args: + hdf5_path (str): path to hdf5 + + obs_keys (tuple, list): keys to observation items (image, object, etc) to be fetched from the dataset + + action_config (dict): TODO + + dataset_keys (tuple, list): keys to dataset items (actions, rewards, etc) to be fetched from the dataset + + frame_stack (int): numbers of stacked frames to fetch. Defaults to 1 (single frame). + + seq_length (int): length of sequences to sample. Defaults to 1 (single frame). + + pad_frame_stack (int): whether to pad sequence for frame stacking at the beginning of a demo. This + ensures that partial frame stacks are observed, such as (s_0, s_0, s_0, s_1). Otherwise, the + first frame stacked observation would be (s_0, s_1, s_2, s_3). + + pad_seq_length (int): whether to pad sequence for sequence fetching at the end of a demo. This + ensures that partial sequences at the end of a demonstration are observed, such as + (s_{T-1}, s_{T}, s_{T}, s_{T}). Otherwise, the last sequence provided would be + (s_{T-3}, s_{T-2}, s_{T-1}, s_{T}). + + get_pad_mask (bool): if True, also provide padding masks as part of the batch. This can be + useful for masking loss functions on padded parts of the data. + + goal_mode (str): either "last" or None. Defaults to None, which is to not fetch goals + + hdf5_cache_mode (str): one of ["all", "low_dim", or None]. Set to "all" to cache entire hdf5 + in memory - this is by far the fastest for data loading. Set to "low_dim" to cache all + non-image data. Set to None to use no caching - in this case, every batch sample is + retrieved via file i/o. You should almost never set this to None, even for large + image datasets. + + hdf5_use_swmr (bool): whether to use swmr feature when opening the hdf5 file. This ensures + that multiple Dataset instances can all access the same hdf5 file without problems. + + hdf5_normalize_obs (bool): if True, normalize observations by computing the mean observation + and std of each observation (in each dimension and modality), and normalizing to unit + mean and variance in each dimension. + + filter_by_attribute (str): if provided, use the provided filter key to look up a subset of + demonstrations to load + + load_next_obs (bool): whether to load next_obs from the dataset + + shuffled_obs_key_groups (list): TODO + + lang: TODO documentation + """ + super(SequenceDataset, self).__init__() + + self.hdf5_path = os.path.expanduser(hdf5_path) + self.hdf5_use_swmr = hdf5_use_swmr + self.hdf5_normalize_obs = hdf5_normalize_obs + self._hdf5_file = None + + assert hdf5_cache_mode in ["all", "low_dim", None] + self.hdf5_cache_mode = hdf5_cache_mode + + self.load_next_obs = load_next_obs + self.filter_by_attribute = filter_by_attribute + + # get all keys that needs to be fetched + self.obs_keys = tuple(obs_keys) + self.action_keys = tuple(action_keys) + self.dataset_keys = tuple(dataset_keys) + # add action keys to dataset keys + if self.action_keys is not None: + self.dataset_keys = tuple(set(self.dataset_keys).union(set(self.action_keys))) + + self.action_config = action_config + + # set up lang and language embedding + self.lang = lang + self._lang_emb = LangUtils.get_lang_emb(self.lang) + + self.n_frame_stack = frame_stack + assert self.n_frame_stack >= 1 + + self.seq_length = seq_length + assert self.seq_length >= 1 + + self.goal_mode = goal_mode + if self.goal_mode is not None: + assert self.goal_mode in ["last"] + if not self.load_next_obs: + assert self.goal_mode != "last" # we use last next_obs as goal + + self.pad_seq_length = pad_seq_length + self.pad_frame_stack = pad_frame_stack + self.get_pad_mask = get_pad_mask + + self.load_demo_info(filter_by_attribute=self.filter_by_attribute) + + # maybe prepare for observation normalization + self.obs_normalization_stats = None + if self.hdf5_normalize_obs: + self.obs_normalization_stats = self.normalize_obs() + + # prepare for action normalization + self.action_normalization_stats = None + + # maybe store dataset in memory for fast access + if self.hdf5_cache_mode in ["all", "low_dim"]: + obs_keys_in_memory = self.obs_keys + if self.hdf5_cache_mode == "low_dim": + # only store low-dim observations + obs_keys_in_memory = [] + for k in self.obs_keys: + if ObsUtils.key_is_obs_modality(k, "low_dim"): + obs_keys_in_memory.append(k) + self.obs_keys_in_memory = obs_keys_in_memory + + self.hdf5_cache = self.load_dataset_in_memory( + demo_list=self.demos, + hdf5_file=self.hdf5_file, + obs_keys=self.obs_keys_in_memory, + dataset_keys=self.dataset_keys, + load_next_obs=self.load_next_obs + ) + + if self.hdf5_cache_mode == "all": + # cache getitem calls for even more speedup. We don't do this for + # "low-dim" since image observations require calls to getitem anyways. + print("SequenceDataset: caching get_item calls...") + self.getitem_cache = [self.get_item(i) for i in LogUtils.custom_tqdm(range(len(self)))] + + # don't need the previous cache anymore + del self.hdf5_cache + self.hdf5_cache = None + else: + self.hdf5_cache = None + + if shuffled_obs_key_groups is None: + self.shuffled_obs_key_groups = list() + else: + self.shuffled_obs_key_groups = shuffled_obs_key_groups + + self.close_and_delete_hdf5_handle() + + def load_demo_info(self, filter_by_attribute=None, demos=None): + """ + Args: + filter_by_attribute (str): if provided, use the provided filter key + to select a subset of demonstration trajectories to load + + demos (list): list of demonstration keys to load from the hdf5 file. If + omitted, all demos in the file (or under the @filter_by_attribute + filter key) are used. + """ + # filter demo trajectory by mask + if demos is not None: + self.demos = demos + elif filter_by_attribute is not None: + self.demos = [elem.decode("utf-8") for elem in np.array(self.hdf5_file["mask/{}".format(filter_by_attribute)][:])] + else: + self.demos = list(self.hdf5_file["data"].keys()) + + # sort demo keys + inds = np.argsort([int(elem[5:]) for elem in self.demos]) + self.demos = [self.demos[i] for i in inds] + + self.n_demos = len(self.demos) + + # keep internal index maps to know which transitions belong to which demos + self._index_to_demo_id = dict() # maps every index to a demo id + self._demo_id_to_start_indices = dict() # gives start index per demo id + self._demo_id_to_demo_length = dict() + + # determine index mapping + self.total_num_sequences = 0 + for ep in self.demos: + demo_length = self.hdf5_file["data/{}".format(ep)].attrs["num_samples"] + self._demo_id_to_start_indices[ep] = self.total_num_sequences + self._demo_id_to_demo_length[ep] = demo_length + + num_sequences = demo_length + # determine actual number of sequences taking into account whether to pad for frame_stack and seq_length + if not self.pad_frame_stack: + num_sequences -= (self.n_frame_stack - 1) + if not self.pad_seq_length: + num_sequences -= (self.seq_length - 1) + + if self.pad_seq_length: + assert demo_length >= 1 # sequence needs to have at least one sample + num_sequences = max(num_sequences, 1) + else: + assert num_sequences >= 1 # assume demo_length >= (self.n_frame_stack - 1 + self.seq_length) + + for _ in range(num_sequences): + self._index_to_demo_id[self.total_num_sequences] = ep + self.total_num_sequences += 1 + + @property + def hdf5_file(self): + """ + This property allows for a lazy hdf5 file open. + """ + if self._hdf5_file is None: + self._hdf5_file = h5py.File(self.hdf5_path, 'r', swmr=self.hdf5_use_swmr, libver='latest') + return self._hdf5_file + + def close_and_delete_hdf5_handle(self): + """ + Maybe close the file handle. + """ + if self._hdf5_file is not None: + self._hdf5_file.close() + self._hdf5_file = None + + @contextmanager + def hdf5_file_opened(self): + """ + Convenient context manager to open the file on entering the scope + and then close it on leaving. + """ + should_close = self._hdf5_file is None + yield self.hdf5_file + if should_close: + self.close_and_delete_hdf5_handle() + + def __del__(self): + self.close_and_delete_hdf5_handle() + + def __repr__(self): + """ + Pretty print the class and important attributes on a call to `print`. + """ + msg = str(self.__class__.__name__) + msg += " (\n\tpath={}\n\tobs_keys={}\n\tseq_length={}\n\tfilter_key={}\n\tframe_stack={}\n" + msg += "\tpad_seq_length={}\n\tpad_frame_stack={}\n\tgoal_mode={}\n" + msg += "\tcache_mode={}\n" + msg += "\tnum_demos={}\n\tnum_sequences={}\n)" + filter_key_str = self.filter_by_attribute if self.filter_by_attribute is not None else "none" + goal_mode_str = self.goal_mode if self.goal_mode is not None else "none" + cache_mode_str = self.hdf5_cache_mode if self.hdf5_cache_mode is not None else "none" + msg = msg.format(self.hdf5_path, self.obs_keys, self.seq_length, filter_key_str, self.n_frame_stack, + self.pad_seq_length, self.pad_frame_stack, goal_mode_str, cache_mode_str, + self.n_demos, self.total_num_sequences) + return msg + + def __len__(self): + """ + Ensure that the torch dataloader will do a complete pass through all sequences in + the dataset before starting a new iteration. + """ + return self.total_num_sequences + + def load_dataset_in_memory(self, demo_list, hdf5_file, obs_keys, dataset_keys, load_next_obs): + """ + Loads the hdf5 dataset into memory, preserving the structure of the file. Note that this + differs from `self.getitem_cache`, which, if active, actually caches the outputs of the + `getitem` operation. + + Args: + demo_list (list): list of demo keys, e.g., 'demo_0' + hdf5_file (h5py.File): file handle to the hdf5 dataset. + obs_keys (list, tuple): observation keys to fetch, e.g., 'images' + dataset_keys (list, tuple): dataset keys to fetch, e.g., 'actions' + load_next_obs (bool): whether to load next_obs from the dataset + + Returns: + all_data (dict): dictionary of loaded data. + """ + all_data = dict() + print("SequenceDataset: loading dataset into memory...") + for ep in LogUtils.custom_tqdm(demo_list): + all_data[ep] = {} + all_data[ep]["attrs"] = {} + all_data[ep]["attrs"]["num_samples"] = hdf5_file["data/{}".format(ep)].attrs["num_samples"] + # get obs + all_data[ep]["obs"] = {k: hdf5_file["data/{}/obs/{}".format(ep, k)][()] for k in obs_keys} + if load_next_obs: + all_data[ep]["next_obs"] = {k: hdf5_file["data/{}/next_obs/{}".format(ep, k)][()] for k in obs_keys} + # get other dataset keys + for k in dataset_keys: + if k in hdf5_file["data/{}".format(ep)]: + all_data[ep][k] = hdf5_file["data/{}/{}".format(ep, k)][()].astype('float32') + else: + all_data[ep][k] = np.zeros((all_data[ep]["attrs"]["num_samples"], 1), dtype=np.float32) + + if "model_file" in hdf5_file["data/{}".format(ep)].attrs: + all_data[ep]["attrs"]["model_file"] = hdf5_file["data/{}".format(ep)].attrs["model_file"] + + return all_data + + def normalize_obs(self): + """ + Computes a dataset-wide mean and standard deviation for the observations + (per dimension and per obs key) and returns it. + """ + + # Run through all trajectories. For each one, compute minimal observation statistics, and then aggregate + # with the previous statistics. + ep = self.demos[0] + obs_traj = {k: self.hdf5_file["data/{}/obs/{}".format(ep, k)][()].astype('float32') for k in self.obs_keys} + obs_traj = ObsUtils.process_obs_dict(obs_traj) + merged_stats = _compute_traj_stats(obs_traj) + print("SequenceDataset: normalizing observations...") + for ep in LogUtils.custom_tqdm(self.demos[1:]): + obs_traj = {k: self.hdf5_file["data/{}/obs/{}".format(ep, k)][()].astype('float32') for k in self.obs_keys} + obs_traj = ObsUtils.process_obs_dict(obs_traj) + traj_stats = _compute_traj_stats(obs_traj) + merged_stats = _aggregate_traj_stats(merged_stats, traj_stats) + + obs_normalization_stats = { k : {} for k in merged_stats } + for k in merged_stats: + # note we add a small tolerance of 1e-3 for std + obs_normalization_stats[k]["mean"] = merged_stats[k]["mean"] + obs_normalization_stats[k]["std"] = np.sqrt(merged_stats[k]["sqdiff"] / merged_stats[k]["n"]) + 1e-3 + return obs_normalization_stats + + def get_obs_normalization_stats(self): + """ + Returns dictionary of mean and std for each observation key if using + observation normalization, otherwise None. + + Returns: + obs_normalization_stats (dict): a dictionary for observation + normalization. This maps observation keys to dicts + with a "mean" and "std" of shape (1, ...) where ... is the default + shape for the observation. + """ + assert self.hdf5_normalize_obs, "not using observation normalization!" + return deepcopy(self.obs_normalization_stats) + + def get_action_traj(self, ep): + action_traj = dict() + for key in self.action_keys: + action_traj[key] = self.hdf5_file["data/{}/{}".format(ep, key)][()].astype('float32') + return action_traj + + def get_action_stats(self): + ep = self.demos[0] + action_traj = self.get_action_traj(ep) + action_stats = _compute_traj_stats(action_traj) + print("SequenceDataset: normalizing actions...") + for ep in LogUtils.custom_tqdm(self.demos[1:]): + action_traj = self.get_action_traj(ep) + traj_stats = _compute_traj_stats(action_traj) + action_stats = _aggregate_traj_stats(action_stats, traj_stats) + return action_stats + + def set_action_normalization_stats(self, action_normalization_stats): + self.action_normalization_stats = action_normalization_stats + + def get_action_normalization_stats(self): + """ + Computes a dataset-wide min, max, mean and standard deviation for the actions + (per dimension) and returns it. + """ + + # Run through all trajectories. For each one, compute minimal observation statistics, and then aggregate + # with the previous statistics. + if self.action_normalization_stats is None: + action_stats = self.get_action_stats() + self.action_normalization_stats = action_stats_to_normalization_stats( + action_stats, self.action_config) + return self.action_normalization_stats + + def get_dataset_for_ep(self, ep, key): + """ + Helper utility to get a dataset for a specific demonstration. + Takes into account whether the dataset has been loaded into memory. + """ + + # check if this key should be in memory + key_should_be_in_memory = (self.hdf5_cache_mode in ["all", "low_dim"]) + if key_should_be_in_memory: + # if key is an observation, it may not be in memory + if '/' in key: + key1, key2 = key.split('/') + assert(key1 in ['obs', 'next_obs', 'action_dict']) + if key2 not in self.obs_keys_in_memory: + key_should_be_in_memory = False + + if key_should_be_in_memory: + # read cache + if '/' in key: + key1, key2 = key.split('/') + assert(key1 in ['obs', 'next_obs', 'action_dict']) + ret = self.hdf5_cache[ep][key1][key2] + else: + ret = self.hdf5_cache[ep][key] + else: + # read from file + hd5key = "data/{}/{}".format(ep, key) + ret = self.hdf5_file[hd5key] + return ret + + def __getitem__(self, index): + """ + Fetch dataset sequence @index (inferred through internal index map), using the getitem_cache if available. + """ + if self.hdf5_cache_mode == "all": + output = self.getitem_cache[index] + else: + output = self.get_item(index) + + for (g1, g2) in self.shuffled_obs_key_groups: + assert len(g1) == len(g2) + if random.random() > 0.5: + # shuffle the keys accordingly + for (o1, o2) in zip(g1, g2): + for otype in ["obs", "next_obs", "goal_obs"]: + if output.get(otype, None) is None: + continue + if o1 not in output[otype] or o2 not in output[otype]: + continue + # swap values + output[otype][o1], output[otype][o2] = output[otype][o2], output[otype][o1] + + return output + + def get_item(self, index): + """ + Main implementation of getitem when not using cache. + """ + + demo_id = self._index_to_demo_id[index] + demo_start_index = self._demo_id_to_start_indices[demo_id] + demo_length = self._demo_id_to_demo_length[demo_id] + + # start at offset index if not padding for frame stacking + demo_index_offset = 0 if self.pad_frame_stack else (self.n_frame_stack - 1) + index_in_demo = index - demo_start_index + demo_index_offset + + # end at offset index if not padding for seq length + demo_length_offset = 0 if self.pad_seq_length else (self.seq_length - 1) + end_index_in_demo = demo_length - demo_length_offset + + meta = self.get_dataset_sequence_from_demo( + demo_id, + index_in_demo=index_in_demo, + keys=self.dataset_keys, + num_frames_to_stack=self.n_frame_stack - 1, # note: need to decrement self.n_frame_stack by one + seq_length=self.seq_length + ) + + # determine goal index + goal_index = None + if self.goal_mode == "last": + goal_index = end_index_in_demo - 1 + + meta["obs"] = self.get_obs_sequence_from_demo( + demo_id, + index_in_demo=index_in_demo, + keys=self.obs_keys, + num_frames_to_stack=self.n_frame_stack - 1, + seq_length=self.seq_length, + prefix="obs" + ) + + if self.load_next_obs: + meta["next_obs"] = self.get_obs_sequence_from_demo( + demo_id, + index_in_demo=index_in_demo, + keys=self.obs_keys, + num_frames_to_stack=self.n_frame_stack - 1, + seq_length=self.seq_length, + prefix="next_obs" + ) + + if goal_index is not None: + goal = self.get_obs_sequence_from_demo( + demo_id, + index_in_demo=goal_index, + keys=self.obs_keys, + num_frames_to_stack=0, + seq_length=1, + prefix="next_obs", + ) + meta["goal_obs"] = {k: goal[k][0] for k in goal} # remove sequence dimension for goal + + # get action components + ac_dict = OrderedDict() + for k in self.action_keys: + ac = meta[k] + # expand action shape if needed + if len(ac.shape) == 1: + ac = ac.reshape(-1, 1) + ac_dict[k] = ac + + # normalize actions + action_normalization_stats = self.get_action_normalization_stats() + ac_dict = ObsUtils.normalize_dict(ac_dict, normalization_stats=action_normalization_stats) + + # concatenate all action components + meta["actions"] = AcUtils.action_dict_to_vector(ac_dict) + + # also return the sampled index + meta["index"] = index + + # language embedding + T = meta["actions"].shape[0] + meta["obs"]["lang_emb"] = np.tile(self._lang_emb, (T, 1)) + + return meta + + def get_sequence_from_demo(self, demo_id, index_in_demo, keys, num_frames_to_stack=0, seq_length=1): + """ + Extract a (sub)sequence of data items from a demo given the @keys of the items. + + Args: + demo_id (str): id of the demo, e.g., demo_0 + index_in_demo (int): beginning index of the sequence wrt the demo + keys (tuple): list of keys to extract + num_frames_to_stack (int): numbers of frame to stack. Seq gets prepended with repeated items if out of range + seq_length (int): sequence length to extract. Seq gets post-pended with repeated items if out of range + + Returns: + a dictionary of extracted items. + """ + assert num_frames_to_stack >= 0 + assert seq_length >= 1 + + demo_length = self._demo_id_to_demo_length[demo_id] + assert index_in_demo < demo_length + + # determine begin and end of sequence + seq_begin_index = max(0, index_in_demo - num_frames_to_stack) + seq_end_index = min(demo_length, index_in_demo + seq_length) + + # determine sequence padding + seq_begin_pad = max(0, num_frames_to_stack - index_in_demo) # pad for frame stacking + seq_end_pad = max(0, index_in_demo + seq_length - demo_length) # pad for sequence length + + # make sure we are not padding if specified. + if not self.pad_frame_stack: + assert seq_begin_pad == 0 + if not self.pad_seq_length: + assert seq_end_pad == 0 + + # fetch observation from the dataset file + seq = dict() + for k in keys: + data = self.get_dataset_for_ep(demo_id, k) + seq[k] = data[seq_begin_index: seq_end_index] + + seq = TensorUtils.pad_sequence(seq, padding=(seq_begin_pad, seq_end_pad), pad_same=True) + pad_mask = np.array([0] * seq_begin_pad + [1] * (seq_end_index - seq_begin_index) + [0] * seq_end_pad) + pad_mask = pad_mask[:, None].astype(bool) + + return seq, pad_mask + + def get_obs_sequence_from_demo(self, demo_id, index_in_demo, keys, num_frames_to_stack=0, seq_length=1, prefix="obs"): + """ + Extract a (sub)sequence of observation items from a demo given the @keys of the items. + + Args: + demo_id (str): id of the demo, e.g., demo_0 + index_in_demo (int): beginning index of the sequence wrt the demo + keys (tuple): list of keys to extract + num_frames_to_stack (int): numbers of frame to stack. Seq gets prepended with repeated items if out of range + seq_length (int): sequence length to extract. Seq gets post-pended with repeated items if out of range + prefix (str): one of "obs", "next_obs" + + Returns: + a dictionary of extracted items. + """ + obs, pad_mask = self.get_sequence_from_demo( + demo_id, + index_in_demo=index_in_demo, + keys=tuple('{}/{}'.format(prefix, k) for k in keys), + num_frames_to_stack=num_frames_to_stack, + seq_length=seq_length, + ) + obs = {'/'.join(k.split('/')[1:]): obs[k] for k in obs} # strip the prefix + if self.get_pad_mask: + obs["pad_mask"] = pad_mask + + return obs + + def get_dataset_sequence_from_demo(self, demo_id, index_in_demo, keys, num_frames_to_stack=0, seq_length=1): + """ + Extract a (sub)sequence of dataset items from a demo given the @keys of the items (e.g., states, actions). + + Args: + demo_id (str): id of the demo, e.g., demo_0 + index_in_demo (int): beginning index of the sequence wrt the demo + keys (tuple): list of keys to extract + num_frames_to_stack (int): numbers of frame to stack. Seq gets prepended with repeated items if out of range + seq_length (int): sequence length to extract. Seq gets post-pended with repeated items if out of range + + Returns: + a dictionary of extracted items. + """ + data, pad_mask = self.get_sequence_from_demo( + demo_id, + index_in_demo=index_in_demo, + keys=keys, + num_frames_to_stack=num_frames_to_stack, + seq_length=seq_length, + ) + if self.get_pad_mask: + data["pad_mask"] = pad_mask + return data + + def get_trajectory_at_index(self, index): + """ + Method provided as a utility to get an entire trajectory, given + the corresponding @index. + """ + demo_id = self.demos[index] + demo_length = self._demo_id_to_demo_length[demo_id] + + meta = self.get_dataset_sequence_from_demo( + demo_id, + index_in_demo=0, + keys=self.dataset_keys, + num_frames_to_stack=self.n_frame_stack - 1, # note: need to decrement self.n_frame_stack by one + seq_length=demo_length + ) + meta["obs"] = self.get_obs_sequence_from_demo( + demo_id, + index_in_demo=0, + keys=self.obs_keys, + seq_length=demo_length + ) + if self.load_next_obs: + meta["next_obs"] = self.get_obs_sequence_from_demo( + demo_id, + index_in_demo=0, + keys=self.obs_keys, + seq_length=demo_length, + prefix="next_obs" + ) + + meta["ep"] = demo_id + return meta + + def get_dataset_sampler(self): + """ + Return instance of torch.utils.data.Sampler or None. Allows + for dataset to define custom sampling logic, such as + re-weighting the probability of samples being drawn. + See the `train` function in scripts/train.py, and torch + `DataLoader` documentation, for more info. + """ + return None + + +class R2D2Dataset(SequenceDataset): + def get_action_traj(self, ep): + action_traj = dict() + for key in self.action_keys: + action_traj[key] = self.hdf5_file[key][()].astype('float32') + if len(action_traj[key].shape) == 1: + action_traj[key] = np.reshape(action_traj[key], (-1, 1)) + + return action_traj + + def load_demo_info(self, filter_by_attribute=None, demos=None, n_demos=None): + """ + Args: + filter_by_attribute (str): if provided, use the provided filter key + to select a subset of demonstration trajectories to load + + demos (list): list of demonstration keys to load from the hdf5 file. If + omitted, all demos in the file (or under the @filter_by_attribute + filter key) are used. + """ + + self.demos = ["demo"] + + self.n_demos = len(self.demos) + + # keep internal index maps to know which transitions belong to which demos + self._index_to_demo_id = dict() # maps every index to a demo id + self._demo_id_to_start_indices = dict() # gives start index per demo id + self._demo_id_to_demo_length = dict() + + # segment time stamps + self._demo_id_to_segments = dict() + + ep = self.demos[0] + + # determine index mapping + self.total_num_sequences = 0 + demo_length = self.hdf5_file["action/cartesian_velocity"].shape[0] + self._demo_id_to_start_indices[ep] = self.total_num_sequences + self._demo_id_to_demo_length[ep] = demo_length + + # seperate demo into segments for better alignment + gripper_actions = list(self.hdf5_file["action/gripper_position"]) + gripper_closed = [1 if x > 0 else 0 for x in gripper_actions] + + try: + # find when the gripper fist opens/closes + gripper_close = gripper_closed.index(1) + gripper_open = gripper_close + gripper_closed[gripper_close:].index(0) + except ValueError: + # special case for (invalid) trajectories + gripper_close, gripper_open = int(demo_length / 3), int(demo_length / 3 * 2) + print("No gripper action:", gripper_actions) + self._demo_id_to_segments[ep] = [0, gripper_close, gripper_open, demo_length - 1] + + num_sequences = demo_length + # determine actual number of sequences taking into account whether to pad for frame_stack and seq_length + if not self.pad_frame_stack: + num_sequences -= (self.n_frame_stack - 1) + if not self.pad_seq_length: + num_sequences -= (self.seq_length - 1) + + if self.pad_seq_length: + assert demo_length >= 1 # sequence needs to have at least one sample + num_sequences = max(num_sequences, 1) + else: + assert num_sequences >= 1 # assume demo_length >= (self.n_frame_stack - 1 + self.seq_length) + + for _ in range(num_sequences): + self._index_to_demo_id[self.total_num_sequences] = ep + self.total_num_sequences += 1 + + def load_dataset_in_memory(self, demo_list, hdf5_file, obs_keys, dataset_keys, load_next_obs): + """ + Loads the hdf5 dataset into memory, preserving the structure of the file. Note that this + differs from `self.getitem_cache`, which, if active, actually caches the outputs of the + `getitem` operation. + + Args: + demo_list (list): list of demo keys, e.g., 'demo_0' + hdf5_file (h5py.File): file handle to the hdf5 dataset. + obs_keys (list, tuple): observation keys to fetch, e.g., 'images' + dataset_keys (list, tuple): dataset keys to fetch, e.g., 'actions' + load_next_obs (bool): whether to load next_obs from the dataset + + Returns: + all_data (dict): dictionary of loaded data. + """ + all_data = dict() + print("SequenceDataset: loading dataset into memory...") + + for ep in LogUtils.custom_tqdm(demo_list): + all_data[ep] = {} + all_data[ep]["attrs"] = {} + all_data[ep]["attrs"]["num_samples"] = hdf5_file["action/cartesian_velocity"].shape[0] # hack to get traj len + # get obs + all_data[ep]["obs"] = {k: hdf5_file["observation/{}".format(k)][()].astype('float32') for k in obs_keys} + if load_next_obs: + raise NotImplementedError + # get other dataset keys + for k in dataset_keys: + if k in hdf5_file.keys(): + all_data[ep][k] = hdf5_file["{}".format(k)][()].astype('float32') + else: + raise NotImplementedError + + return all_data + + def get_dataset_for_ep(self, ep, key, try_to_use_cache=True): + """ + Helper utility to get a dataset for a specific demonstration. + Takes into account whether the dataset has been loaded into memory. + """ + + # check if this key should be in memory + key_should_be_in_memory = try_to_use_cache and (self.hdf5_cache_mode in ["all", "low_dim"]) + if key_should_be_in_memory: + # if key is an observation, it may not be in memory + if '/' in key: + key_splits = key.split('/') + key1 = key_splits[0] + key2 = "/".join(key_splits[1:]) + if key1 == "observation" and key2 not in self.obs_keys_in_memory: + key_should_be_in_memory = False + + if key_should_be_in_memory: + # read cache + if '/' in key: + key_splits = key.split('/') + key1 = key_splits[0] + key2 = "/".join(key_splits[1:]) + if key1 == "observation": + ret = self.hdf5_cache[ep]["obs"][key2] + else: + ret = self.hdf5_cache[ep][key] + else: + ret = self.hdf5_cache[ep][key] + else: + # read from file + hd5key = "{}".format(key) #"data/{}/{}".format(ep, key) + ret = self.hdf5_file[hd5key] + return ret + + + def get_sequence_from_demo(self, demo_id, index_in_demo, keys, num_frames_to_stack=0, seq_length=1): + """ + Extract a (sub)sequence of data items from a demo given the @keys of the items. + + Args: + demo_id (str): id of the demo, e.g., demo_0 + index_in_demo (int): beginning index of the sequence wrt the demo + keys (tuple): list of keys to extract + num_frames_to_stack (int): numbers of frame to stack. Seq gets prepended with repeated items if out of range + seq_length (int): sequence length to extract. Seq gets post-pended with repeated items if out of range + + Returns: + a dictionary of extracted items. + """ + assert num_frames_to_stack >= 0 + assert seq_length >= 1 + + demo_length = self._demo_id_to_demo_length[demo_id] + assert index_in_demo < demo_length + + # determine begin and end of sequence + seq_begin_index = max(0, index_in_demo - num_frames_to_stack) + seq_end_index = min(demo_length, index_in_demo + seq_length) + + # determine sequence padding + seq_begin_pad = max(0, num_frames_to_stack - index_in_demo) # pad for frame stacking + seq_end_pad = max(0, index_in_demo + seq_length - demo_length) # pad for sequence length + + # make sure we are not padding if specified. + if not self.pad_frame_stack: + assert seq_begin_pad == 0 + if not self.pad_seq_length: + assert seq_end_pad == 0 + + # fetch observation from the dataset file + seq = dict() + for k in keys: + data = self.get_dataset_for_ep(demo_id, k) + seq[k] = data[seq_begin_index: seq_end_index].astype("float32") + + seq = TensorUtils.pad_sequence(seq, padding=(seq_begin_pad, seq_end_pad), pad_same=True) + pad_mask = np.array([0] * seq_begin_pad + [1] * (seq_end_index - seq_begin_index) + [0] * seq_end_pad) + pad_mask = pad_mask[:, None].astype(np.bool_) + + return seq, pad_mask + + + def get_item(self, index): + """ + Main implementation of getitem when not using cache. + """ + + demo_id = self._index_to_demo_id[index] + demo_start_index = self._demo_id_to_start_indices[demo_id] + demo_length = self._demo_id_to_demo_length[demo_id] + + # start at offset index if not padding for frame stacking + demo_index_offset = 0 if self.pad_frame_stack else (self.n_frame_stack - 1) + index_in_demo = index - demo_start_index + demo_index_offset + + # end at offset index if not padding for seq length + demo_length_offset = 0 if self.pad_seq_length else (self.seq_length - 1) + end_index_in_demo = demo_length - demo_length_offset + + meta = self.get_dataset_sequence_from_demo( + demo_id, + index_in_demo=index_in_demo, + keys=self.dataset_keys, + num_frames_to_stack=self.n_frame_stack - 1, + seq_length=self.seq_length, + ) + + # determine goal index + goal_index = None + if self.goal_mode == "last": + goal_index = end_index_in_demo - 1 + + meta["obs"] = self.get_obs_sequence_from_demo( + demo_id, + index_in_demo=index_in_demo, + keys=self.obs_keys, + num_frames_to_stack=self.n_frame_stack - 1, + seq_length=self.seq_length, + prefix="observation" + ) + + if self.load_next_obs: + meta["next_obs"] = self.get_obs_sequence_from_demo( + demo_id, + index_in_demo=index_in_demo, + keys=self.obs_keys, + num_frames_to_stack=self.n_frame_stack - 1, + seq_length=self.seq_length, + prefix="next_obs" + ) + + if goal_index is not None: + goal = self.get_obs_sequence_from_demo( + demo_id, + index_in_demo=goal_index, + keys=self.obs_keys, + num_frames_to_stack=0, + seq_length=1, + prefix="next_obs", + ) + meta["goal_obs"] = {k: goal[k][0] for k in goal} # remove sequence dimension for goal + + # get action components + ac_dict = OrderedDict() + for k in self.action_keys: + ac = meta[k] + # expand action shape if needed + if len(ac.shape) == 1: + ac = ac.reshape(-1, 1) + ac_dict[k] = ac + + # normalize actions + action_normalization_stats = self.get_action_normalization_stats() + ac_dict = ObsUtils.normalize_dict(ac_dict, normalization_stats=action_normalization_stats) + + # concatenate all action components + meta["actions"] = AcUtils.action_dict_to_vector(ac_dict) + + # keys to reshape + for k in meta["obs"]: + if len(meta["obs"][k].shape) == 1: + meta["obs"][k] = np.expand_dims(meta["obs"][k], axis=1) + + # also return the sampled index + meta["index"] = index + + # language embedding + T = meta["actions"].shape[0] + meta["obs"]["lang_emb"] = np.tile(self._lang_emb, (T, 1)) + + return meta + + +class MetaDataset(torch.utils.data.Dataset): + def __init__( + self, + datasets, + ds_weights, + normalize_weights_by_ds_size=False, + ): + super(MetaDataset, self).__init__() + self.datasets = datasets + ds_lens = np.array([len(ds) for ds in self.datasets]) + if normalize_weights_by_ds_size: + self.ds_weights = np.array(ds_weights) / ds_lens + else: + self.ds_weights = ds_weights + self._ds_ind_bins = np.cumsum([0] + list(ds_lens)) + + # cache mode "all" not supported! The action normalization stats of each + # dataset will change after the datasets are already initialized + for ds in self.datasets: + assert ds.hdf5_cache_mode != "all" + + # TODO: comment + action_stats = self.get_action_stats() + self.action_normalization_stats = action_stats_to_normalization_stats( + action_stats, self.datasets[0].action_config) + self.set_action_normalization_stats(self.action_normalization_stats) + + def __len__(self): + return np.sum([len(ds) for ds in self.datasets]) + + def __getitem__(self, idx): + ds_ind = np.digitize(idx, self._ds_ind_bins) - 1 + ind_in_ds = idx - self._ds_ind_bins[ds_ind] + meta = self.datasets[ds_ind].__getitem__(ind_in_ds) + meta["index"] = idx + return meta + + def get_ds_label(self, idx): + ds_ind = np.digitize(idx, self._ds_ind_bins) - 1 + ds_label = self.ds_labels[ds_ind] + return ds_label + + def get_ds_id(self, idx): + ds_ind = np.digitize(idx, self._ds_ind_bins) - 1 + ds_label = self.ds_labels[ds_ind] + return self.ds_labels_to_ids[ds_label] + + def __repr__(self): + str_output = '\n'.join([ds.__repr__() for ds in self.datasets]) + return str_output + + def get_dataset_sampler(self): + weights = np.ones(len(self)) + for i, (start, end) in enumerate(zip(self._ds_ind_bins[:-1], self._ds_ind_bins[1:])): + weights[start:end] = self.ds_weights[i] + + sampler = torch.utils.data.WeightedRandomSampler( + weights=weights, + num_samples=len(self), + replacement=True, + ) + return sampler + + def get_action_stats(self): + meta_action_stats = self.datasets[0].get_action_stats() + for dataset in self.datasets[1:]: + ds_action_stats = dataset.get_action_stats() + meta_action_stats = _aggregate_traj_stats(meta_action_stats, ds_action_stats) + + return meta_action_stats + + def set_action_normalization_stats(self, action_normalization_stats): + self.action_normalization_stats = action_normalization_stats + for ds in self.datasets: + ds.set_action_normalization_stats(self.action_normalization_stats) + + def get_action_normalization_stats(self): + """ + Computes a dataset-wide min, max, mean and standard deviation for the actions + (per dimension) and returns it. + """ + + # Run through all trajectories. For each one, compute minimal observation statistics, and then aggregate + # with the previous statistics. + if self.action_normalization_stats is None: + action_stats = self.get_action_stats() + self.action_normalization_stats = action_stats_to_normalization_stats( + action_stats, self.datasets[0].action_config) + return self.action_normalization_stats + +def _compute_traj_stats(traj_obs_dict): + """ + Helper function to compute statistics over a single trajectory of observations. + """ + traj_stats = { k : {} for k in traj_obs_dict } + for k in traj_obs_dict: + traj_stats[k]["n"] = traj_obs_dict[k].shape[0] + traj_stats[k]["mean"] = traj_obs_dict[k].mean(axis=0, keepdims=True) # [1, ...] + traj_stats[k]["sqdiff"] = ((traj_obs_dict[k] - traj_stats[k]["mean"]) ** 2).sum(axis=0, keepdims=True) # [1, ...] + traj_stats[k]["min"] = traj_obs_dict[k].min(axis=0, keepdims=True) + traj_stats[k]["max"] = traj_obs_dict[k].max(axis=0, keepdims=True) + return traj_stats + +def _aggregate_traj_stats(traj_stats_a, traj_stats_b): + """ + Helper function to aggregate trajectory statistics. + See https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm + for more information. + """ + merged_stats = {} + for k in traj_stats_a: + n_a, avg_a, M2_a, min_a, max_a = traj_stats_a[k]["n"], traj_stats_a[k]["mean"], traj_stats_a[k]["sqdiff"], traj_stats_a[k]["min"], traj_stats_a[k]["max"] + n_b, avg_b, M2_b, min_b, max_b = traj_stats_b[k]["n"], traj_stats_b[k]["mean"], traj_stats_b[k]["sqdiff"], traj_stats_b[k]["min"], traj_stats_b[k]["max"] + n = n_a + n_b + mean = (n_a * avg_a + n_b * avg_b) / n + delta = (avg_b - avg_a) + M2 = M2_a + M2_b + (delta ** 2) * (n_a * n_b) / n + min_ = np.minimum(min_a, min_b) + max_ = np.maximum(max_a, max_b) + merged_stats[k] = dict(n=n, mean=mean, sqdiff=M2, min=min_, max=max_) + return merged_stats + +def action_stats_to_normalization_stats(action_stats, action_config): + action_normalization_stats = OrderedDict() + for action_key in action_stats.keys(): + # get how this action should be normalized from config, default to None + norm_method = action_config[action_key].get("normalization", None) + if norm_method is None: + # no normalization, unit scale, zero offset + action_normalization_stats[action_key] = { + "scale": np.ones_like(action_stats[action_key]["mean"], dtype=np.float32), + "offset": np.zeros_like(action_stats[action_key]["mean"], dtype=np.float32) + } + elif norm_method == "min_max": + # normalize min to -1 and max to 1 + range_eps = 1e-4 + input_min = action_stats[action_key]["min"].astype(np.float32) + input_max = action_stats[action_key]["max"].astype(np.float32) + # instead of -1 and 1 use numbers just below threshold to prevent numerical instability issues + output_min = -0.999999 + output_max = 0.999999 + + # ignore input dimentions that is too small to prevent division by zero + input_range = input_max - input_min + ignore_dim = input_range < range_eps + input_range[ignore_dim] = output_max - output_min + + # expected usage of scale and offset + # normalized_action = (raw_action - offset) / scale + # raw_action = scale * normalized_action + offset + + # eq1: input_max = scale * output_max + offset + # eq2: input_min = scale * output_min + offset + + # solution for scale and offset + # eq1 - eq2: + # input_max - input_min = scale * (output_max - output_min) + # (input_max - input_min) / (output_max - output_min) = scale <- eq3 + # offset = input_min - scale * output_min <- eq4 + scale = input_range / (output_max - output_min) + offset = input_min - scale * output_min + + offset[ignore_dim] = input_min[ignore_dim] - (output_max + output_min) / 2 + + action_normalization_stats[action_key] = { + "scale": scale, + "offset": offset + } + elif norm_method == "gaussian": + # normalize to zero mean unit variance + input_mean = action_stats[action_key]["mean"].astype(np.float32) + input_std = np.sqrt(action_stats[action_key]["sqdiff"] / action_stats[action_key]["n"]).astype(np.float32) + + # ignore input dimentions that is too small to prevent division by zero + std_eps = 1e-6 + ignore_dim = input_std < std_eps + input_std[ignore_dim] = 1.0 + + action_normalization_stats[action_key] = { + "scale": input_mean, + "offset": input_std + } + else: + raise NotImplementedError( + 'action_config.actions.normalization: "{}" is not supported'.format(norm_method)) + + return action_normalization_stats diff --git a/aloha-devel/robomimic/utils/hyperparam_utils.py b/aloha-devel/robomimic/utils/hyperparam_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0ff6397ebaabb433a6d1719fe9cebdf4fb2a0667 --- /dev/null +++ b/aloha-devel/robomimic/utils/hyperparam_utils.py @@ -0,0 +1,373 @@ +""" +A collection of utility functions and classes for generating config jsons for hyperparameter sweeps. +""" +import argparse +import os +import json +import re +import itertools + +from collections import OrderedDict +from copy import deepcopy + + +class ConfigGenerator(object): + """ + Useful class to keep track of hyperparameters to sweep, and to generate + the json configs for each experiment run. + """ + def __init__(self, base_config_file, wandb_proj_name="debug", script_file=None, generated_config_dir=None): + """ + Args: + base_config_file (str): path to a base json config to use as a starting point + for the parameter sweep. + + script_file (str): script filename to write as output + """ + assert isinstance(base_config_file, str) + self.base_config_file = base_config_file + assert generated_config_dir is None or isinstance(generated_config_dir, str) + if generated_config_dir is not None: + generated_config_dir = os.path.expanduser(generated_config_dir) + self.generated_config_dir = generated_config_dir + assert script_file is None or isinstance(script_file, str) + if script_file is None: + self.script_file = os.path.join('~', 'tmp/tmpp.sh') + else: + self.script_file = script_file + self.script_file = os.path.expanduser(self.script_file) + self.parameters = OrderedDict() + + assert isinstance(wandb_proj_name, str) + self.wandb_proj_name = wandb_proj_name + + def add_param(self, key, name, group, values, value_names=None, hidename=False, prepend=False): + """ + Add parameter to the hyperparameter sweep. + + Args: + key (str): location of parameter in the config, using hierarchical key format + (ex. train/data = config.train.data) + + name (str): name, as it will appear in the experiment name + + group (int): group id - parameters with the same ID have their values swept + together + + values (list): list of values to sweep over for this parameter + + value_names ([str]): if provided, strings to use in experiment name for + each value, instead of the parameter value. This is helpful for parameters + that may have long or large values (for example, dataset path). + """ + if value_names is not None: + assert len(values) == len(value_names) + self.parameters[key] = argparse.Namespace( + key=key, + name=name, + group=group, + values=values, + value_names=value_names, + hidename=hidename, + ) + if prepend: + self.parameters.move_to_end(key, last=False) + + def generate(self, override_base_name=False): + """ + Generates json configs for the hyperparameter sweep using attributes + @self.parameters, @self.base_config_file, and @self.script_file, + all of which should have first been set externally by calling + @add_param, @set_base_config_file, and @set_script_file. + """ + assert len(self.parameters) > 0, "must add parameters using add_param first!" + generated_json_paths = self._generate_jsons(override_base_name=override_base_name) + self._script_from_jsons(generated_json_paths) + + def _name_for_experiment(self, base_name, parameter_values, parameter_value_names): + """ + This function generates the name for an experiment, given one specific + parameter setting. + + Args: + base_name (str): base experiment name + parameter_values (OrderedDict): dictionary that maps parameter name to + the parameter value for this experiment run + parameter_value_names (dict): dictionary that maps parameter name to + the name to use for its value in the experiment name + + Returns: + name (str): generated experiment name + """ + name = base_name + for k in parameter_values: + # append parameter name and value to end of base name + if len(self.parameters[k].name) == 0 or self.parameters[k].hidename: + # empty string indicates that naming should be skipped + continue + if parameter_value_names[k] is not None: + # take name from passed dictionary + val_str = parameter_value_names[k] + else: + val_str = parameter_values[k] + if isinstance(parameter_values[k], list) or isinstance(parameter_values[k], tuple): + # convert list to string to avoid weird spaces and naming problems + val_str = "_".join([str(x) for x in parameter_values[k]]) + val_str = str(val_str) + if len(name) > 0: + name += "_" + name += '{}'.format(self.parameters[k].name) + if len(val_str) > 0: + name += '_{}'.format(val_str) + return name + + def _get_parameter_ranges(self): + """ + Extract parameter ranges from base json file. Also takes all possible + combinations of the parameter ranges to generate an expanded set of values. + + Returns: + parameter_ranges (dict): dictionary that maps the parameter to a list + of all values it should take for each generated config. The length + of the list will be the total number of configs that will be + generated from this scan. + + parameter_names (dict): dictionary that maps the parameter to a list + of all name strings that should contribute to each invididual + experiment's name. The length of the list will be the total + number of configs that will be generated from this scan. + """ + + # mapping from group id to list of indices to grab from each parameter's list + # of values in the parameter group + parameter_group_indices = OrderedDict() + for k in self.parameters: + group_id = self.parameters[k].group + assert isinstance(self.parameters[k].values, list) + num_param_values = len(self.parameters[k].values) + if group_id not in parameter_group_indices: + parameter_group_indices[group_id] = list(range(num_param_values)) + else: + assert len(parameter_group_indices[group_id]) == num_param_values, \ + "error: inconsistent number of parameter values in group with id {}".format(group_id) + + keys = list(parameter_group_indices.keys()) + inds = list(parameter_group_indices.values()) + new_parameter_group_indices = OrderedDict( + { k : [] for k in keys } + ) + # get all combinations of the different parameter group indices + # and then use these indices to determine the new parameter ranges + # per member of each parameter group. + # + # e.g. with two parameter groups, one with two values, and another with three values + # we have [0, 1] x [0, 1, 2] = [0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2] + # so the corresponding parameter group indices are [0, 0, 0, 1, 1, 1] and + # [0, 1, 2, 0, 1, 2], and all parameters in each parameter group are indexed + # together using these indices, to get each parameter range. + for comb in itertools.product(*inds): + for i in range(len(comb)): + new_parameter_group_indices[keys[i]].append(comb[i]) + parameter_group_indices = new_parameter_group_indices + + # use the indices to gather the parameter values to sweep per parameter + parameter_ranges = OrderedDict() + parameter_names = OrderedDict() + for k in self.parameters: + parameter_values = self.parameters[k].values + group_id = self.parameters[k].group + inds = parameter_group_indices[group_id] + parameter_ranges[k] = [parameter_values[ind] for ind in inds] + + # add in parameter names if supplied + parameter_names[k] = None + if self.parameters[k].value_names is not None: + par_names = self.parameters[k].value_names + assert isinstance(par_names, list) + assert len(par_names) == len(parameter_values) + parameter_names[k] = [par_names[ind] for ind in inds] + + # ensure that the number of parameter settings is the same per parameter + first_key = list(parameter_ranges.keys())[0] + num_settings = len(parameter_ranges[first_key]) + for k in parameter_ranges: + assert len(parameter_ranges[k]) == num_settings, "inconsistent number of values" + + return parameter_ranges, parameter_names + + def _generate_jsons(self, override_base_name=False): + """ + Generates json configs for the hyperparameter sweep, using @self.parameters and + @self.base_config_file. + + Returns: + json_paths (list): list of paths to created json files, one per experiment + """ + + # base directory for saving jsons + if self.generated_config_dir: + base_dir = self.generated_config_dir + if not os.path.exists(base_dir): + os.makedirs(base_dir) + else: + base_dir = os.path.abspath(os.path.dirname(self.base_config_file)) + + # read base json + base_config = load_json(self.base_config_file, verbose=False) + + # base exp name from this base config + if override_base_name: + base_exp_name = "" + else: + base_exp_name = base_config['experiment']['name'] + + # use base json to determine the parameter ranges + parameter_ranges, parameter_names = self._get_parameter_ranges() + + # iterate through each parameter setting to create each json + first_key = list(parameter_ranges.keys())[0] + num_settings = len(parameter_ranges[first_key]) + + # keep track of path to generated jsons + json_paths = [] + + for i in range(num_settings): + # the specific parameter setting for this experiment + setting = { k : parameter_ranges[k][i] for k in parameter_ranges } + maybe_parameter_names = OrderedDict() + for k in parameter_names: + maybe_parameter_names[k] = None + if parameter_names[k] is not None: + maybe_parameter_names[k] = parameter_names[k][i] + + # experiment name from setting + exp_name = self._name_for_experiment( + base_name=base_exp_name, + parameter_values=setting, + parameter_value_names=maybe_parameter_names, + ) + + # copy old json, but override name, and parameter values + json_dict = deepcopy(base_config) + json_dict['experiment']['name'] = exp_name + for k in parameter_ranges: + set_value_for_key(json_dict, k, v=parameter_ranges[k][i]) + + # populate list of identifying meta for logger; + # see meta_config method in base_config.py for more info + json_dict["experiment"]["logging"]["wandb_proj_name"] = self.wandb_proj_name + if "meta" not in json_dict: + json_dict["meta"] = dict() + json_dict["meta"].update( + hp_base_config_file=self.base_config_file, + hp_keys=list(), + hp_values=list(), + ) + # logging: keep track of hyp param names and values as meta info + for k in parameter_ranges.keys(): + key_name = self.parameters[k].name + if key_name is not None and len(key_name) > 0: + if maybe_parameter_names[k] is not None: + value_name = maybe_parameter_names[k] + else: + value_name = setting[k] + + json_dict["meta"]["hp_keys"].append(key_name) + json_dict["meta"]["hp_values"].append(value_name) + + # save file in same directory as old json + json_path = os.path.join(base_dir, "{}.json".format(exp_name)) + save_json(json_dict, json_path) + json_paths.append(json_path) + + print("Num exps:", len(json_paths)) + + return json_paths + + def _script_from_jsons(self, json_paths): + """ + Generates a bash script to run the experiments that correspond to + the input jsons. + """ + with open(self.script_file, 'w') as f: + f.write("#!/bin/bash\n\n") + for path in json_paths: + # write python command to file + import robomimic + cmd = "python {}/scripts/train.py --config {}\n".format(robomimic.__path__[0], path) + + print() + print(cmd) + f.write(cmd) + + +def load_json(json_file, verbose=True): + """ + Simple utility function to load a json file as a dict. + + Args: + json_file (str): path to json file to load + verbose (bool): if True, pretty print the loaded json dictionary + + Returns: + config (dict): json dictionary + """ + with open(json_file, 'r') as f: + config = json.load(f) + if verbose: + print('loading external config: =================') + print(json.dumps(config, indent=4)) + print('==========================================') + return config + + +def save_json(config, json_file): + """ + Simple utility function to save a dictionary to a json file on disk. + + Args: + config (dict): dictionary to save + json_file (str): path to json file to write + """ + with open(json_file, 'w') as f: + # preserve original key ordering + json.dump(config, f, sort_keys=False, indent=4) + + +def get_value_for_key(dic, k): + """ + Get value for nested dictionary with levels denoted by "/" or ".". + For example, if @k is "a/b", then this function returns + @dic["a"]["b"]. + + Args: + dic (dict): a nested dictionary + k (str): a single string meant to index several levels down into + the nested dictionary, where levels can be denoted by "/" or + by ".". + Returns: + val: the nested dictionary value for the provided key + """ + val = dic + subkeys = re.split('/|\.', k) + for s in subkeys[:-1]: + val = val[s] + return val[subkeys[-1]] + + +def set_value_for_key(dic, k, v): + """ + Set value for hierarchical dictionary with levels denoted by "/" or ".". + + Args: + dic (dict): a nested dictionary + k (str): a single string meant to index several levels down into + the nested dictionary, where levels can be denoted by "/" or + by ".". + v: the value to set at the provided key + """ + val = dic + subkeys = re.split('/|\.', k) #k.split('/') + for s in subkeys[:-1]: + val = val[s] + val[subkeys[-1]] = v diff --git a/aloha-devel/robomimic/utils/lang_utils.py b/aloha-devel/robomimic/utils/lang_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a28a1f2c372e9615c6d03a30680d1cd1731ac362 --- /dev/null +++ b/aloha-devel/robomimic/utils/lang_utils.py @@ -0,0 +1,27 @@ +import os +from transformers import AutoModel, pipeline, AutoTokenizer, CLIPTextModelWithProjection + +os.environ["TOKENIZERS_PARALLELISM"] = "true" # needed to suppress warning about potential deadlock +tokenizer = "openai/clip-vit-large-patch14" #"openai/clip-vit-base-patch32" +lang_emb_model = CLIPTextModelWithProjection.from_pretrained( + tokenizer, + cache_dir=os.path.expanduser("~/tmp/clip") +).eval() +tz = AutoTokenizer.from_pretrained(tokenizer, TOKENIZERS_PARALLELISM=True) + +def get_lang_emb(lang): + if lang is None: + return None + + tokens = tz( + text=lang, # the sentence to be encoded + add_special_tokens=True, # Add [CLS] and [SEP] + max_length=25, # maximum length of a sentence + padding="max_length", + return_attention_mask=True, # Generate the attention mask + return_tensors="pt", # ask the function to return PyTorch tensors + ) + lang_emb = lang_emb_model(**tokens)['text_embeds'].detach()[0] + + return lang_emb + diff --git a/aloha-devel/robomimic/utils/loss_utils.py b/aloha-devel/robomimic/utils/loss_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b3f5bf223ed7dfbd510b4b8a8edf2e98b0567613 --- /dev/null +++ b/aloha-devel/robomimic/utils/loss_utils.py @@ -0,0 +1,208 @@ +""" +This file contains a collection of useful loss functions for use with torch tensors. +""" + +import math +import numpy as np +import torch +import torch.nn.functional as F + + +def cosine_loss(preds, labels): + """ + Cosine loss between two tensors. + + Args: + preds (torch.Tensor): torch tensor + labels (torch.Tensor): torch tensor + + Returns: + loss (torch.Tensor): cosine loss + """ + sim = torch.nn.CosineSimilarity(dim=len(preds.shape) - 1)(preds, labels) + return -torch.mean(sim - 1.0) + + +def KLD_0_1_loss(mu, logvar): + """ + KL divergence loss. Computes D_KL( N(mu, sigma) || N(0, 1) ). Note that + this function averages across the batch dimension, but sums across dimension. + + Args: + mu (torch.Tensor): mean tensor of shape (B, D) + logvar (torch.Tensor): logvar tensor of shape (B, D) + + Returns: + loss (torch.Tensor): KL divergence loss between the input gaussian distribution + and N(0, 1) + """ + return -0.5 * (1. + logvar - mu.pow(2) - logvar.exp()).sum(dim=1).mean() + + +def KLD_gaussian_loss(mu_1, logvar_1, mu_2, logvar_2): + """ + KL divergence loss between two Gaussian distributions. This function + computes the average loss across the batch. + + Args: + mu_1 (torch.Tensor): first means tensor of shape (B, D) + logvar_1 (torch.Tensor): first logvars tensor of shape (B, D) + mu_2 (torch.Tensor): second means tensor of shape (B, D) + logvar_2 (torch.Tensor): second logvars tensor of shape (B, D) + + Returns: + loss (torch.Tensor): KL divergence loss between the two gaussian distributions + """ + return -0.5 * (1. + \ + logvar_1 - logvar_2 \ + - ((mu_2 - mu_1).pow(2) / logvar_2.exp()) \ + - (logvar_1.exp() / logvar_2.exp()) \ + ).sum(dim=1).mean() + + +def log_normal(x, m, v): + """ + Log probability of tensor x under diagonal multivariate normal with + mean m and variance v. The last dimension of the tensors is treated + as the dimension of the Gaussian distribution - all other dimensions + are treated as independent Gaussians. Adapted from CS 236 at Stanford. + + Args: + x (torch.Tensor): tensor with shape (B, ..., D) + m (torch.Tensor): means tensor with shape (B, ..., D) or (1, ..., D) + v (torch.Tensor): variances tensor with shape (B, ..., D) or (1, ..., D) + + Returns: + log_prob (torch.Tensor): log probabilities of shape (B, ...) + """ + element_wise = -0.5 * (torch.log(v) + (x - m).pow(2) / v + np.log(2 * np.pi)) + log_prob = element_wise.sum(-1) + return log_prob + + +def log_normal_mixture(x, m, v, w=None, log_w=None): + """ + Log probability of tensor x under a uniform mixture of Gaussians. + Adapted from CS 236 at Stanford. + + Args: + x (torch.Tensor): tensor with shape (B, D) + m (torch.Tensor): means tensor with shape (B, M, D) or (1, M, D), where + M is number of mixture components + v (torch.Tensor): variances tensor with shape (B, M, D) or (1, M, D) where + M is number of mixture components + w (torch.Tensor): weights tensor - if provided, should be + shape (B, M) or (1, M) + log_w (torch.Tensor): log-weights tensor - if provided, should be + shape (B, M) or (1, M) + + Returns: + log_prob (torch.Tensor): log probabilities of shape (B,) + """ + + # (B , D) -> (B , 1, D) + x = x.unsqueeze(1) + # (B, 1, D) -> (B, M, D) -> (B, M) + log_prob = log_normal(x, m, v) + if w is not None or log_w is not None: + # this weights the log probabilities by the mixture weights so we have log(w_i * N(x | m_i, v_i)) + if w is not None: + assert log_w is None + log_w = torch.log(w) + log_prob += log_w + # then compute log sum_i exp [log(w_i * N(x | m_i, v_i))] + # (B, M) -> (B,) + log_prob = log_sum_exp(log_prob , dim=1) + else: + # (B, M) -> (B,) + log_prob = log_mean_exp(log_prob , dim=1) # mean accounts for uniform weights + return log_prob + + +def log_mean_exp(x, dim): + """ + Compute the log(mean(exp(x), dim)) in a numerically stable manner. + Adapted from CS 236 at Stanford. + + Args: + x (torch.Tensor): a tensor + dim (int): dimension along which mean is computed + + Returns: + y (torch.Tensor): log(mean(exp(x), dim)) + """ + return log_sum_exp(x, dim) - np.log(x.size(dim)) + + +def log_sum_exp(x, dim=0): + """ + Compute the log(sum(exp(x), dim)) in a numerically stable manner. + Adapted from CS 236 at Stanford. + + Args: + x (torch.Tensor): a tensor + dim (int): dimension along which sum is computed + + Returns: + y (torch.Tensor): log(sum(exp(x), dim)) + """ + max_x = torch.max(x, dim)[0] + new_x = x - max_x.unsqueeze(dim).expand_as(x) + return max_x + (new_x.exp().sum(dim)).log() + + +def project_values_onto_atoms(values, probabilities, atoms): + """ + Project the categorical distribution given by @probabilities on the + grid of values given by @values onto a grid of values given by @atoms. + This is useful when computing a bellman backup where the backed up + values from the original grid will not be in the original support, + requiring L2 projection. + + Each value in @values has a corresponding probability in @probabilities - + this probability mass is shifted to the closest neighboring grid points in + @atoms in proportion. For example, if the value in question is 0.2, and the + neighboring atoms are 0 and 1, then 0.8 of the probability weight goes to + atom 0 and 0.2 of the probability weight will go to 1. + + Adapted from https://github.com/deepmind/acme/blob/master/acme/tf/losses/distributional.py#L42 + + Args: + values: value grid to project, of shape (batch_size, n_atoms) + probabilities: probabilities for categorical distribution on @values, shape (batch_size, n_atoms) + atoms: value grid to project onto, of shape (n_atoms,) or (1, n_atoms) + + Returns: + new probability vectors that correspond to the L2 projection of the categorical distribution + onto @atoms + """ + + # make sure @atoms is shape (n_atoms,) + if len(atoms.shape) > 1: + atoms = atoms.squeeze(0) + + # helper tensors from @atoms + vmin, vmax = atoms[0], atoms[1] + d_pos = torch.cat([atoms, vmin[None]], dim=0)[1:] + d_neg = torch.cat([vmax[None], atoms], dim=0)[:-1] + + # ensure that @values grid is within the support of @atoms + clipped_values = values.clamp(min=vmin, max=vmax)[:, None, :] # (batch_size, 1, n_atoms) + clipped_atoms = atoms[None, :, None] # (1, n_atoms, 1) + + # distance between atom values in support + d_pos = (d_pos - atoms)[None, :, None] # atoms[i + 1] - atoms[i], shape (1, n_atoms, 1) + d_neg = (atoms - d_neg)[None, :, None] # atoms[i] - atoms[i - 1], shape (1, n_atoms, 1) + + # distances between all pairs of grid values + deltas = clipped_values - clipped_atoms # (batch_size, n_atoms, n_atoms) + + # computes eqn (7) in distributional RL paper by doing the following - for each + # output atom in @atoms, consider values that are close enough, and weight their + # probability mass contribution by the normalized distance in [0, 1] given + # by (1. - (z_j - z_i) / (delta_z)). + d_sign = (deltas >= 0.).float() + delta_hat = (d_sign * deltas / d_pos) - ((1. - d_sign) * deltas / d_neg) + delta_hat = (1. - delta_hat).clamp(min=0., max=1.) + probabilities = probabilities[:, None, :] + return (delta_hat * probabilities).sum(dim=2) diff --git a/aloha-devel/robomimic/utils/obs_utils.py b/aloha-devel/robomimic/utils/obs_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e901e436a331238b68aacbee95f041bee7fa4637 --- /dev/null +++ b/aloha-devel/robomimic/utils/obs_utils.py @@ -0,0 +1,1002 @@ +""" +A collection of utilities for working with observation dictionaries and +different kinds of modalities such as images. +""" +import numpy as np +from copy import deepcopy +from collections import OrderedDict + +import torch +import torch.nn.functional as F + +import robomimic.utils.tensor_utils as TU + +# MACRO FOR VALID IMAGE CHANNEL SIZES +VALID_IMAGE_CHANNEL_DIMS = {1, 3} # depth, rgb + +# DO NOT MODIFY THIS! +# This keeps track of observation types (modalities) - and is populated on call to @initialize_obs_utils_with_obs_specs. +# This will be a dictionary that maps observation modality (e.g. low_dim, rgb) to a list of observation +# keys under that observation modality. +OBS_MODALITIES_TO_KEYS = None + +# DO NOT MODIFY THIS! +# This keeps track of observation types (modalities) - and is populated on call to @initialize_obs_utils_with_obs_specs. +# This will be a dictionary that maps observation keys to their corresponding observation modality +# (e.g. low_dim, rgb) +OBS_KEYS_TO_MODALITIES = None + +# DO NOT MODIFY THIS +# This holds the default encoder kwargs that will be used if none are passed at runtime for any given network +DEFAULT_ENCODER_KWARGS = None + +# DO NOT MODIFY THIS +# This holds the registered observation modality classes +OBS_MODALITY_CLASSES = {} + +# DO NOT MODIFY THIS +# This global dict stores mapping from observation encoder / randomizer network name to class. +# We keep track of these registries to enable automated class inference at runtime, allowing +# users to simply extend our base encoder / randomizer class and refer to that class in string form +# in their config, without having to manually register their class internally. +# This also future-proofs us for any additional encoder / randomizer classes we would +# like to add ourselves. +OBS_ENCODER_CORES = {"None": None} # Include default None +OBS_RANDOMIZERS = {"None": None} # Include default None + + +def register_obs_key(target_class): + assert target_class not in OBS_MODALITY_CLASSES, f"Already registered modality {target_class}!" + OBS_MODALITY_CLASSES[target_class.name] = target_class + + +def register_encoder_core(target_class): + assert target_class not in OBS_ENCODER_CORES, f"Already registered obs encoder core {target_class}!" + OBS_ENCODER_CORES[target_class.__name__] = target_class + + +def register_randomizer(target_class): + assert target_class not in OBS_RANDOMIZERS, f"Already registered obs randomizer {target_class}!" + OBS_RANDOMIZERS[target_class.__name__] = target_class + + +class ObservationKeyToModalityDict(dict): + """ + Custom dictionary class with the sole additional purpose of automatically registering new "keys" at runtime + without breaking. This is mainly for backwards compatibility, where certain keys such as "latent", "actions", etc. + are used automatically by certain models (e.g.: VAEs) but were never specified by the user externally in their + config. Thus, this dictionary will automatically handle those keys by implicitly associating them with the low_dim + modality. + """ + def __getitem__(self, item): + # If a key doesn't already exist, warn the user and add default mapping + if item not in self.keys(): + print(f"ObservationKeyToModalityDict: {item} not found," + f" adding {item} to mapping with assumed low_dim modality!") + self.__setitem__(item, "low_dim") + return super(ObservationKeyToModalityDict, self).__getitem__(item) + + +def obs_encoder_kwargs_from_config(obs_encoder_config): + """ + Generate a set of args used to create visual backbones for networks + from the observation encoder config. + + Args: + obs_encoder_config (Config): Config object containing relevant encoder information. Should be equivalent to + config.observation.encoder + + Returns: + dict: Processed encoder kwargs + """ + # Loop over each obs modality + # Unlock encoder config + obs_encoder_config.unlock() + for obs_modality, encoder_kwargs in obs_encoder_config.items(): + ### TODO: fix sanity checks. Disabling this code snippet to allow chaining multiple randomizers. + # # First run some sanity checks and store the classes + # for cls_name, cores in zip(("core", "obs_randomizer"), (OBS_ENCODER_CORES, OBS_RANDOMIZERS)): + # # Make sure the requested encoder for each obs_modality exists + # cfg_cls = encoder_kwargs[f"{cls_name}_class"] + # if cfg_cls is not None: + # assert cfg_cls in cores, f"No {cls_name} class with name {cfg_cls} found, must register this class before" \ + # f"creating model!" + # # encoder_kwargs[f"{cls_name}_class"] = cores[cfg_cls] + + # Process core and randomizer kwargs + encoder_kwargs.core_kwargs = dict() if encoder_kwargs.core_kwargs is None else \ + deepcopy(encoder_kwargs.core_kwargs) + encoder_kwargs.obs_randomizer_kwargs = dict() if encoder_kwargs.obs_randomizer_kwargs is None else \ + deepcopy(encoder_kwargs.obs_randomizer_kwargs) + + # Re-lock keys + obs_encoder_config.lock() + + return dict(obs_encoder_config) + + +def initialize_obs_modality_mapping_from_dict(modality_mapping): + """ + This function is an alternative to @initialize_obs_utils_with_obs_specs, that allows manually setting of modalities. + NOTE: Only one of these should be called at runtime -- not both! (Note that all training scripts that use a config) + automatically handle obs modality mapping, so using this function is usually unnecessary) + + Args: + modality_mapping (dict): Maps modality string names (e.g.: rgb, low_dim, etc.) to a list of observation + keys that should belong to that modality + """ + global OBS_KEYS_TO_MODALITIES, OBS_MODALITIES_TO_KEYS + + OBS_KEYS_TO_MODALITIES = ObservationKeyToModalityDict() + OBS_MODALITIES_TO_KEYS = dict() + + for mod, keys in modality_mapping.items(): + OBS_MODALITIES_TO_KEYS[mod] = deepcopy(keys) + OBS_KEYS_TO_MODALITIES.update({k: mod for k in keys}) + + +def initialize_obs_utils_with_obs_specs(obs_modality_specs): + """ + This function should be called before using any observation key-specific + functions in this file, in order to make sure that all utility + functions are aware of the observation modalities (e.g. which ones + are low-dimensional, which ones are rgb, etc.). + + It constructs two dictionaries: (1) that map observation modality (e.g. low_dim, rgb) to + a list of observation keys under that modality, and (2) that maps the inverse, specific + observation keys to their corresponding observation modality. + + Input should be a nested dictionary (or list of such dicts) with the following structure: + + obs_variant (str): + obs_modality (str): observation keys (list) + ... + ... + + Example: + { + "obs": { + "low_dim": ["robot0_eef_pos", "robot0_eef_quat"], + "rgb": ["agentview_image", "robot0_eye_in_hand"], + } + "goal": { + "low_dim": ["robot0_eef_pos"], + "rgb": ["agentview_image"] + } + } + + In the example, raw observations consist of low-dim and rgb modalities, with + the robot end effector pose under low-dim, and the agentview and wrist camera + images under rgb, while goal observations also consist of low-dim and rgb modalities, + with a subset of the raw observation keys per modality. + + Args: + obs_modality_specs (dict or list): A nested dictionary (see docstring above for an example) + or a list of nested dictionaries. Accepting a list as input makes it convenient for + situations where multiple modules may each have their own modality spec. + """ + global OBS_KEYS_TO_MODALITIES, OBS_MODALITIES_TO_KEYS + + OBS_KEYS_TO_MODALITIES = ObservationKeyToModalityDict() + + # accept one or more spec dictionaries - if it's just one, account for this + if isinstance(obs_modality_specs, dict): + obs_modality_spec_list = [obs_modality_specs] + else: + obs_modality_spec_list = obs_modality_specs + + # iterates over observation specs + obs_modality_mapping = {} + for obs_modality_spec in obs_modality_spec_list: + # iterates over observation variants (e.g. observations, goals, subgoals) + for obs_modalities in obs_modality_spec.values(): + for obs_modality, obs_keys in obs_modalities.items(): + # add all keys for each obs modality to the corresponding list in obs_modality_mapping + if obs_modality not in obs_modality_mapping: + obs_modality_mapping[obs_modality] = [] + obs_modality_mapping[obs_modality] += obs_keys + # loop over each modality, and add to global dict if it doesn't exist yet + for obs_key in obs_keys: + if obs_key not in OBS_KEYS_TO_MODALITIES: + OBS_KEYS_TO_MODALITIES[obs_key] = obs_modality + # otherwise, run sanity check to make sure we don't have conflicting, duplicate entries + else: + assert OBS_KEYS_TO_MODALITIES[obs_key] == obs_modality, \ + f"Cannot register obs key {obs_key} with modality {obs_modality}; " \ + f"already exists with corresponding modality {OBS_KEYS_TO_MODALITIES[obs_key]}" + + # remove duplicate entries and store in global mapping + OBS_MODALITIES_TO_KEYS = { obs_modality : list(set(obs_modality_mapping[obs_modality])) for obs_modality in obs_modality_mapping } + + print("\n============= Initialized Observation Utils with Obs Spec =============\n") + for obs_modality, obs_keys in OBS_MODALITIES_TO_KEYS.items(): + print("using obs modality: {} with keys: {}".format(obs_modality, obs_keys)) + + +def initialize_default_obs_encoder(obs_encoder_config): + """ + Initializes the default observation encoder kwarg information to be used by all networks if no values are manually + specified at runtime. + + Args: + obs_encoder_config (Config): Observation encoder config to use. + Should be equivalent to config.observation.encoder + """ + global DEFAULT_ENCODER_KWARGS + DEFAULT_ENCODER_KWARGS = obs_encoder_kwargs_from_config(obs_encoder_config) + + +def initialize_obs_utils_with_config(config): + """ + Utility function to parse config and call @initialize_obs_utils_with_obs_specs and + @initialize_default_obs_encoder_kwargs with the correct arguments. + + Args: + config (BaseConfig instance): config object + """ + if config.algo_name == "hbc": + obs_modality_specs = [ + config.observation.planner.modalities, + config.observation.actor.modalities, + ] + obs_encoder_config = config.observation.actor.encoder + elif config.algo_name == "iris": + obs_modality_specs = [ + config.observation.value_planner.planner.modalities, + config.observation.value_planner.value.modalities, + config.observation.actor.modalities, + ] + obs_encoder_config = config.observation.actor.encoder + else: + obs_modality_specs = [config.observation.modalities] + obs_encoder_config = config.observation.encoder + initialize_obs_utils_with_obs_specs(obs_modality_specs=obs_modality_specs) + initialize_default_obs_encoder(obs_encoder_config=obs_encoder_config) + + +def key_is_obs_modality(key, obs_modality): + """ + Check if observation key corresponds to modality @obs_modality. + + Args: + key (str): obs key name to check + obs_modality (str): observation modality - e.g.: "low_dim", "rgb" + """ + assert OBS_KEYS_TO_MODALITIES is not None, "error: must call ObsUtils.initialize_obs_utils_with_obs_config first" + return OBS_KEYS_TO_MODALITIES[key] == obs_modality + + +def center_crop(im, t_h, t_w): + """ + Takes a center crop of an image. + + Args: + im (np.array or torch.Tensor): image of shape (..., height, width, channel) + t_h (int): height of crop + t_w (int): width of crop + + Returns: + im (np.array or torch.Tensor): center cropped image + """ + assert(im.shape[-3] >= t_h and im.shape[-2] >= t_w) + assert(im.shape[-1] in [1, 3]) + crop_h = int((im.shape[-3] - t_h) / 2) + crop_w = int((im.shape[-2] - t_w) / 2) + return im[..., crop_h:crop_h + t_h, crop_w:crop_w + t_w, :] + + +def batch_image_hwc_to_chw(im): + """ + Channel swap for images - useful for preparing images for + torch training. + + Args: + im (np.array or torch.Tensor): image of shape (batch, height, width, channel) + or (height, width, channel) + + Returns: + im (np.array or torch.Tensor): image of shape (batch, channel, height, width) + or (channel, height, width) + """ + start_dims = np.arange(len(im.shape) - 3).tolist() + s = start_dims[-1] if len(start_dims) > 0 else -1 + if isinstance(im, np.ndarray): + return im.transpose(start_dims + [s + 3, s + 1, s + 2]) + else: + return im.permute(start_dims + [s + 3, s + 1, s + 2]) + + +def batch_image_chw_to_hwc(im): + """ + Inverse of channel swap in @batch_image_hwc_to_chw. + + Args: + im (np.array or torch.Tensor): image of shape (batch, channel, height, width) + or (channel, height, width) + + Returns: + im (np.array or torch.Tensor): image of shape (batch, height, width, channel) + or (height, width, channel) + """ + start_dims = np.arange(len(im.shape) - 3).tolist() + s = start_dims[-1] if len(start_dims) > 0 else -1 + if isinstance(im, np.ndarray): + return im.transpose(start_dims + [s + 2, s + 3, s + 1]) + else: + return im.permute(start_dims + [s + 2, s + 3, s + 1]) + + +def process_obs(obs, obs_modality=None, obs_key=None): + """ + Process observation @obs corresponding to @obs_modality modality (or implicitly inferred from @obs_key) + to prepare for network input. + + Note that either obs_modality OR obs_key must be specified! + + If both are specified, obs_key will override obs_modality + + Args: + obs (np.array or torch.Tensor): Observation to process. Leading batch dimension is optional + obs_modality (str): Observation modality (e.g.: depth, image, low_dim, etc.) + obs_key (str): Name of observation from which to infer @obs_modality + + Returns: + processed_obs (np.array or torch.Tensor): processed observation + """ + assert obs_modality is not None or obs_key is not None, "Either obs_modality or obs_key must be specified!" + if obs_key is not None: + obs_modality = OBS_KEYS_TO_MODALITIES[obs_key] + return OBS_MODALITY_CLASSES[obs_modality].process_obs(obs) + + +def process_obs_dict(obs_dict): + """ + Process observations in observation dictionary to prepare for network input. + + Args: + obs_dict (dict): dictionary mapping observation keys to np.array or + torch.Tensor. Leading batch dimensions are optional. + + Returns: + new_dict (dict): dictionary where observation keys have been processed by their corresponding processors + """ + return { k : process_obs(obs=obs, obs_key=k) for k, obs in obs_dict.items() } # shallow copy + + +def process_frame(frame, channel_dim, scale): + """ + Given frame fetched from dataset, process for network input. Converts array + to float (from uint8), normalizes pixels from range [0, @scale] to [0, 1], and channel swaps + from (H, W, C) to (C, H, W). + + Args: + frame (np.array or torch.Tensor): frame array + channel_dim (int): Number of channels to sanity check for + scale (float): Value to normalize inputs by + + Returns: + processed_frame (np.array or torch.Tensor): processed frame + """ + # Channel size should either be 3 (RGB) or 1 (depth) + assert (frame.shape[-1] == channel_dim) + frame = TU.to_float(frame) + frame /= scale + frame = frame.clip(0.0, 1.0) + frame = batch_image_hwc_to_chw(frame) + + return frame + + +def unprocess_obs(obs, obs_modality=None, obs_key=None): + """ + Prepare observation @obs corresponding to @obs_modality modality (or implicitly inferred from @obs_key) + to prepare for deployment. + + Note that either obs_modality OR obs_key must be specified! + + If both are specified, obs_key will override obs_modality + + Args: + obs (np.array or torch.Tensor): Observation to unprocess. Leading batch dimension is optional + obs_modality (str): Observation modality (e.g.: depth, image, low_dim, etc.) + obs_key (str): Name of observation from which to infer @obs_modality + + Returns: + unprocessed_obs (np.array or torch.Tensor): unprocessed observation + """ + assert obs_modality is not None or obs_key is not None, "Either obs_modality or obs_key must be specified!" + if obs_key is not None: + obs_modality = OBS_KEYS_TO_MODALITIES[obs_key] + return OBS_MODALITY_CLASSES[obs_modality].unprocess_obs(obs) + + +def unprocess_obs_dict(obs_dict): + """ + Prepare processed observation dictionary for saving to dataset. Inverse of + @process_obs. + + Args: + obs_dict (dict): dictionary mapping observation keys to np.array or + torch.Tensor. Leading batch dimensions are optional. + + Returns: + new_dict (dict): dictionary where observation keys have been unprocessed by + their respective unprocessor methods + """ + return { k : unprocess_obs(obs=obs, obs_key=k) for k, obs in obs_dict.items() } # shallow copy + + +def unprocess_frame(frame, channel_dim, scale): + """ + Given frame prepared for network input, prepare for saving to dataset. + Inverse of @process_frame. + + Args: + frame (np.array or torch.Tensor): frame array + channel_dim (int): What channel dimension should be (used for sanity check) + scale (float): Scaling factor to apply during denormalization + + Returns: + unprocessed_frame (np.array or torch.Tensor): frame passed through + inverse operation of @process_frame + """ + assert frame.shape[-3] == channel_dim # check for channel dimension + frame = batch_image_chw_to_hwc(frame) + frame *= scale + return frame + + +def get_processed_shape(obs_modality, input_shape): + """ + Given observation modality @obs_modality and expected inputs of shape @input_shape (excluding batch dimension), return the + expected processed observation shape resulting from process_{obs_modality}. + + Args: + obs_modality (str): Observation modality to use (e.g.: low_dim, rgb, depth, etc...) + input_shape (list of int): Expected input dimensions, excluding the batch dimension + + Returns: + list of int: expected processed input shape + """ + return list(process_obs(obs=np.zeros(input_shape), obs_modality=obs_modality).shape) + + +def normalize_dict(dict, normalization_stats): + """ + Normalize dict using the provided "offset" and "scale" entries + for each observation key. The dictionary will be + modified in-place. + + Args: + dict (dict): dictionary mapping key to np.array or + torch.Tensor. Leading batch dimensions are optional. + + normalization_stats (dict): this should map keys to dicts + with a "offset" and "scale" of shape (1, ...) where ... is the default + shape for the dict value. + + Returns: + dict (dict): obs dict with normalized arrays + """ + + # ensure we have statistics for each modality key in the dict + assert set(dict.keys()).issubset(normalization_stats) + + for m in dict: + offset = normalization_stats[m]["offset"] + scale = normalization_stats[m]["scale"] + + # check shape consistency + shape_len_diff = len(offset.shape) - len(dict[m].shape) + assert shape_len_diff in [0, 1], "shape length mismatch in @normalize_dict" + # if dict has no leading batch dim, check shapes match exactly, else allow first dim to broadcast + assert offset.shape[1:] == dict[m].shape[(1 - shape_len_diff):], "shape mismatch in @normalize_dict" + + # handle case where obs dict is not batched by removing stats batch dimension + if shape_len_diff == 1: + offset = offset[0] + scale = scale[0] + + dict[m] = (dict[m] - offset) / scale + + return dict + + +def unnormalize_dict(dict, normalization_stats): + """ + Unnormalize dict using the provided "offset" and "scale" entries + for each observation key. The dictionary will be + modified in-place. + + Args: + dict (dict): dictionary mapping key to np.array or + torch.Tensor. Leading batch dimensions are optional. + + normalization_stats (dict): this should map keys to dicts + with a "offset" and "scale" of shape (1, ...) where ... is the default + shape for the dict value. + + Returns: + dict (dict): obs dict with normalized arrays + """ + + # ensure we have statistics for each modality key in the dict + assert set(dict.keys()).issubset(normalization_stats) + + for m in dict: + offset = normalization_stats[m]["offset"] + scale = normalization_stats[m]["scale"] + + # check shape consistency + shape_len_diff = len(offset.shape) - len(dict[m].shape) + assert shape_len_diff in [0, 1], "shape length mismatch in @unnormalize_dict" + # if dict has no leading batch dim, check shapes match exactly, else allow first dim to broadcast + assert offset.shape[1:] == dict[m].shape[(1 - shape_len_diff):], "shape mismatch in @unnormalize_dict" + + # handle case where obs dict is not batched by removing stats batch dimension + if shape_len_diff == 1: + offset = offset[0] + scale = scale[0] + + dict[m] = (dict[m] * scale) + offset + + return dict + + +def has_modality(modality, obs_keys): + """ + Returns True if @modality is present in the list of observation keys @obs_keys. + + Args: + modality (str): modality to check for, e.g.: rgb, depth, etc. + obs_keys (list): list of observation keys + """ + for k in obs_keys: + if key_is_obs_modality(k, obs_modality=modality): + return True + return False + + +def repeat_and_stack_observation(obs_dict, n): + """ + Given an observation dictionary and a desired repeat value @n, + this function will return a new observation dictionary where + each modality is repeated @n times and the copies are + stacked in the first dimension. + + For example, if a batch of 3 observations comes in, and n is 2, + the output will look like [ob1; ob1; ob2; ob2; ob3; ob3] in + each modality. + + Args: + obs_dict (dict): dictionary mapping observation key to np.array or + torch.Tensor. Leading batch dimensions are optional. + + n (int): number to repeat by + + Returns: + repeat_obs_dict (dict): repeated obs dict + """ + return TU.repeat_by_expand_at(obs_dict, repeats=n, dim=0) + + +def crop_image_from_indices(images, crop_indices, crop_height, crop_width): + """ + Crops images at the locations specified by @crop_indices. Crops will be + taken across all channels. + + Args: + images (torch.Tensor): batch of images of shape [..., C, H, W] + + crop_indices (torch.Tensor): batch of indices of shape [..., N, 2] where + N is the number of crops to take per image and each entry corresponds + to the pixel height and width of where to take the crop. Note that + the indices can also be of shape [..., 2] if only 1 crop should + be taken per image. Leading dimensions must be consistent with + @images argument. Each index specifies the top left of the crop. + Values must be in range [0, H - CH - 1] x [0, W - CW - 1] where + H and W are the height and width of @images and CH and CW are + @crop_height and @crop_width. + + crop_height (int): height of crop to take + + crop_width (int): width of crop to take + + Returns: + crops (torch.Tesnor): cropped images of shape [..., C, @crop_height, @crop_width] + """ + + # make sure length of input shapes is consistent + assert crop_indices.shape[-1] == 2 + ndim_im_shape = len(images.shape) + ndim_indices_shape = len(crop_indices.shape) + assert (ndim_im_shape == ndim_indices_shape + 1) or (ndim_im_shape == ndim_indices_shape + 2) + + # maybe pad so that @crop_indices is shape [..., N, 2] + is_padded = False + if ndim_im_shape == ndim_indices_shape + 2: + crop_indices = crop_indices.unsqueeze(-2) + is_padded = True + + # make sure leading dimensions between images and indices are consistent + assert images.shape[:-3] == crop_indices.shape[:-2] + + device = images.device + image_c, image_h, image_w = images.shape[-3:] + num_crops = crop_indices.shape[-2] + + # make sure @crop_indices are in valid range + assert (crop_indices[..., 0] >= 0).all().item() + assert (crop_indices[..., 0] < (image_h - crop_height)).all().item() + assert (crop_indices[..., 1] >= 0).all().item() + assert (crop_indices[..., 1] < (image_w - crop_width)).all().item() + + # convert each crop index (ch, cw) into a list of pixel indices that correspond to the entire window. + + # 2D index array with columns [0, 1, ..., CH - 1] and shape [CH, CW] + crop_ind_grid_h = torch.arange(crop_height).to(device) + crop_ind_grid_h = TU.unsqueeze_expand_at(crop_ind_grid_h, size=crop_width, dim=-1) + # 2D index array with rows [0, 1, ..., CW - 1] and shape [CH, CW] + crop_ind_grid_w = torch.arange(crop_width).to(device) + crop_ind_grid_w = TU.unsqueeze_expand_at(crop_ind_grid_w, size=crop_height, dim=0) + # combine into shape [CH, CW, 2] + crop_in_grid = torch.cat((crop_ind_grid_h.unsqueeze(-1), crop_ind_grid_w.unsqueeze(-1)), dim=-1) + + # Add above grid with the offset index of each sampled crop to get 2d indices for each crop. + # After broadcasting, this will be shape [..., N, CH, CW, 2] and each crop has a [CH, CW, 2] + # shape array that tells us which pixels from the corresponding source image to grab. + grid_reshape = [1] * len(crop_indices.shape[:-1]) + [crop_height, crop_width, 2] + all_crop_inds = crop_indices.unsqueeze(-2).unsqueeze(-2) + crop_in_grid.reshape(grid_reshape) + + # For using @torch.gather, convert to flat indices from 2D indices, and also + # repeat across the channel dimension. To get flat index of each pixel to grab for + # each sampled crop, we just use the mapping: ind = h_ind * @image_w + w_ind + all_crop_inds = all_crop_inds[..., 0] * image_w + all_crop_inds[..., 1] # shape [..., N, CH, CW] + all_crop_inds = TU.unsqueeze_expand_at(all_crop_inds, size=image_c, dim=-3) # shape [..., N, C, CH, CW] + all_crop_inds = TU.flatten(all_crop_inds, begin_axis=-2) # shape [..., N, C, CH * CW] + + # Repeat and flatten the source images -> [..., N, C, H * W] and then use gather to index with crop pixel inds + images_to_crop = TU.unsqueeze_expand_at(images, size=num_crops, dim=-4) + images_to_crop = TU.flatten(images_to_crop, begin_axis=-2) + crops = torch.gather(images_to_crop, dim=-1, index=all_crop_inds) + # [..., N, C, CH * CW] -> [..., N, C, CH, CW] + reshape_axis = len(crops.shape) - 1 + crops = TU.reshape_dimensions(crops, begin_axis=reshape_axis, end_axis=reshape_axis, + target_dims=(crop_height, crop_width)) + + if is_padded: + # undo padding -> [..., C, CH, CW] + crops = crops.squeeze(-4) + return crops + + +def sample_random_image_crops(images, crop_height, crop_width, num_crops, pos_enc=False): + """ + For each image, randomly sample @num_crops crops of size (@crop_height, @crop_width), from + @images. + + Args: + images (torch.Tensor): batch of images of shape [..., C, H, W] + + crop_height (int): height of crop to take + + crop_width (int): width of crop to take + + num_crops (n): number of crops to sample + + pos_enc (bool): if True, also add 2 channels to the outputs that gives a spatial + encoding of the original source pixel locations. This means that the + output crops will contain information about where in the source image + it was sampled from. + + Returns: + crops (torch.Tensor): crops of shape (..., @num_crops, C, @crop_height, @crop_width) + if @pos_enc is False, otherwise (..., @num_crops, C + 2, @crop_height, @crop_width) + + crop_inds (torch.Tensor): sampled crop indices of shape (..., N, 2) + """ + device = images.device + + # maybe add 2 channels of spatial encoding to the source image + source_im = images + if pos_enc: + # spatial encoding [y, x] in [0, 1] + h, w = source_im.shape[-2:] + pos_y, pos_x = torch.meshgrid(torch.arange(h), torch.arange(w)) + pos_y = pos_y.float().to(device) / float(h) + pos_x = pos_x.float().to(device) / float(w) + position_enc = torch.stack((pos_y, pos_x)) # shape [C, H, W] + + # unsqueeze and expand to match leading dimensions -> shape [..., C, H, W] + leading_shape = source_im.shape[:-3] + position_enc = position_enc[(None,) * len(leading_shape)] + position_enc = position_enc.expand(*leading_shape, -1, -1, -1) + + # concat across channel dimension with input + source_im = torch.cat((source_im, position_enc), dim=-3) + + # make sure sample boundaries ensure crops are fully within the images + image_c, image_h, image_w = source_im.shape[-3:] + max_sample_h = image_h - crop_height + max_sample_w = image_w - crop_width + + # Sample crop locations for all tensor dimensions up to the last 3, which are [C, H, W]. + # Each gets @num_crops samples - typically this will just be the batch dimension (B), so + # we will sample [B, N] indices, but this supports having more than one leading dimension, + # or possibly no leading dimension. + # + # Trick: sample in [0, 1) with rand, then re-scale to [0, M) and convert to long to get sampled ints + crop_inds_h = (max_sample_h * torch.rand(*source_im.shape[:-3], num_crops).to(device)).long() + crop_inds_w = (max_sample_w * torch.rand(*source_im.shape[:-3], num_crops).to(device)).long() + crop_inds = torch.cat((crop_inds_h.unsqueeze(-1), crop_inds_w.unsqueeze(-1)), dim=-1) # shape [..., N, 2] + + crops = crop_image_from_indices( + images=source_im, + crop_indices=crop_inds, + crop_height=crop_height, + crop_width=crop_width, + ) + + return crops, crop_inds + + +class Modality: + """ + Observation Modality class to encapsulate necessary functions needed to + process observations of this modality + """ + # observation keys to associate with this modality + keys = set() + + # Custom processing function that should prepare raw observations of this modality for training + _custom_obs_processor = None + + # Custom unprocessing function that should prepare observations of this modality used during training for deployment + _custom_obs_unprocessor = None + + # Name of this modality -- must be set by subclass! + name = None + + def __init_subclass__(cls, **kwargs): + """ + Hook method to automatically register all valid subclasses so we can keep track of valid modalities + """ + assert cls.name is not None, f"Name of modality {cls.__name__} must be specified!" + register_obs_key(cls) + + @classmethod + def set_keys(cls, keys): + """ + Sets the observation keys associated with this modality. + + Args: + keys (list or set): observation keys to associate with this modality + """ + cls.keys = {k for k in keys} + + @classmethod + def add_keys(cls, keys): + """ + Adds the observation @keys associated with this modality to the current set of keys. + + Args: + keys (list or set): observation keys to add to associate with this modality + """ + for key in keys: + cls.keys.add(key) + + @classmethod + def set_obs_processor(cls, processor=None): + """ + Sets the processor for this observation modality. If @processor is set to None, then + the obs processor will use the default one (self.process_obs(...)). Otherwise, @processor + should be a function to process this corresponding observation modality. + + Args: + processor (function or None): If not None, should be function that takes in either a + np.array or torch.Tensor and output the processed array / tensor. If None, will reset + to the default processor (self.process_obs(...)) + """ + cls._custom_obs_processor = processor + + @classmethod + def set_obs_unprocessor(cls, unprocessor=None): + """ + Sets the unprocessor for this observation modality. If @unprocessor is set to None, then + the obs unprocessor will use the default one (self.unprocess_obs(...)). Otherwise, @unprocessor + should be a function to process this corresponding observation modality. + + Args: + unprocessor (function or None): If not None, should be function that takes in either a + np.array or torch.Tensor and output the unprocessed array / tensor. If None, will reset + to the default unprocessor (self.unprocess_obs(...)) + """ + cls._custom_obs_unprocessor = unprocessor + + @classmethod + def _default_obs_processor(cls, obs): + """ + Default processing function for this obs modality. + + Note that this function is overridden by self.custom_obs_processor (a function with identical inputs / outputs) + if it is not None. + + Args: + obs (np.array or torch.Tensor): raw observation, which may include a leading batch dimension + + Returns: + np.array or torch.Tensor: processed observation + """ + raise NotImplementedError + + @classmethod + def _default_obs_unprocessor(cls, obs): + """ + Default unprocessing function for this obs modality. + + Note that this function is overridden by self.custom_obs_unprocessor + (a function with identical inputs / outputs) if it is not None. + + Args: + obs (np.array or torch.Tensor): processed observation, which may include a leading batch dimension + + Returns: + np.array or torch.Tensor: unprocessed observation + """ + raise NotImplementedError + + @classmethod + def process_obs(cls, obs): + """ + Prepares an observation @obs of this modality for network input. + + Args: + obs (np.array or torch.Tensor): raw observation, which may include a leading batch dimension + + Returns: + np.array or torch.Tensor: processed observation + """ + processor = cls._custom_obs_processor if \ + cls._custom_obs_processor is not None else cls._default_obs_processor + return processor(obs) + + @classmethod + def unprocess_obs(cls, obs): + """ + Prepares an observation @obs of this modality for deployment. + + Args: + obs (np.array or torch.Tensor): processed observation, which may include a leading batch dimension + + Returns: + np.array or torch.Tensor: unprocessed observation + """ + unprocessor = cls._custom_obs_unprocessor if \ + cls._custom_obs_unprocessor is not None else cls._default_obs_unprocessor + return unprocessor(obs) + + @classmethod + def process_obs_from_dict(cls, obs_dict, inplace=True): + """ + Receives a dictionary of keyword mapped observations @obs_dict, and processes the observations with keys + corresponding to this modality. A copy will be made of the received dictionary unless @inplace is True + + Args: + obs_dict (dict): Dictionary mapping observation keys to observations + inplace (bool): If True, will modify @obs_dict in place, otherwise, will create a copy + + Returns: + dict: observation dictionary with processed observations corresponding to this modality + """ + if inplace: + obs_dict = deepcopy(obs_dict) + # Loop over all keys and process the ones corresponding to this modality + for key, obs in obs_dict.values(): + if key in cls.keys: + obs_dict[key] = cls.process_obs(obs) + + return obs_dict + + +class ImageModality(Modality): + """ + Modality for RGB image observations + """ + name = "rgb" + + @classmethod + def _default_obs_processor(cls, obs): + """ + Given image fetched from dataset, process for network input. Converts array + to float (from uint8), normalizes pixels from range [0, 255] to [0, 1], and channel swaps + from (H, W, C) to (C, H, W). + + Args: + obs (np.array or torch.Tensor): image array + + Returns: + processed_obs (np.array or torch.Tensor): processed image + """ + return process_frame(frame=obs, channel_dim=3, scale=255.) + + @classmethod + def _default_obs_unprocessor(cls, obs): + """ + Given image prepared for network input, prepare for saving to dataset. + Inverse of @process_frame. + + Args: + obs (np.array or torch.Tensor): image array + + Returns: + unprocessed_obs (np.array or torch.Tensor): image passed through + inverse operation of @process_frame + """ + return TU.to_uint8(unprocess_frame(frame=obs, channel_dim=3, scale=255.)) + + +class DepthModality(Modality): + """ + Modality for depth observations + """ + name = "depth" + + @classmethod + def _default_obs_processor(cls, obs): + """ + Given depth fetched from dataset, process for network input. Converts array + to float (from uint8), normalizes pixels from range [0, 1] to [0, 1], and channel swaps + from (H, W, C) to (C, H, W). + + Args: + obs (np.array or torch.Tensor): depth array + + Returns: + processed_obs (np.array or torch.Tensor): processed depth + """ + return process_frame(frame=obs, channel_dim=1, scale=1.) + + @classmethod + def _default_obs_unprocessor(cls, obs): + """ + Given depth prepared for network input, prepare for saving to dataset. + Inverse of @process_depth. + + Args: + obs (np.array or torch.Tensor): depth array + + Returns: + unprocessed_obs (np.array or torch.Tensor): depth passed through + inverse operation of @process_depth + """ + return unprocess_frame(frame=obs, channel_dim=1, scale=1.) + + +class ScanModality(Modality): + """ + Modality for scan observations + """ + name = "scan" + + @classmethod + def _default_obs_processor(cls, obs): + return obs + + @classmethod + def _default_obs_unprocessor(cls, obs): + return obs + + +class LowDimModality(Modality): + """ + Modality for low dimensional observations + """ + name = "low_dim" + + @classmethod + def _default_obs_processor(cls, obs): + return obs + + @classmethod + def _default_obs_unprocessor(cls, obs): + return obs diff --git a/aloha-devel/robomimic/utils/python_utils.py b/aloha-devel/robomimic/utils/python_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5bc71bd1aaaf08bb406f3a72e83886d86c0d19a6 --- /dev/null +++ b/aloha-devel/robomimic/utils/python_utils.py @@ -0,0 +1,73 @@ +""" +Set of general purpose utility functions for easier interfacing with Python API +""" +import inspect +from copy import deepcopy +import robomimic.macros as Macros + + +def get_class_init_kwargs(cls): + """ + Helper function to return a list of all valid keyword arguments (excluding "self") for the given @cls class. + + Args: + cls (object): Class from which to grab __init__ kwargs + + Returns: + list: All keyword arguments (excluding "self") specified by @cls __init__ constructor method + """ + return list(inspect.signature(cls.__init__).parameters.keys())[1:] + + +def extract_subset_dict(dic, keys, copy=False): + """ + Helper function to extract a subset of dictionary key-values from a current dictionary. Optionally (deep)copies + the values extracted from the original @dic if @copy is True. + + Args: + dic (dict): Dictionary containing multiple key-values + keys (Iterable): Specific keys to extract from @dic. If the key doesn't exist in @dic, then the key is skipped + copy (bool): If True, will deepcopy all values corresponding to the specified @keys + + Returns: + dict: Extracted subset dictionary containing only the specified @keys and their corresponding values + """ + subset = {k: dic[k] for k in keys if k in dic} + return deepcopy(subset) if copy else subset + + +def extract_class_init_kwargs_from_dict(cls, dic, copy=False, verbose=False): + """ + Helper function to return a dictionary of key-values that specifically correspond to @cls class's __init__ + constructor method, from @dic which may or may not contain additional, irrelevant kwargs. + + Note that @dic may possibly be missing certain kwargs as specified by cls.__init__. No error will be raised. + + Args: + cls (object): Class from which to grab __init__ kwargs that will be be used as filtering keys for @dic + dic (dict): Dictionary containing multiple key-values + copy (bool): If True, will deepcopy all values corresponding to the specified @keys + verbose (bool): If True (or if macro DEBUG is True), then will print out mismatched keys + + Returns: + dict: Extracted subset dictionary possibly containing only the specified keys from cls.__init__ and their + corresponding values + """ + # extract only relevant kwargs for this specific backbone + cls_keys = get_class_init_kwargs(cls) + subdic = extract_subset_dict( + dic=dic, + keys=cls_keys, + copy=copy, + ) + + # Run sanity check if verbose or debugging + if verbose or Macros.DEBUG: + keys_not_in_cls = [k for k in dic if k not in cls_keys] + keys_not_in_dic = [k for k in cls_keys if k not in list(dic.keys())] + if len(keys_not_in_cls) > 0: + print(f"Warning: For class {cls.__name__}, got unknown keys: {keys_not_in_cls} ") + if len(keys_not_in_dic) > 0: + print(f"Warning: For class {cls.__name__}, got missing keys: {keys_not_in_dic} ") + + return subdic \ No newline at end of file diff --git a/aloha-devel/robomimic/utils/tensor_utils.py b/aloha-devel/robomimic/utils/tensor_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ec2063b24c5d39333645edcf53a8468ab28a24df --- /dev/null +++ b/aloha-devel/robomimic/utils/tensor_utils.py @@ -0,0 +1,960 @@ +""" +A collection of utilities for working with nested tensor structures consisting +of numpy arrays and torch tensors. +""" +import collections +import numpy as np +import torch + + +def recursive_dict_list_tuple_apply(x, type_func_dict): + """ + Recursively apply functions to a nested dictionary or list or tuple, given a dictionary of + {data_type: function_to_apply}. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + type_func_dict (dict): a mapping from data types to the functions to be + applied for each data type. + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + assert(list not in type_func_dict) + assert(tuple not in type_func_dict) + assert(dict not in type_func_dict) + + if isinstance(x, (dict, collections.OrderedDict)): + new_x = collections.OrderedDict() if isinstance(x, collections.OrderedDict) else dict() + for k, v in x.items(): + new_x[k] = recursive_dict_list_tuple_apply(v, type_func_dict) + return new_x + elif isinstance(x, (list, tuple)): + ret = [recursive_dict_list_tuple_apply(v, type_func_dict) for v in x] + if isinstance(x, tuple): + ret = tuple(ret) + return ret + else: + for t, f in type_func_dict.items(): + if isinstance(x, t): + return f(x) + else: + raise NotImplementedError( + 'Cannot handle data type %s' % str(type(x))) + + +def map_tensor(x, func): + """ + Apply function @func to torch.Tensor objects in a nested dictionary or + list or tuple. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + func (function): function to apply to each tensor + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: func, + type(None): lambda x: x, + } + ) + + +def map_ndarray(x, func): + """ + Apply function @func to np.ndarray objects in a nested dictionary or + list or tuple. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + func (function): function to apply to each array + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + np.ndarray: func, + type(None): lambda x: x, + } + ) + + +def map_tensor_ndarray(x, tensor_func, ndarray_func): + """ + Apply function @tensor_func to torch.Tensor objects and @ndarray_func to + np.ndarray objects in a nested dictionary or list or tuple. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + tensor_func (function): function to apply to each tensor + ndarray_Func (function): function to apply to each array + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: tensor_func, + np.ndarray: ndarray_func, + type(None): lambda x: x, + } + ) + + +def clone(x): + """ + Clones all torch tensors and numpy arrays in nested dictionary or list + or tuple and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x.clone(), + np.ndarray: lambda x: x.copy(), + type(None): lambda x: x, + } + ) + + +def detach(x): + """ + Detaches all torch tensors in nested dictionary or list + or tuple and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x.detach(), + } + ) + + +def to_batch(x): + """ + Introduces a leading batch dimension of 1 for all torch tensors and numpy + arrays in nested dictionary or list or tuple and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x[None, ...], + np.ndarray: lambda x: x[None, ...], + type(None): lambda x: x, + } + ) + + +def to_sequence(x): + """ + Introduces a time dimension of 1 at dimension 1 for all torch tensors and numpy + arrays in nested dictionary or list or tuple and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x[:, None, ...], + np.ndarray: lambda x: x[:, None, ...], + type(None): lambda x: x, + } + ) + + +def index_at_time(x, ind): + """ + Indexes all torch tensors and numpy arrays in dimension 1 with index @ind in + nested dictionary or list or tuple and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + ind (int): index + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x[:, ind, ...], + np.ndarray: lambda x: x[:, ind, ...], + type(None): lambda x: x, + } + ) + + +def unsqueeze(x, dim): + """ + Adds dimension of size 1 at dimension @dim in all torch tensors and numpy arrays + in nested dictionary or list or tuple and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + dim (int): dimension + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x.unsqueeze(dim=dim), + np.ndarray: lambda x: np.expand_dims(x, axis=dim), + type(None): lambda x: x, + } + ) + + +def contiguous(x): + """ + Makes all torch tensors and numpy arrays contiguous in nested dictionary or + list or tuple and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x.contiguous(), + np.ndarray: lambda x: np.ascontiguousarray(x), + type(None): lambda x: x, + } + ) + + +def to_device(x, device): + """ + Sends all torch tensors in nested dictionary or list or tuple to device + @device, and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + device (torch.Device): device to send tensors to + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x, d=device: x.to(d), + type(None): lambda x: x, + } + ) + + +def to_tensor(x): + """ + Converts all numpy arrays in nested dictionary or list or tuple to + torch tensors (and leaves existing torch Tensors as-is), and returns + a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x, + np.ndarray: lambda x: torch.from_numpy(x), + type(None): lambda x: x, + } + ) + + +def to_numpy(x): + """ + Converts all torch tensors in nested dictionary or list or tuple to + numpy (and leaves existing numpy arrays as-is), and returns + a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + def f(tensor): + if tensor.is_cuda: + return tensor.detach().cpu().numpy() + else: + return tensor.detach().numpy() + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: f, + np.ndarray: lambda x: x, + type(None): lambda x: x, + } + ) + + +def to_list(x): + """ + Converts all torch tensors and numpy arrays in nested dictionary or list + or tuple to a list, and returns a new nested structure. Useful for + json encoding. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + def f(tensor): + if tensor.is_cuda: + return tensor.detach().cpu().numpy().tolist() + else: + return tensor.detach().numpy().tolist() + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: f, + np.ndarray: lambda x: x.tolist(), + type(None): lambda x: x, + } + ) + + +def to_float(x): + """ + Converts all torch tensors and numpy arrays in nested dictionary or list + or tuple to float type entries, and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x.float(), + np.ndarray: lambda x: x.astype(np.float32), + type(None): lambda x: x, + } + ) + + +def to_uint8(x): + """ + Converts all torch tensors and numpy arrays in nested dictionary or list + or tuple to uint8 type entries, and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x.byte(), + np.ndarray: lambda x: x.astype(np.uint8), + type(None): lambda x: x, + } + ) + + +def to_torch(x, device): + """ + Converts all numpy arrays and torch tensors in nested dictionary or list or tuple to + torch tensors on device @device and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + device (torch.Device): device to send tensors to + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return to_device(to_float(to_tensor(x)), device) + + +def to_one_hot_single(tensor, num_class): + """ + Convert tensor to one-hot representation, assuming a certain number of total class labels. + + Args: + tensor (torch.Tensor): tensor containing integer labels + num_class (int): number of classes + + Returns: + x (torch.Tensor): tensor containing one-hot representation of labels + """ + x = torch.zeros(tensor.size() + (num_class,)).to(tensor.device) + x.scatter_(-1, tensor.unsqueeze(-1), 1) + return x + + +def to_one_hot(tensor, num_class): + """ + Convert all tensors in nested dictionary or list or tuple to one-hot representation, + assuming a certain number of total class labels. + + Args: + tensor (dict or list or tuple): a possibly nested dictionary or list or tuple + num_class (int): number of classes + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return map_tensor(tensor, func=lambda x, nc=num_class: to_one_hot_single(x, nc)) + + +def flatten_single(x, begin_axis=1): + """ + Flatten a tensor in all dimensions from @begin_axis onwards. + + Args: + x (torch.Tensor): tensor to flatten + begin_axis (int): which axis to flatten from + + Returns: + y (torch.Tensor): flattened tensor + """ + fixed_size = x.size()[:begin_axis] + _s = list(fixed_size) + [-1] + return x.reshape(*_s) + + +def flatten(x, begin_axis=1): + """ + Flatten all tensors in nested dictionary or list or tuple, from @begin_axis onwards. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + begin_axis (int): which axis to flatten from + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x, b=begin_axis: flatten_single(x, begin_axis=b), + } + ) + + +def reshape_dimensions_single(x, begin_axis, end_axis, target_dims): + """ + Reshape selected dimensions in a tensor to a target dimension. + + Args: + x (torch.Tensor): tensor to reshape + begin_axis (int): begin dimension + end_axis (int): end dimension (inclusive) + target_dims (tuple or list): target shape for the range of dimensions + (@begin_axis, @end_axis) + + Returns: + y (torch.Tensor): reshaped tensor + """ + assert(begin_axis <= end_axis) + assert(begin_axis >= 0) + assert(end_axis < len(x.shape)) + assert(isinstance(target_dims, (tuple, list))) + s = x.shape + final_s = [] + for i in range(len(s)): + if i == begin_axis: + final_s.extend(target_dims) + elif i < begin_axis or i > end_axis: + final_s.append(s[i]) + return x.reshape(*final_s) + + +def reshape_dimensions(x, begin_axis, end_axis, target_dims): + """ + Reshape selected dimensions for all tensors in nested dictionary or list or tuple + to a target dimension. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + begin_axis (int): begin dimension + end_axis (int): end dimension (inclusive) + target_dims (tuple or list): target shape for the range of dimensions + (@begin_axis, @end_axis) + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x, b=begin_axis, e=end_axis, t=target_dims: reshape_dimensions_single( + x, begin_axis=b, end_axis=e, target_dims=t), + np.ndarray: lambda x, b=begin_axis, e=end_axis, t=target_dims: reshape_dimensions_single( + x, begin_axis=b, end_axis=e, target_dims=t), + type(None): lambda x: x, + } + ) + + +def join_dimensions(x, begin_axis, end_axis): + """ + Joins all dimensions between dimensions (@begin_axis, @end_axis) into a flat dimension, for + all tensors in nested dictionary or list or tuple. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + begin_axis (int): begin dimension + end_axis (int): end dimension + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x, b=begin_axis, e=end_axis: reshape_dimensions_single( + x, begin_axis=b, end_axis=e, target_dims=[-1]), + np.ndarray: lambda x, b=begin_axis, e=end_axis: reshape_dimensions_single( + x, begin_axis=b, end_axis=e, target_dims=[-1]), + type(None): lambda x: x, + } + ) + + +def expand_at_single(x, size, dim): + """ + Expand a tensor at a single dimension @dim by @size + + Args: + x (torch.Tensor): input tensor + size (int): size to expand + dim (int): dimension to expand + + Returns: + y (torch.Tensor): expanded tensor + """ + assert dim < x.ndimension() + assert x.shape[dim] == 1 + expand_dims = [-1] * x.ndimension() + expand_dims[dim] = size + return x.expand(*expand_dims) + + +def expand_at(x, size, dim): + """ + Expand all tensors in nested dictionary or list or tuple at a single + dimension @dim by @size. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + size (int): size to expand + dim (int): dimension to expand + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return map_tensor(x, lambda t, s=size, d=dim: expand_at_single(t, s, d)) + + +def unsqueeze_expand_at(x, size, dim): + """ + Unsqueeze and expand a tensor at a dimension @dim by @size. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + size (int): size to expand + dim (int): dimension to unsqueeze and expand + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + x = unsqueeze(x, dim) + return expand_at(x, size, dim) + + +def repeat_by_expand_at(x, repeats, dim): + """ + Repeat a dimension by combining expand and reshape operations. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + repeats (int): number of times to repeat the target dimension + dim (int): dimension to repeat on + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + x = unsqueeze_expand_at(x, repeats, dim + 1) + return join_dimensions(x, dim, dim + 1) + + +def named_reduce_single(x, reduction, dim): + """ + Reduce tensor at a dimension by named reduction functions. + + Args: + x (torch.Tensor): tensor to be reduced + reduction (str): one of ["sum", "max", "mean", "flatten"] + dim (int): dimension to be reduced (or begin axis for flatten) + + Returns: + y (torch.Tensor): reduced tensor + """ + assert x.ndimension() > dim + assert reduction in ["sum", "max", "mean", "flatten"] + if reduction == "flatten": + x = flatten(x, begin_axis=dim) + elif reduction == "max": + x = torch.max(x, dim=dim)[0] # [B, D] + elif reduction == "sum": + x = torch.sum(x, dim=dim) + else: + x = torch.mean(x, dim=dim) + return x + + +def named_reduce(x, reduction, dim): + """ + Reduces all tensors in nested dictionary or list or tuple at a dimension + using a named reduction function. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + reduction (str): one of ["sum", "max", "mean", "flatten"] + dim (int): dimension to be reduced (or begin axis for flatten) + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return map_tensor(x, func=lambda t, r=reduction, d=dim: named_reduce_single(t, r, d)) + + +def gather_along_dim_with_dim_single(x, target_dim, source_dim, indices): + """ + This function indexes out a target dimension of a tensor in a structured way, + by allowing a different value to be selected for each member of a flat index + tensor (@indices) corresponding to a source dimension. This can be interpreted + as moving along the source dimension, using the corresponding index value + in @indices to select values for all other dimensions outside of the + source and target dimensions. A common use case is to gather values + in target dimension 1 for each batch member (target dimension 0). + + Args: + x (torch.Tensor): tensor to gather values for + target_dim (int): dimension to gather values along + source_dim (int): dimension to hold constant and use for gathering values + from the other dimensions + indices (torch.Tensor): flat index tensor with same shape as tensor @x along + @source_dim + + Returns: + y (torch.Tensor): gathered tensor, with dimension @target_dim indexed out + """ + assert len(indices.shape) == 1 + assert x.shape[source_dim] == indices.shape[0] + + # unsqueeze in all dimensions except the source dimension + new_shape = [1] * x.ndimension() + new_shape[source_dim] = -1 + indices = indices.reshape(*new_shape) + + # repeat in all dimensions - but preserve shape of source dimension, + # and make sure target_dimension has singleton dimension + expand_shape = list(x.shape) + expand_shape[source_dim] = -1 + expand_shape[target_dim] = 1 + indices = indices.expand(*expand_shape) + + out = x.gather(dim=target_dim, index=indices) + return out.squeeze(target_dim) + + +def gather_along_dim_with_dim(x, target_dim, source_dim, indices): + """ + Apply @gather_along_dim_with_dim_single to all tensors in a nested + dictionary or list or tuple. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + target_dim (int): dimension to gather values along + source_dim (int): dimension to hold constant and use for gathering values + from the other dimensions + indices (torch.Tensor): flat index tensor with same shape as tensor @x along + @source_dim + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return map_tensor(x, + lambda y, t=target_dim, s=source_dim, i=indices: gather_along_dim_with_dim_single(y, t, s, i)) + + +def gather_sequence_single(seq, indices): + """ + Given a tensor with leading dimensions [B, T, ...], gather an element from each sequence in + the batch given an index for each sequence. + + Args: + seq (torch.Tensor): tensor with leading dimensions [B, T, ...] + indices (torch.Tensor): tensor indices of shape [B] + + Return: + y (torch.Tensor): indexed tensor of shape [B, ....] + """ + return gather_along_dim_with_dim_single(seq, target_dim=1, source_dim=0, indices=indices) + + +def gather_sequence(seq, indices): + """ + Given a nested dictionary or list or tuple, gathers an element from each sequence of the batch + for tensors with leading dimensions [B, T, ...]. + + Args: + seq (dict or list or tuple): a possibly nested dictionary or list or tuple with tensors + of leading dimensions [B, T, ...] + indices (torch.Tensor): tensor indices of shape [B] + + Returns: + y (dict or list or tuple): new nested dict-list-tuple with tensors of shape [B, ...] + """ + return gather_along_dim_with_dim(seq, target_dim=1, source_dim=0, indices=indices) + + +def pad_sequence_single(seq, padding, batched=False, pad_same=True, pad_values=None): + """ + Pad input tensor or array @seq in the time dimension (dimension 1). + + Args: + seq (np.ndarray or torch.Tensor): sequence to be padded + padding (tuple): begin and end padding, e.g. [1, 1] pads both begin and end of the sequence by 1 + batched (bool): if sequence has the batch dimension + pad_same (bool): if pad by duplicating + pad_values (scalar or (ndarray, Tensor)): values to be padded if not pad_same + + Returns: + padded sequence (np.ndarray or torch.Tensor) + """ + assert isinstance(seq, (np.ndarray, torch.Tensor)) + assert pad_same or pad_values is not None + if pad_values is not None: + assert isinstance(pad_values, float) + repeat_func = np.repeat if isinstance(seq, np.ndarray) else torch.repeat_interleave + concat_func = np.concatenate if isinstance(seq, np.ndarray) else torch.cat + ones_like_func = np.ones_like if isinstance(seq, np.ndarray) else torch.ones_like + seq_dim = 1 if batched else 0 + + begin_pad = [] + end_pad = [] + + if padding[0] > 0: + pad = seq[[0]] if pad_same else ones_like_func(seq[[0]]) * pad_values + begin_pad.append(repeat_func(pad, padding[0], seq_dim)) + if padding[1] > 0: + pad = seq[[-1]] if pad_same else ones_like_func(seq[[-1]]) * pad_values + end_pad.append(repeat_func(pad, padding[1], seq_dim)) + + return concat_func(begin_pad + [seq] + end_pad, seq_dim) + + +def pad_sequence(seq, padding, batched=False, pad_same=True, pad_values=None): + """ + Pad a nested dictionary or list or tuple of sequence tensors in the time dimension (dimension 1). + + Args: + seq (dict or list or tuple): a possibly nested dictionary or list or tuple with tensors + of leading dimensions [B, T, ...] + padding (tuple): begin and end padding, e.g. [1, 1] pads both begin and end of the sequence by 1 + batched (bool): if sequence has the batch dimension + pad_same (bool): if pad by duplicating + pad_values (scalar or (ndarray, Tensor)): values to be padded if not pad_same + + Returns: + padded sequence (dict or list or tuple) + """ + return recursive_dict_list_tuple_apply( + seq, + { + torch.Tensor: lambda x, p=padding, b=batched, ps=pad_same, pv=pad_values: + pad_sequence_single(x, p, b, ps, pv), + np.ndarray: lambda x, p=padding, b=batched, ps=pad_same, pv=pad_values: + pad_sequence_single(x, p, b, ps, pv), + type(None): lambda x: x, + } + ) + + +def assert_size_at_dim_single(x, size, dim, msg): + """ + Ensure that array or tensor @x has size @size in dim @dim. + + Args: + x (np.ndarray or torch.Tensor): input array or tensor + size (int): size that tensors should have at @dim + dim (int): dimension to check + msg (str): text to display if assertion fails + """ + assert x.shape[dim] == size, msg + + +def assert_size_at_dim(x, size, dim, msg): + """ + Ensure that arrays and tensors in nested dictionary or list or tuple have + size @size in dim @dim. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + size (int): size that tensors should have at @dim + dim (int): dimension to check + """ + map_tensor(x, lambda t, s=size, d=dim, m=msg: assert_size_at_dim_single(t, s, d, m)) + + +def get_shape(x): + """ + Get all shapes of arrays and tensors in nested dictionary or list or tuple. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple that contains each array or + tensor's shape + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x.shape, + np.ndarray: lambda x: x.shape, + type(None): lambda x: x, + } + ) + + +def list_of_flat_dict_to_dict_of_list(list_of_dict): + """ + Helper function to go from a list of flat dictionaries to a dictionary of lists. + By "flat" we mean that none of the values are dictionaries, but are numpy arrays, + floats, etc. + + Args: + list_of_dict (list): list of flat dictionaries + + Returns: + dict_of_list (dict): dictionary of lists + """ + assert isinstance(list_of_dict, list) + dic = collections.OrderedDict() + for i in range(len(list_of_dict)): + for k in list_of_dict[i]: + if k not in dic: + dic[k] = [] + dic[k].append(list_of_dict[i][k]) + return dic + + +def flatten_nested_dict_list(d, parent_key='', sep='_', item_key=''): + """ + Flatten a nested dict or list to a list. + + For example, given a dict + { + a: 1 + b: { + c: 2 + } + c: 3 + } + + the function would return [(a, 1), (b_c, 2), (c, 3)] + + Args: + d (dict, list): a nested dict or list to be flattened + parent_key (str): recursion helper + sep (str): separator for nesting keys + item_key (str): recursion helper + Returns: + list: a list of (key, value) tuples + """ + items = [] + if isinstance(d, (tuple, list)): + new_key = parent_key + sep + item_key if len(parent_key) > 0 else item_key + for i, v in enumerate(d): + items.extend(flatten_nested_dict_list(v, new_key, sep=sep, item_key=str(i))) + return items + elif isinstance(d, dict): + new_key = parent_key + sep + item_key if len(parent_key) > 0 else item_key + for k, v in d.items(): + assert isinstance(k, str) + items.extend(flatten_nested_dict_list(v, new_key, sep=sep, item_key=k)) + return items + else: + new_key = parent_key + sep + item_key if len(parent_key) > 0 else item_key + return [(new_key, d)] + + +def time_distributed(inputs, op, activation=None, inputs_as_kwargs=False, inputs_as_args=False, **kwargs): + """ + Apply function @op to all tensors in nested dictionary or list or tuple @inputs in both the + batch (B) and time (T) dimension, where the tensors are expected to have shape [B, T, ...]. + Will do this by reshaping tensors to [B * T, ...], passing through the op, and then reshaping + outputs to [B, T, ...]. + + Args: + inputs (list or tuple or dict): a possibly nested dictionary or list or tuple with tensors + of leading dimensions [B, T, ...] + op: a layer op that accepts inputs + activation: activation to apply at the output + inputs_as_kwargs (bool): whether to feed input as a kwargs dict to the op + inputs_as_args (bool) whether to feed input as a args list to the op + kwargs (dict): other kwargs to supply to the op + + Returns: + outputs (dict or list or tuple): new nested dict-list-tuple with tensors of leading dimension [B, T]. + """ + batch_size, seq_len = flatten_nested_dict_list(inputs)[0][1].shape[:2] + inputs = join_dimensions(inputs, 0, 1) + if inputs_as_kwargs: + outputs = op(**inputs, **kwargs) + elif inputs_as_args: + outputs = op(*inputs, **kwargs) + else: + outputs = op(inputs, **kwargs) + + if activation is not None: + outputs = map_tensor(outputs, activation) + outputs = reshape_dimensions(outputs, begin_axis=0, end_axis=0, target_dims=(batch_size, seq_len)) + return outputs diff --git a/aloha-devel/robomimic/utils/test_utils.py b/aloha-devel/robomimic/utils/test_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..86f125e0769951a21753d3f7f864b7d7184fb0a2 --- /dev/null +++ b/aloha-devel/robomimic/utils/test_utils.py @@ -0,0 +1,264 @@ +""" +Utilities for testing algorithm implementations - used mainly by scripts in tests directory. +""" +import os +import json +import shutil +import traceback +from termcolor import colored + +import numpy as np +import torch + +import robomimic +import robomimic.utils.file_utils as FileUtils +import robomimic.utils.torch_utils as TorchUtils +from robomimic.config import Config, config_factory +from robomimic.scripts.train import train + + +def maybe_remove_dir(dir_to_remove): + """ + Remove directory if it exists. + + Args: + dir_to_remove (str): path to directory to remove + """ + if os.path.exists(dir_to_remove): + shutil.rmtree(dir_to_remove) + + +def maybe_remove_file(file_to_remove): + """ + Remove file if it exists. + + Args: + file_to_remove (str): path to file to remove + """ + if os.path.exists(file_to_remove): + os.remove(file_to_remove) + + +def example_dataset_path(): + """ + Path to dataset to use for testing and example purposes. It should + exist under the tests/assets directory, and will be downloaded + from a server if it does not exist. + """ + dataset_folder = os.path.join(robomimic.__path__[0], "../tests/assets/") + dataset_path = os.path.join(dataset_folder, "test_v141.hdf5") + if not os.path.exists(dataset_path): + print("\nWARNING: test hdf5 does not exist! Downloading from server...") + os.makedirs(dataset_folder, exist_ok=True) + FileUtils.download_url( + url="http://downloads.cs.stanford.edu/downloads/rt_benchmark/test_v141.hdf5", + download_dir=dataset_folder, + ) + return dataset_path + + +def example_momart_dataset_path(): + """ + Path to momart dataset to use for testing and example purposes. It should + exist under the tests/assets directory, and will be downloaded + from a server if it does not exist. + """ + dataset_folder = os.path.join(robomimic.__path__[0], "../tests/assets/") + dataset_path = os.path.join(dataset_folder, "test_momart.hdf5") + if not os.path.exists(dataset_path): + user_response = input("\nWARNING: momart test hdf5 does not exist! We will download sample dataset. " + "This will take 0.6GB space. Proceed? y/n\n") + assert user_response.lower() in {"yes", "y"}, f"Did not receive confirmation. Aborting download." + + print("\nDownloading from server...") + + os.makedirs(dataset_folder, exist_ok=True) + FileUtils.download_url( + url="http://downloads.cs.stanford.edu/downloads/rt_mm/sample/test_momart.hdf5", + download_dir=dataset_folder, + ) + return dataset_path + + +def temp_model_dir_path(): + """ + Path to a temporary model directory to write to for testing and example purposes. + """ + return os.path.join(robomimic.__path__[0], "../tests/tmp_model_dir") + + +def temp_dataset_path(): + """ + Defines default dataset path to write to for testing. + """ + return os.path.join(robomimic.__path__[0], "../tests/", "tmp.hdf5") + + +def temp_video_path(): + """ + Defines default video path to write to for testing. + """ + return os.path.join(robomimic.__path__[0], "../tests/", "tmp.mp4") + + +def get_base_config(algo_name): + """ + Base config for testing algorithms. + + Args: + algo_name (str): name of algorithm - loads the corresponding json + from the config templates directory + """ + + # we will load and override defaults from template config + base_config_path = os.path.join(robomimic.__path__[0], "exps/templates/{}.json".format(algo_name)) + with open(base_config_path, 'r') as f: + config = Config(json.load(f)) + + # small dataset with a handful of trajectories + config.train.data = example_dataset_path() + + # temporary model dir + model_dir = temp_model_dir_path() + maybe_remove_dir(model_dir) + config.train.output_dir = model_dir + + # train and validate for 3 gradient steps + config.experiment.name = "test" + config.experiment.validate = True + config.experiment.epoch_every_n_steps = 3 + config.experiment.validation_epoch_every_n_steps = 3 + config.train.num_epochs = 1 + + # default train and validation filter keys + config.train.hdf5_filter_key = "train" + config.train.hdf5_validation_filter_key = "valid" + + # ensure model saving, rollout, and offscreen video rendering are tested too + config.experiment.save.enabled = True + config.experiment.save.every_n_epochs = 1 + config.experiment.rollout.enabled = True + config.experiment.rollout.rate = 1 + config.experiment.rollout.n = 1 + config.experiment.rollout.horizon = 10 + config.experiment.render_video = True + + # turn off logging to stdout, since that can interfere with testing code outputs + config.experiment.logging.terminal_output_to_txt = False + + # test cuda (if available) + config.train.cuda = True + + return config + + +def config_from_modifier(base_config, config_modifier): + """ + Helper function to load a base config, modify it using + the passed @config modifier function, and finalize it + for training. + + Args: + base_config (BaseConfig instance): starting config object that is + loaded (to change algorithm config defaults), and then modified + with @config_modifier + + config_modifier (function): function that takes a config object as + input, and modifies it + """ + + # algo name to default config for this algorithm + algo_name = base_config["algo_name"] + config = config_factory(algo_name) + + # update config with the settings specified in the base config + with config.unlocked(): + config.update(base_config) + + # modify the config and finalize it for training (no more modifications allowed) + config = config_modifier(config) + + return config + + +def checkpoint_path_from_test_run(): + """ + Helper function that gets the path of a model checkpoint after a test training run is finished. + """ + exp_dir = os.path.join(temp_model_dir_path(), "test") + time_dir_names = [f.name for f in os.scandir(exp_dir) if f.is_dir()] + assert len(time_dir_names) == 1 + path_to_models = os.path.join(exp_dir, time_dir_names[0], "models") + epoch_name = [f.name for f in os.scandir(path_to_models) if f.name.startswith("model")][0] + return os.path.join(path_to_models, epoch_name) + + +def test_eval_agent_from_checkpoint(ckpt_path, device): + """ + Test loading a model from checkpoint and running a rollout with the + trained agent for a small number of steps. + + Args: + ckpt_path (str): path to a checkpoint pth file + + device (torch.Device): torch device + """ + + # get policy and env from checkpoint + policy, ckpt_dict = FileUtils.policy_from_checkpoint(ckpt_path=ckpt_path, device=device, verbose=True) + env, _ = FileUtils.env_from_checkpoint(ckpt_dict=ckpt_dict, verbose=True) + + # run a test rollout + ob_dict = env.reset() + policy.start_episode() + for _ in range(15): + ac = policy(ob=ob_dict) + ob_dict, r, done, _ = env.step(ac) + + +def test_run(base_config, config_modifier): + """ + Takes a base_config and config_modifier (function that modifies a passed Config object) + and runs training as a test. It also takes the trained checkpoint, tries to load the + policy and environment from the checkpoint, and run an evaluation rollout. Returns + a string that is colored green if the run finished successfully without any issues, + and colored red if an error occurred. If an error occurs, the traceback is included + in the string. + + Args: + base_config (BaseConfig instance): starting config object that is + loaded (to change algorithm config defaults), and then modified + with @config_modifier + + config_modifier (function): function that takes a config object as + input, and modifies it + + Returns: + ret (str): a green "passed!" string, or a red "failed with error" string that contains + the traceback + """ + try: + # get config + config = config_from_modifier(base_config=base_config, config_modifier=config_modifier) + + # set torch device + device = TorchUtils.get_torch_device(try_to_use_cuda=config.train.cuda) + + # run training + train(config, device=device) + + # test evaluating a trained agent using saved checkpoint + ckpt_path = checkpoint_path_from_test_run() + test_eval_agent_from_checkpoint(ckpt_path, device=device) + + # indicate success + ret = colored("passed!", "green") + + except Exception as e: + # indicate failure by returning error string + ret = colored("failed with error:\n{}\n\n{}".format(e, traceback.format_exc()), "red") + + # make sure model directory is cleaned up before returning from this function + maybe_remove_dir(temp_model_dir_path()) + + return ret diff --git a/aloha-devel/robomimic/utils/torch_utils.py b/aloha-devel/robomimic/utils/torch_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b604621b20e3fb2d4ed33a1bfa3fb74fe3e88fcf --- /dev/null +++ b/aloha-devel/robomimic/utils/torch_utils.py @@ -0,0 +1,636 @@ +""" +This file contains some PyTorch utilities. +""" +import numpy as np +import torch +import torch.optim as optim +import torch.nn.functional as F + + +def soft_update(source, target, tau): + """ + Soft update from the parameters of a @source torch module to a @target torch module + with strength @tau. The update follows target = target * (1 - tau) + source * tau. + + Args: + source (torch.nn.Module): source network to push target network parameters towards + target (torch.nn.Module): target network to update + """ + for target_param, param in zip(target.parameters(), source.parameters()): + target_param.copy_( + target_param * (1.0 - tau) + param * tau + ) + + +def hard_update(source, target): + """ + Hard update @target parameters to match @source. + + Args: + source (torch.nn.Module): source network to provide parameters + target (torch.nn.Module): target network to update parameters for + """ + for target_param, param in zip(target.parameters(), source.parameters()): + target_param.copy_(param) + + +def get_torch_device(try_to_use_cuda): + """ + Return torch device. If using cuda (GPU), will also set cudnn.benchmark to True + to optimize CNNs. + + Args: + try_to_use_cuda (bool): if True and cuda is available, will use GPU + + Returns: + device (torch.Device): device to use for models + """ + if try_to_use_cuda and torch.cuda.is_available(): + torch.backends.cudnn.benchmark = True + device = torch.device("cuda:0") + else: + device = torch.device("cpu") + return device + + +def reparameterize(mu, logvar): + """ + Reparameterize for the backpropagation of z instead of q. + This makes it so that we can backpropagate through the sampling of z from + our encoder when feeding the sampled variable to the decoder. + + (See "The reparameterization trick" section of https://arxiv.org/abs/1312.6114) + + Args: + mu (torch.Tensor): batch of means from the encoder distribution + logvar (torch.Tensor): batch of log variances from the encoder distribution + + Returns: + z (torch.Tensor): batch of sampled latents from the encoder distribution that + support backpropagation + """ + # logvar = \log(\sigma^2) = 2 * \log(\sigma) + # \sigma = \exp(0.5 * logvar) + + # clamped for numerical stability + logstd = (0.5 * logvar).clamp(-4, 15) + std = torch.exp(logstd) + + # Sample \epsilon from normal distribution + # use std to create a new tensor, so we don't have to care + # about running on GPU or not + eps = std.new(std.size()).normal_() + + # Then multiply with the standard deviation and add the mean + z = eps.mul(std).add_(mu) + + return z + + +def optimizer_from_optim_params(net_optim_params, net): + """ + Helper function to return a torch Optimizer from the optim_params + section of the config for a particular network. + + Args: + optim_params (Config): optim_params part of algo_config corresponding + to @net. This determines the optimizer that is created. + + net (torch.nn.Module): module whose parameters this optimizer will be + responsible + + Returns: + optimizer (torch.optim.Optimizer): optimizer + """ + optimizer_type = net_optim_params.get("optimizer_type", "adam") + lr = net_optim_params["learning_rate"]["initial"] + + if optimizer_type == "adam": + return optim.Adam( + params=net.parameters(), + lr=lr, + weight_decay=net_optim_params["regularization"]["L2"], + ) + elif optimizer_type == "adamw": + return optim.AdamW( + params=net.parameters(), + lr=lr, + weight_decay=net_optim_params["regularization"]["L2"], + ) + + +def lr_scheduler_from_optim_params(net_optim_params, net, optimizer): + """ + Helper function to return a LRScheduler from the optim_params + section of the config for a particular network. Returns None + if a scheduler is not needed. + + Args: + optim_params (Config): optim_params part of algo_config corresponding + to @net. This determines whether a learning rate scheduler is created. + + net (torch.nn.Module): module whose parameters this optimizer will be + responsible + + optimizer (torch.optim.Optimizer): optimizer for this net + + Returns: + lr_scheduler (torch.optim.lr_scheduler or None): learning rate scheduler + """ + lr_scheduler_type = net_optim_params["learning_rate"].get("scheduler_type", "multistep") + epoch_schedule = net_optim_params["learning_rate"]["epoch_schedule"] + + lr_scheduler = None + if len(epoch_schedule) > 0: + if lr_scheduler_type == "linear": + assert len(epoch_schedule) == 1 + end_epoch = epoch_schedule[0] + + return optim.lr_scheduler.LinearLR( + optimizer, + start_factor=1.0, + end_factor=net_optim_params["learning_rate"]["decay_factor"], + total_iters=end_epoch, + ) + elif lr_scheduler_type == "multistep": + return optim.lr_scheduler.MultiStepLR( + optimizer=optimizer, + milestones=epoch_schedule, + gamma=net_optim_params["learning_rate"]["decay_factor"], + ) + else: + raise ValueError("Invalid LR scheduler type: {}".format(lr_scheduler_type)) + + return lr_scheduler + + +def backprop_for_loss(net, optim, loss, max_grad_norm=None, retain_graph=False): + """ + Backpropagate loss and update parameters for network with + name @name. + + Args: + net (torch.nn.Module): network to update + + optim (torch.optim.Optimizer): optimizer to use + + loss (torch.Tensor): loss to use for backpropagation + + max_grad_norm (float): if provided, used to clip gradients + + retain_graph (bool): if True, graph is not freed after backward call + + Returns: + grad_norms (float): average gradient norms from backpropagation + """ + + # backprop + optim.zero_grad() + loss.backward(retain_graph=retain_graph) + + # gradient clipping + if max_grad_norm is not None: + torch.nn.utils.clip_grad_norm_(net.parameters(), max_grad_norm) + + # compute grad norms + grad_norms = 0. + for p in net.parameters(): + # only clip gradients for parameters for which requires_grad is True + if p.grad is not None: + grad_norms += p.grad.data.norm(2).pow(2).item() + + # step + optim.step() + + return grad_norms + + +def rot_6d_to_axis_angle(rot_6d): + """ + Converts tensor with rot_6d representation to axis-angle representation. + """ + rot_mat = rotation_6d_to_matrix(rot_6d) + rot = matrix_to_axis_angle(rot_mat) + return rot + + +def rot_6d_to_euler_angles(rot_6d, convention="XYZ"): + """ + Converts tensor with rot_6d representation to euler representation. + """ + rot_mat = rotation_6d_to_matrix(rot_6d) + rot = matrix_to_euler_angles(rot_mat, convention=convention) + return rot + + +def axis_angle_to_rot_6d(axis_angle): + """ + Converts tensor with rot_6d representation to axis-angle representation. + """ + rot_mat = axis_angle_to_matrix(axis_angle) + rot_6d = matrix_to_rotation_6d(rot_mat) + return rot_6d + + +def euler_angles_to_rot_6d(euler_angles, convention="XYZ"): + """ + Converts tensor with rot_6d representation to euler representation. + """ + rot_mat = euler_angles_to_matrix(euler_angles, convention="XYZ") + rot_6d = matrix_to_rotation_6d(rot_mat) + return rot_6d + + +class dummy_context_mgr(): + """ + A dummy context manager - useful for having conditional scopes (such + as @maybe_no_grad). Nothing happens in this scope. + """ + def __enter__(self): + return None + def __exit__(self, exc_type, exc_value, traceback): + return False + + +def maybe_no_grad(no_grad): + """ + Args: + no_grad (bool): if True, the returned context will be torch.no_grad(), otherwise + it will be a dummy context + """ + return torch.no_grad() if no_grad else dummy_context_mgr() + + +""" +The following utility functions were taken from PyTorch3D: +https://github.com/facebookresearch/pytorch3d/blob/d84f274a0822da969668d00e831870fd88327845/pytorch3d/transforms/rotation_conversions.py +""" +def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor: + """ + Returns torch.sqrt(torch.max(0, x)) + but with a zero subgradient where x is 0. + """ + ret = torch.zeros_like(x) + positive_mask = x > 0 + ret[positive_mask] = torch.sqrt(x[positive_mask]) + return ret + + +def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as quaternions to rotation matrices. + Args: + quaternions: quaternions with real part first, + as tensor of shape (..., 4). + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + r, i, j, k = torch.unbind(quaternions, -1) + # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`. + two_s = 2.0 / (quaternions * quaternions).sum(-1) + + o = torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + -1, + ) + return o.reshape(quaternions.shape[:-1] + (3, 3)) + + +def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as rotation matrices to quaternions. + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + Returns: + quaternions with real part first, as tensor of shape (..., 4). + """ + if matrix.size(-1) != 3 or matrix.size(-2) != 3: + raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.") + + batch_dim = matrix.shape[:-2] + m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( + matrix.reshape(batch_dim + (9,)), dim=-1 + ) + + q_abs = _sqrt_positive_part( + torch.stack( + [ + 1.0 + m00 + m11 + m22, + 1.0 + m00 - m11 - m22, + 1.0 - m00 + m11 - m22, + 1.0 - m00 - m11 + m22, + ], + dim=-1, + ) + ) + + # we produce the desired quaternion multiplied by each of r, i, j, k + quat_by_rijk = torch.stack( + [ + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1), + ], + dim=-2, + ) + + # We floor here at 0.1 but the exact level is not important; if q_abs is small, + # the candidate won't be picked. + flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device) + quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr)) + + # if not for numerical problems, quat_candidates[i] should be same (up to a sign), + # forall i; we pick the best-conditioned one (with the largest denominator) + + return quat_candidates[ + F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, : + ].reshape(batch_dim + (4,)) + + +def axis_angle_to_matrix(axis_angle: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as axis/angle to rotation matrices. + Args: + axis_angle: Rotations given as a vector in axis angle form, + as a tensor of shape (..., 3), where the magnitude is + the angle turned anticlockwise in radians around the + vector's direction. + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle)) + + +def matrix_to_axis_angle(matrix: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as rotation matrices to axis/angle. + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + Returns: + Rotations given as a vector in axis angle form, as a tensor + of shape (..., 3), where the magnitude is the angle + turned anticlockwise in radians around the vector's + direction. + """ + return quaternion_to_axis_angle(matrix_to_quaternion(matrix)) + + +def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as axis/angle to quaternions. + Args: + axis_angle: Rotations given as a vector in axis angle form, + as a tensor of shape (..., 3), where the magnitude is + the angle turned anticlockwise in radians around the + vector's direction. + Returns: + quaternions with real part first, as tensor of shape (..., 4). + """ + angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True) + half_angles = angles * 0.5 + eps = 1e-6 + small_angles = angles.abs() < eps + sin_half_angles_over_angles = torch.empty_like(angles) + sin_half_angles_over_angles[~small_angles] = ( + torch.sin(half_angles[~small_angles]) / angles[~small_angles] + ) + # for x small, sin(x/2) is about x/2 - (x/2)^3/6 + # so sin(x/2)/x is about 1/2 - (x*x)/48 + sin_half_angles_over_angles[small_angles] = ( + 0.5 - (angles[small_angles] * angles[small_angles]) / 48 + ) + quaternions = torch.cat( + [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1 + ) + return quaternions + + +def quaternion_to_axis_angle(quaternions: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as quaternions to axis/angle. + Args: + quaternions: quaternions with real part first, + as tensor of shape (..., 4). + Returns: + Rotations given as a vector in axis angle form, as a tensor + of shape (..., 3), where the magnitude is the angle + turned anticlockwise in radians around the vector's + direction. + """ + norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True) + half_angles = torch.atan2(norms, quaternions[..., :1]) + angles = 2 * half_angles + eps = 1e-6 + small_angles = angles.abs() < eps + sin_half_angles_over_angles = torch.empty_like(angles) + sin_half_angles_over_angles[~small_angles] = ( + torch.sin(half_angles[~small_angles]) / angles[~small_angles] + ) + # for x small, sin(x/2) is about x/2 - (x/2)^3/6 + # so sin(x/2)/x is about 1/2 - (x*x)/48 + sin_half_angles_over_angles[small_angles] = ( + 0.5 - (angles[small_angles] * angles[small_angles]) / 48 + ) + return quaternions[..., 1:] / sin_half_angles_over_angles + + +def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor: + """ + Converts 6D rotation representation by Zhou et al. [1] to rotation matrix + using Gram--Schmidt orthogonalization per Section B of [1]. + Args: + d6: 6D rotation representation, of size (*, 6) + Returns: + batch of rotation matrices of size (*, 3, 3) + [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H. + On the Continuity of Rotation Representations in Neural Networks. + IEEE Conference on Computer Vision and Pattern Recognition, 2019. + Retrieved from http://arxiv.org/abs/1812.07035 + """ + + a1, a2 = d6[..., :3], d6[..., 3:] + b1 = F.normalize(a1, dim=-1) + b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1 + b2 = F.normalize(b2, dim=-1) + b3 = torch.cross(b1, b2, dim=-1) + return torch.stack((b1, b2, b3), dim=-2) + + +def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor: + """ + Converts rotation matrices to 6D rotation representation by Zhou et al. [1] + by dropping the last row. Note that 6D representation is not unique. + Args: + matrix: batch of rotation matrices of size (*, 3, 3) + Returns: + 6D rotation representation, of size (*, 6) + [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H. + On the Continuity of Rotation Representations in Neural Networks. + IEEE Conference on Computer Vision and Pattern Recognition, 2019. + Retrieved from http://arxiv.org/abs/1812.07035 + """ + batch_dim = matrix.size()[:-2] + return matrix[..., :2, :].clone().reshape(batch_dim + (6,)) + + +def matrix_to_euler_angles(matrix: torch.Tensor, convention: str) -> torch.Tensor: + """ + Convert rotations given as rotation matrices to Euler angles in radians. + + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + convention: Convention string of three uppercase letters. + + Returns: + Euler angles in radians as tensor of shape (..., 3). + """ + if len(convention) != 3: + raise ValueError("Convention must have 3 letters.") + if convention[1] in (convention[0], convention[2]): + raise ValueError(f"Invalid convention {convention}.") + for letter in convention: + if letter not in ("X", "Y", "Z"): + raise ValueError(f"Invalid letter {letter} in convention string.") + if matrix.size(-1) != 3 or matrix.size(-2) != 3: + raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.") + i0 = _index_from_letter(convention[0]) + i2 = _index_from_letter(convention[2]) + tait_bryan = i0 != i2 + if tait_bryan: + central_angle = torch.asin( + matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0) + ) + else: + central_angle = torch.acos(matrix[..., i0, i0]) + + o = ( + _angle_from_tan( + convention[0], convention[1], matrix[..., i2], False, tait_bryan + ), + central_angle, + _angle_from_tan( + convention[2], convention[1], matrix[..., i0, :], True, tait_bryan + ), + ) + return torch.stack(o, -1) + + +def euler_angles_to_matrix(euler_angles: torch.Tensor, convention: str) -> torch.Tensor: + """ + Convert rotations given as Euler angles in radians to rotation matrices. + + Args: + euler_angles: Euler angles in radians as tensor of shape (..., 3). + convention: Convention string of three uppercase letters from + {"X", "Y", and "Z"}. + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3: + raise ValueError("Invalid input euler angles.") + if len(convention) != 3: + raise ValueError("Convention must have 3 letters.") + if convention[1] in (convention[0], convention[2]): + raise ValueError(f"Invalid convention {convention}.") + for letter in convention: + if letter not in ("X", "Y", "Z"): + raise ValueError(f"Invalid letter {letter} in convention string.") + matrices = [ + _axis_angle_rotation(c, e) + for c, e in zip(convention, torch.unbind(euler_angles, -1)) + ] + # return functools.reduce(torch.matmul, matrices) + return torch.matmul(torch.matmul(matrices[0], matrices[1]), matrices[2]) + + +def _index_from_letter(letter: str) -> int: + if letter == "X": + return 0 + if letter == "Y": + return 1 + if letter == "Z": + return 2 + raise ValueError("letter must be either X, Y or Z.") + + +def _angle_from_tan( + axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool +) -> torch.Tensor: + """ + Extract the first or third Euler angle from the two members of + the matrix which are positive constant times its sine and cosine. + + Args: + axis: Axis label "X" or "Y or "Z" for the angle we are finding. + other_axis: Axis label "X" or "Y or "Z" for the middle axis in the + convention. + data: Rotation matrices as tensor of shape (..., 3, 3). + horizontal: Whether we are looking for the angle for the third axis, + which means the relevant entries are in the same row of the + rotation matrix. If not, they are in the same column. + tait_bryan: Whether the first and third axes in the convention differ. + + Returns: + Euler Angles in radians for each matrix in data as a tensor + of shape (...). + """ + + i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis] + if horizontal: + i2, i1 = i1, i2 + even = (axis + other_axis) in ["XY", "YZ", "ZX"] + if horizontal == even: + return torch.atan2(data[..., i1], data[..., i2]) + if tait_bryan: + return torch.atan2(-data[..., i2], data[..., i1]) + return torch.atan2(data[..., i2], -data[..., i1]) + + +def _axis_angle_rotation(axis: str, angle: torch.Tensor) -> torch.Tensor: + """ + Return the rotation matrices for one of the rotations about an axis + of which Euler angles describe, for each value of the angle given. + + Args: + axis: Axis label "X" or "Y or "Z". + angle: any shape tensor of Euler angles in radians + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + + cos = torch.cos(angle) + sin = torch.sin(angle) + one = torch.ones_like(angle) + zero = torch.zeros_like(angle) + + if axis == "X": + R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos) + elif axis == "Y": + R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos) + elif axis == "Z": + R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one) + else: + raise ValueError("letter must be either X, Y or Z.") + + return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3)) diff --git a/aloha-devel/robomimic/utils/train_utils.py b/aloha-devel/robomimic/utils/train_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b00e1905e4d373f91aa3326f11e1ba6a27901410 --- /dev/null +++ b/aloha-devel/robomimic/utils/train_utils.py @@ -0,0 +1,811 @@ +""" +This file contains several utility functions used to define the main training loop. It +mainly consists of functions to assist with logging, rollouts, and the @run_epoch function, +which is the core training logic for models in this repository. +""" +import os +import time +import datetime +import shutil +import json +import h5py +import imageio +import numpy as np +from copy import deepcopy +from collections import OrderedDict + +import torch + +import robomimic +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.log_utils as LogUtils +import robomimic.utils.file_utils as FileUtils + +from robomimic.utils.dataset import SequenceDataset, R2D2Dataset, MetaDataset +from robomimic.envs.env_base import EnvBase +from robomimic.envs.wrappers import EnvWrapper +from robomimic.algo import RolloutPolicy +from tianshou.env import SubprocVectorEnv + + +def get_exp_dir(config, auto_remove_exp_dir=False): + """ + Create experiment directory from config. If an identical experiment directory + exists and @auto_remove_exp_dir is False (default), the function will prompt + the user on whether to remove and replace it, or keep the existing one and + add a new subdirectory with the new timestamp for the current run. + + Args: + auto_remove_exp_dir (bool): if True, automatically remove the existing experiment + folder if it exists at the same path. + + Returns: + log_dir (str): path to created log directory (sub-folder in experiment directory) + output_dir (str): path to created models directory (sub-folder in experiment directory) + to store model checkpoints + video_dir (str): path to video directory (sub-folder in experiment directory) + to store rollout videos + """ + # timestamp for directory names + t_now = time.time() + time_str = datetime.datetime.fromtimestamp(t_now).strftime('%Y%m%d%H%M%S') + + # create directory for where to dump model parameters, tensorboard logs, and videos + base_output_dir = os.path.expanduser(config.train.output_dir) + if not os.path.isabs(base_output_dir): + # relative paths are specified relative to robomimic module location + base_output_dir = os.path.join(robomimic.__path__[0], base_output_dir) + base_output_dir = os.path.join(base_output_dir, config.experiment.name) + if os.path.exists(base_output_dir): + if not auto_remove_exp_dir: + ans = input("WARNING: model directory ({}) already exists! \noverwrite? (y/n)\n".format(base_output_dir)) + else: + ans = "y" + if ans == "y": + print("REMOVING") + shutil.rmtree(base_output_dir) + + # only make model directory if model saving is enabled + output_dir = None + if config.experiment.save.enabled: + output_dir = os.path.join(base_output_dir, time_str, "models") + os.makedirs(output_dir) + + # tensorboard directory + log_dir = os.path.join(base_output_dir, time_str, "logs") + os.makedirs(log_dir) + + # video directory + video_dir = os.path.join(base_output_dir, time_str, "videos") + os.makedirs(video_dir) + + # vis directory + vis_dir = os.path.join(base_output_dir, time_str, "vis") + os.makedirs(vis_dir) + + return log_dir, output_dir, video_dir, vis_dir + + +def load_data_for_training(config, obs_keys): + """ + Data loading at the start of an algorithm. + + Args: + config (BaseConfig instance): config object + obs_keys (list): list of observation modalities that are required for + training (this will inform the dataloader on what modalities to load) + + Returns: + train_dataset (SequenceDataset instance): train dataset object + valid_dataset (SequenceDataset instance): valid dataset object (only if using validation) + """ + + # config can contain an attribute to filter on + train_filter_by_attribute = config.train.hdf5_filter_key + valid_filter_by_attribute = config.train.hdf5_validation_filter_key + if valid_filter_by_attribute is not None: + assert config.experiment.validate, "specified validation filter key {}, but config.experiment.validate is not set".format(valid_filter_by_attribute) + + # load the dataset into memory + if config.experiment.validate: + assert not config.train.hdf5_normalize_obs, "no support for observation normalization with validation data yet" + assert (train_filter_by_attribute is not None) and (valid_filter_by_attribute is not None), \ + "did not specify filter keys corresponding to train and valid split in dataset" \ + " - please fill config.train.hdf5_filter_key and config.train.hdf5_validation_filter_key" + train_demo_keys = FileUtils.get_demos_for_filter_key( + hdf5_path=os.path.expanduser(config.train.data), + filter_key=train_filter_by_attribute, + ) + valid_demo_keys = FileUtils.get_demos_for_filter_key( + hdf5_path=os.path.expanduser(config.train.data), + filter_key=valid_filter_by_attribute, + ) + assert set(train_demo_keys).isdisjoint(set(valid_demo_keys)), "training demonstrations overlap with " \ + "validation demonstrations!" + train_dataset = dataset_factory(config, obs_keys, filter_by_attribute=train_filter_by_attribute) + valid_dataset = dataset_factory(config, obs_keys, filter_by_attribute=valid_filter_by_attribute) + else: + train_dataset = dataset_factory(config, obs_keys, filter_by_attribute=train_filter_by_attribute) + valid_dataset = None + + return train_dataset, valid_dataset + + +def dataset_factory(config, obs_keys, filter_by_attribute=None, dataset_path=None): + """ + Create a SequenceDataset instance to pass to a torch DataLoader. + + Args: + config (BaseConfig instance): config object + + obs_keys (list): list of observation modalities that are required for + training (this will inform the dataloader on what modalities to load) + + filter_by_attribute (str): if provided, use the provided filter key + to select a subset of demonstration trajectories to load + + dataset_path (str): if provided, the SequenceDataset instance should load + data from this dataset path. Defaults to config.train.data. + + Returns: + dataset (SequenceDataset instance): dataset object + """ + if dataset_path is None: + dataset_path = config.train.data + + ds_kwargs = dict( + hdf5_path=dataset_path, + obs_keys=obs_keys, + action_keys=config.train.action_keys, + dataset_keys=config.train.dataset_keys, + action_config=config.train.action_config, + load_next_obs=config.train.hdf5_load_next_obs, # whether to load next observations (s') from dataset + frame_stack=config.train.frame_stack, + seq_length=config.train.seq_length, + pad_frame_stack=config.train.pad_frame_stack, + pad_seq_length=config.train.pad_seq_length, + get_pad_mask=True, + goal_mode=config.train.goal_mode, + hdf5_cache_mode=config.train.hdf5_cache_mode, + hdf5_use_swmr=config.train.hdf5_use_swmr, + hdf5_normalize_obs=config.train.hdf5_normalize_obs, + filter_by_attribute=filter_by_attribute, + shuffled_obs_key_groups=config.train.shuffled_obs_key_groups, + ) + + ds_kwargs["hdf5_path"] = [ds_cfg["path"] for ds_cfg in config.train.data] + ds_kwargs["filter_by_attribute"] = [ds_cfg.get("filter_key", filter_by_attribute) for ds_cfg in config.train.data] + ds_weights = [ds_cfg.get("weight", 1.0) for ds_cfg in config.train.data] + ds_langs = [ds_cfg.get("lang", "dummy") for ds_cfg in config.train.data] + + meta_ds_kwargs = dict() + + dataset = get_dataset( + ds_class=R2D2Dataset if config.train.data_format == "r2d2" else SequenceDataset, + ds_kwargs=ds_kwargs, + ds_weights=ds_weights, + ds_langs=ds_langs, + normalize_weights_by_ds_size=False, + meta_ds_class=MetaDataset, + meta_ds_kwargs=meta_ds_kwargs, + ) + + return dataset + + +def get_dataset( + ds_class, + ds_kwargs, + ds_weights, + ds_langs, + normalize_weights_by_ds_size, + meta_ds_class=MetaDataset, + meta_ds_kwargs=None, +): + ds_list = [] + for i in range(len(ds_weights)): + + ds_kwargs_copy = deepcopy(ds_kwargs) + + keys = ["hdf5_path", "filter_by_attribute"] + + for k in keys: + ds_kwargs_copy[k] = ds_kwargs[k][i] + + ds_kwargs_copy["lang"] = ds_langs[i] + + ds_list.append(ds_class(**ds_kwargs_copy)) + + if len(ds_weights) == 1: + ds = ds_list[0] + else: + if meta_ds_kwargs is None: + meta_ds_kwargs = dict() + ds = meta_ds_class( + datasets=ds_list, + ds_weights=ds_weights, + normalize_weights_by_ds_size=normalize_weights_by_ds_size, + **meta_ds_kwargs + ) + + return ds + + +def batchify_obs(obs_list): + """ + TODO: add comments + """ + keys = list(obs_list[0].keys()) + obs = { + k: np.stack([obs_list[i][k] for i in range(len(obs_list))]) for k in keys + } + + return obs + + +def run_rollout( + policy, + env, + horizon, + use_goals=False, + render=False, + video_writer=None, + video_skip=5, + terminate_on_success=False, + ): + """ + Runs a rollout in an environment with the current network parameters. + + Args: + policy (RolloutPolicy instance): policy to use for rollouts. + + env (EnvBase instance): environment to use for rollouts. + + horizon (int): maximum number of steps to roll the agent out for + + use_goals (bool): if True, agent is goal-conditioned, so provide goal observations from env + + render (bool): if True, render the rollout to the screen + + video_writer (imageio Writer instance): if not None, use video writer object to append frames at + rate given by @video_skip + + video_skip (int): how often to write video frame + + terminate_on_success (bool): if True, terminate episode early as soon as a success is encountered + + Returns: + results (dict): dictionary containing return, success rate, etc. + """ + assert isinstance(policy, RolloutPolicy) + assert isinstance(env, EnvBase) or isinstance(env, EnvWrapper) or isinstance(env, SubprocVectorEnv) + + batched = isinstance(env, SubprocVectorEnv) + + policy.start_episode() + + ob_dict = env.reset() + goal_dict = None + if use_goals: + # retrieve goal from the environment + goal_dict = env.get_goal() + + results = {} + video_count = 0 # video frame counter + + rews = [] + success = None #{ k: False for k in env.is_success() } # success metrics + + if batched: + end_step = [None for _ in range(len(env))] + else: + end_step = None + + if batched: + video_frames = [[] for _ in range(len(env))] + else: + video_frames = [] + + try: + for step_i in range(horizon): #LogUtils.tqdm(range(horizon)): + # get action from policy + if batched: + policy_ob = batchify_obs(ob_dict) + ac = policy(ob=policy_ob, goal=goal_dict, batched=True) #, return_ob=True) + else: + policy_ob = ob_dict + ac = policy(ob=policy_ob, goal=goal_dict) #, return_ob=True) + + # play action + ob_dict, r, done, info = env.step(ac) + + # render to screen + if render: + env.render(mode="human") + + # compute reward + rews.append(r) + + # cur_success_metrics = env.is_success() + if batched: + cur_success_metrics = TensorUtils.list_of_flat_dict_to_dict_of_list([info[i]["is_success"] for i in range(len(info))]) + cur_success_metrics = {k: np.array(v) for (k, v) in cur_success_metrics.items()} + else: + cur_success_metrics = info["is_success"] + + if success is None: + success = deepcopy(cur_success_metrics) + else: + for k in success: + success[k] = success[k] | cur_success_metrics[k] + + # visualization + if video_writer is not None: + if video_count % video_skip == 0: + if batched: + # frames = env.render(mode="rgb_array", height=video_height, width=video_width) + + frames = [] + policy_ob = deepcopy(policy_ob) + for env_i in range(len(env)): + cam_imgs = [] + for im_name in ["robot0_agentview_left_image", "robot0_agentview_right_image", "robot0_eye_in_hand_image"]: + im = TensorUtils.to_numpy( + policy_ob[im_name][env_i, -1] + ) + im = np.transpose(im, (1, 2, 0)) + if policy_ob.get("ret", None) is not None: + im_ret = TensorUtils.to_numpy( + policy_ob["ret"]["obs"][im_name][env_i,:,-1] + ) + im_ret = np.transpose(im_ret, (0, 2, 3, 1)) + im = np.concatenate((im, *im_ret), axis=0) + cam_imgs.append(im) + frame = np.concatenate(cam_imgs, axis=1) + frame = (frame * 255.0).astype(np.uint8) + frames.append(frame) + + for env_i in range(len(env)): + frame = frames[env_i] + video_frames[env_i].append(frame) + else: + frame = env.render(mode="rgb_array", height=512, width=512) + + # cam_imgs = [] + # for im_name in ["robot0_eye_in_hand_image", "robot0_agentview_right_image", "robot0_agentview_left_image"]: + # im_input = TensorUtils.to_numpy( + # policy_ob_dict[im_name][0,-1] + # ) + # im_ret = TensorUtils.to_numpy( + # policy_ob_dict["ret"]["obs"][im_name][0,:,-1] + # ) + # im_input = np.transpose(im_input, (1, 2, 0)) + # im_input = add_border_to_frame(im_input, border_size=3, color="black") + # im_ret = np.transpose(im_ret, (0, 2, 3, 1)) + # im = np.concatenate((im_input, *im_ret), axis=1) + # cam_imgs.append(im) + + # frame = np.concatenate(cam_imgs, axis=0) + video_frames.append(frame) + + video_count += 1 + + # break if done + if batched: + for env_i in range(len(env)): + if end_step[env_i] is not None: + continue + + if done[env_i] or (terminate_on_success and success["task"][env_i]): + end_step[env_i] = step_i + else: + if done or (terminate_on_success and success["task"]): + end_step = step_i + break + + except Exception as e: + print("WARNING: got rollout exception {}".format(e)) + + + if video_writer is not None: + if batched: + for env_i in range(len(video_frames)): + for frame in video_frames[env_i]: + video_writer.append_data(frame) + else: + for frame in video_frames: + video_writer.append_data(frame) + + if batched: + total_reward = np.zeros(len(env)) + rews = np.array(rews) + for env_i in range(len(env)): + end_step_env_i = end_step[env_i] or step_i + total_reward[env_i] = np.sum(rews[:end_step_env_i+1, env_i]) + end_step[env_i] = end_step_env_i + + results["Return"] = total_reward + results["Horizon"] = np.array(end_step) + 1 + results["Success_Rate"] = success["task"].astype(float) + else: + end_step = end_step or step_i + total_reward = np.sum(rews[:end_step + 1]) + + results["Return"] = total_reward + results["Horizon"] = end_step + 1 + results["Success_Rate"] = float(success["task"]) + + # log additional success metrics + for k in success: + if k != "task": + if batched: + results["{}_Success_Rate".format(k)] = success[k].astype(float) + else: + results["{}_Success_Rate".format(k)] = float(success[k]) + + return results + + +def rollout_with_stats( + policy, + envs, + horizon, + use_goals=False, + num_episodes=None, + render=False, + video_dir=None, + video_path=None, + epoch=None, + video_skip=5, + terminate_on_success=False, + verbose=False, + ): + """ + A helper function used in the train loop to conduct evaluation rollouts per environment + and summarize the results. + + Can specify @video_dir (to dump a video per environment) or @video_path (to dump a single video + for all environments). + + Args: + policy (RolloutPolicy instance): policy to use for rollouts. + + envs (dict): dictionary that maps env_name (str) to EnvBase instance. The policy will + be rolled out in each env. + + horizon (int): maximum number of steps to roll the agent out for + + use_goals (bool): if True, agent is goal-conditioned, so provide goal observations from env + + num_episodes (int): number of rollout episodes per environment + + render (bool): if True, render the rollout to the screen + + video_dir (str): if not None, dump rollout videos to this directory (one per environment) + + video_path (str): if not None, dump a single rollout video for all environments + + epoch (int): epoch number (used for video naming) + + video_skip (int): how often to write video frame + + terminate_on_success (bool): if True, terminate episode early as soon as a success is encountered + + verbose (bool): if True, print results of each rollout + + Returns: + all_rollout_logs (dict): dictionary of rollout statistics (e.g. return, success rate, ...) + averaged across all rollouts + + video_paths (dict): path to rollout videos for each environment + """ + assert isinstance(policy, RolloutPolicy) + + all_rollout_logs = OrderedDict() + + # handle paths and create writers for video writing + assert (video_path is None) or (video_dir is None), "rollout_with_stats: can't specify both video path and dir" + write_video = (video_path is not None) or (video_dir is not None) + video_paths = OrderedDict() + video_writers = OrderedDict() + if video_path is not None: + # a single video is written for all envs + video_paths = { k : video_path for k in envs } + video_writer = imageio.get_writer(video_path, fps=20) + video_writers = { k : video_writer for k in envs } + if video_dir is not None: + # video is written per env + video_str = "_epoch_{}.mp4".format(epoch) if epoch is not None else ".mp4" + video_paths = { k : os.path.join(video_dir, "{}{}".format(k, video_str)) for k in envs } + video_writers = { k : imageio.get_writer(video_paths[k], fps=20) for k in envs } + + for env_name, env in envs.items(): + env_video_writer = None + if write_video: + print("video writes to " + video_paths[env_name]) + env_video_writer = video_writers[env_name] + + batched = isinstance(env, SubprocVectorEnv) + + if batched: + env_name = env.get_env_attr(key="name", id=0)[0] + else: + env_name = env.name + + print("rollout: env={}, horizon={}, use_goals={}, num_episodes={}".format( + env_name, horizon, use_goals, num_episodes, + )) + rollout_logs = [] + if batched: + iterator = range(0, num_episodes, len(env)) + else: + iterator = range(num_episodes) + if not verbose: + iterator = LogUtils.custom_tqdm(iterator, total=num_episodes) + + num_success = 0 + for ep_i in iterator: + rollout_timestamp = time.time() + rollout_info = run_rollout( + policy=policy, + env=env, + horizon=horizon, + render=render, + use_goals=use_goals, + video_writer=env_video_writer, + video_skip=video_skip, + terminate_on_success=terminate_on_success, + ) + if batched: + rollout_info["time"] = [(time.time() - rollout_timestamp) / len(env)] * len(env) + + for env_i in range(len(env)): + rollout_logs.append({k: rollout_info[k][env_i] for k in rollout_info}) + num_success += np.sum(rollout_info["Success_Rate"]) + else: + rollout_info["time"] = time.time() - rollout_timestamp + + rollout_logs.append(rollout_info) + num_success += rollout_info["Success_Rate"] + + if verbose: + if batched: + raise NotImplementedError + print("Episode {}, horizon={}, num_success={}".format(ep_i + 1, horizon, num_success)) + print(json.dumps(rollout_info, sort_keys=True, indent=4)) + + if video_dir is not None: + # close this env's video writer (next env has it's own) + env_video_writer.close() + + # average metric across all episodes + rollout_logs = dict((k, [rollout_logs[i][k] for i in range(len(rollout_logs))]) for k in rollout_logs[0]) + rollout_logs_mean = dict((k, np.mean(v)) for k, v in rollout_logs.items()) + rollout_logs_mean["Time_Episode"] = np.sum(rollout_logs["time"]) / 60. # total time taken for rollouts in minutes + all_rollout_logs[env_name] = rollout_logs_mean + + if video_path is not None: + # close video writer that was used for all envs + video_writer.close() + + return all_rollout_logs, video_paths + + +def should_save_from_rollout_logs( + all_rollout_logs, + best_return, + best_success_rate, + epoch_ckpt_name, + save_on_best_rollout_return, + save_on_best_rollout_success_rate, + ): + """ + Helper function used during training to determine whether checkpoints and videos + should be saved. It will modify input attributes appropriately (such as updating + the best returns and success rates seen and modifying the epoch ckpt name), and + returns a dict with the updated statistics. + + Args: + all_rollout_logs (dict): dictionary of rollout results that should be consistent + with the output of @rollout_with_stats + + best_return (dict): dictionary that stores the best average rollout return seen so far + during training, for each environment + + best_success_rate (dict): dictionary that stores the best average success rate seen so far + during training, for each environment + + epoch_ckpt_name (str): what to name the checkpoint file - this name might be modified + by this function + + save_on_best_rollout_return (bool): if True, should save checkpoints that achieve a + new best rollout return + + save_on_best_rollout_success_rate (bool): if True, should save checkpoints that achieve a + new best rollout success rate + + Returns: + save_info (dict): dictionary that contains updated input attributes @best_return, + @best_success_rate, @epoch_ckpt_name, along with two additional attributes + @should_save_ckpt (True if should save this checkpoint), and @ckpt_reason + (string that contains the reason for saving the checkpoint) + """ + should_save_ckpt = False + ckpt_reason = None + for env_name in all_rollout_logs: + rollout_logs = all_rollout_logs[env_name] + + if rollout_logs["Return"] > best_return[env_name]: + best_return[env_name] = rollout_logs["Return"] + if save_on_best_rollout_return: + # save checkpoint if achieve new best return + epoch_ckpt_name += "_{}_return_{}".format(env_name, best_return[env_name]) + should_save_ckpt = True + ckpt_reason = "return" + + if rollout_logs["Success_Rate"] > best_success_rate[env_name]: + best_success_rate[env_name] = rollout_logs["Success_Rate"] + if save_on_best_rollout_success_rate: + # save checkpoint if achieve new best success rate + epoch_ckpt_name += "_{}_success_{}".format(env_name, best_success_rate[env_name]) + should_save_ckpt = True + ckpt_reason = "success" + + # return the modified input attributes + return dict( + best_return=best_return, + best_success_rate=best_success_rate, + epoch_ckpt_name=epoch_ckpt_name, + should_save_ckpt=should_save_ckpt, + ckpt_reason=ckpt_reason, + ) + + +def save_model(model, config, env_meta, shape_meta, ckpt_path, obs_normalization_stats=None, action_normalization_stats=None): + """ + Save model to a torch pth file. + + Args: + model (Algo instance): model to save + + config (BaseConfig instance): config to save + + env_meta (dict): env metadata for this training run + + shape_meta (dict): shape metdata for this training run + + ckpt_path (str): writes model checkpoint to this path + + obs_normalization_stats (dict): optionally pass a dictionary for observation + normalization. This should map observation keys to dicts + with a "mean" and "std" of shape (1, ...) where ... is the default + shape for the observation. + + action_normalization_stats (dict): TODO + """ + env_meta = deepcopy(env_meta) + shape_meta = deepcopy(shape_meta) + params = dict( + model=model.serialize(), + config=config.dump(), + algo_name=config.algo_name, + env_metadata=env_meta, + shape_metadata=shape_meta, + ) + if obs_normalization_stats is not None: + assert config.train.hdf5_normalize_obs + obs_normalization_stats = deepcopy(obs_normalization_stats) + params["obs_normalization_stats"] = TensorUtils.to_list(obs_normalization_stats) + if action_normalization_stats is not None: + action_normalization_stats = deepcopy(action_normalization_stats) + params["action_normalization_stats"] = TensorUtils.to_list(action_normalization_stats) + torch.save(params, ckpt_path) + print("save checkpoint to {}".format(ckpt_path)) + + +def run_epoch(model, data_loader, epoch, validate=False, num_steps=None, obs_normalization_stats=None): + """ + Run an epoch of training or validation. + + Args: + model (Algo instance): model to train + + data_loader (DataLoader instance): data loader that will be used to serve batches of data + to the model + + epoch (int): epoch number + + validate (bool): whether this is a training epoch or validation epoch. This tells the model + whether to do gradient steps or purely do forward passes. + + num_steps (int): if provided, this epoch lasts for a fixed number of batches (gradient steps), + otherwise the epoch is a complete pass through the training dataset + + obs_normalization_stats (dict or None): if provided, this should map observation keys to dicts + with a "mean" and "std" of shape (1, ...) where ... is the default + shape for the observation. + + Returns: + step_log_all (dict): dictionary of logged training metrics averaged across all batches + """ + epoch_timestamp = time.time() + if validate: + model.set_eval() + else: + model.set_train() + if num_steps is None: + num_steps = len(data_loader) + + step_log_all = [] + timing_stats = dict(Data_Loading=[], Process_Batch=[], Train_Batch=[], Log_Info=[]) + start_time = time.time() + + data_loader_iter = iter(data_loader) + for _ in LogUtils.custom_tqdm(range(num_steps)): + + # load next batch from data loader + try: + t = time.time() + batch = next(data_loader_iter) + except StopIteration: + # reset for next dataset pass + data_loader_iter = iter(data_loader) + t = time.time() + batch = next(data_loader_iter) + timing_stats["Data_Loading"].append(time.time() - t) + + # process batch for training + t = time.time() + input_batch = model.process_batch_for_training(batch) + input_batch = model.postprocess_batch_for_training(input_batch, obs_normalization_stats=obs_normalization_stats) + timing_stats["Process_Batch"].append(time.time() - t) + + # forward and backward pass + t = time.time() + info = model.train_on_batch(input_batch, epoch, validate=validate) + timing_stats["Train_Batch"].append(time.time() - t) + + # tensorboard logging + t = time.time() + step_log = model.log_info(info) + step_log_all.append(step_log) + timing_stats["Log_Info"].append(time.time() - t) + + # flatten and take the mean of the metrics + step_log_dict = {} + for i in range(len(step_log_all)): + for k in step_log_all[i]: + if k not in step_log_dict: + step_log_dict[k] = [] + step_log_dict[k].append(step_log_all[i][k]) + step_log_all = dict((k, float(np.mean(v))) for k, v in step_log_dict.items()) + + # add in timing stats + for k in timing_stats: + # sum across all training steps, and convert from seconds to minutes + step_log_all["Time_{}".format(k)] = np.sum(timing_stats[k]) / 60. + step_log_all["Time_Epoch"] = (time.time() - epoch_timestamp) / 60. + + return step_log_all + + +def is_every_n_steps(interval, current_step, skip_zero=False): + """ + Convenient function to check whether current_step is at the interval. + Returns True if current_step % interval == 0 and asserts a few corner cases (e.g., interval <= 0) + + Args: + interval (int): target interval + current_step (int): current step + skip_zero (bool): whether to skip 0 (return False at 0) + + Returns: + is_at_interval (bool): whether current_step is at the interval + """ + if interval is None: + return False + assert isinstance(interval, int) and interval > 0 + assert isinstance(current_step, int) and current_step >= 0 + if skip_zero and current_step == 0: + return False + return current_step % interval == 0 diff --git a/aloha-devel/robomimic/utils/vis_utils.py b/aloha-devel/robomimic/utils/vis_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..712e75e9e5d88a2d825f984b76838d8508373970 --- /dev/null +++ b/aloha-devel/robomimic/utils/vis_utils.py @@ -0,0 +1,146 @@ +""" +This file contains utility functions for visualizing image observations in the training pipeline. +These functions can be a useful debugging tool. +""" +import numpy as np +import matplotlib.pyplot as plt +import os + +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.obs_utils as ObsUtils + + +def image_tensor_to_numpy(image): + """ + Converts processed image tensors to numpy so that they can be saved to disk or video. + A useful utility function for visualizing images in the middle of training. + + Args: + image (torch.Tensor): images of shape [..., C, H, W] + + Returns: + image (np.array): converted images of shape [..., H, W, C] and type uint8 + """ + return TensorUtils.to_numpy( + ObsUtils.unprocess_image(image) + ).astype(np.uint8) + + +def image_to_disk(image, fname): + """ + Writes an image to disk. + + Args: + image (np.array): image of shape [H, W, 3] + fname (str): path to save image to + """ + image = Image.fromarray(image) + image.save(fname) + + +def image_tensor_to_disk(image, fname): + """ + Writes an image tensor to disk. Any leading batch dimensions are indexed out + with the first element. + + Args: + image (torch.Tensor): image of shape [..., C, H, W]. All leading dimensions + will be indexed out with the first element + fname (str): path to save image to + """ + # index out all leading dimensions before [C, H, W] + num_leading_dims = len(image.shape[:-3]) + for _ in range(num_leading_dims): + image = image[0] + image = image_tensor_to_numpy(image) + image_to_disk(image, fname) + + +def visualize_image_randomizer(original_image, randomized_image, randomizer_name=None): + """ + A function that visualizes the before and after of an image-based input randomizer + Args: + original_image: batch of original image shaped [B, H, W, 3] + randomized_image: randomized image shaped [B, N, H, W, 3]. N is the number of randomization per input sample + randomizer_name: (Optional) name of the randomizer + Returns: + None + """ + + B, N, H, W, C = randomized_image.shape + + # Create a grid of subplots with B rows and N+1 columns (1 for the original image, N for the randomized images) + fig, axes = plt.subplots(B, N + 1, figsize=(4 * (N + 1), 4 * B)) + + for i in range(B): + # Display the original image in the first column of each row + axes[i, 0].imshow(original_image[i]) + axes[i, 0].set_title("Original") + axes[i, 0].axis("off") + + # Display the randomized images in the remaining columns of each row + for j in range(N): + axes[i, j + 1].imshow(randomized_image[i, j]) + axes[i, j + 1].axis("off") + + title = randomizer_name if randomizer_name is not None else "Randomized" + fig.suptitle(title, fontsize=16) + + # Adjust the space between subplots for better visualization + plt.subplots_adjust(wspace=0.5, hspace=0.5) + + # Show the entire grid of subplots + plt.show() + + +def make_model_prediction_plot( + hdf5_path, + save_path, + images, + action_names, + actual_actions, + predicted_actions, +): + """ + TODO: documentation + actual_actions: (T, D) + predicted_actions: (T, D) + """ + image_keys = sorted(list(images.keys())) + action_dim = actual_actions.shape[1] + traj_length = len(actual_actions) + + # Plot + fig, axs = plt.subplots(len(images) + action_dim, 1, figsize=(30, (len(images) + action_dim) * 3)) + for i, image_key in enumerate(image_keys): + interval = int(traj_length/15) # plot `5` images + images[image_key] = images[image_key][::interval] + combined_images = np.concatenate(images[image_key], axis=1) + axs[i].imshow(combined_images) + if i == 0: + axs[i].set_title(hdf5_path + '\n' + image_key, fontsize=30) + else: + axs[i].set_title(image_key, fontsize=30) + axs[i].axis("off") + for dim in range(action_dim): + ax = axs[len(images)+dim] + ax.plot(range(traj_length), actual_actions[:, dim], label='Actual Action', color='blue') + ax.plot(range(traj_length), predicted_actions[:, dim], label='Predicted Action', color='red') + # ax.set_xlabel('Timestep') + # ax.set_ylabel('Action Dimension {}'.format(dim + 1)) + ax.set_title(action_names[dim], fontsize=30) + ax.xaxis.set_tick_params(labelsize=24) + ax.yaxis.set_tick_params(labelsize=24) + ax.legend(fontsize=20) + plt.subplots_adjust(left=0.05, right=0.95, top=0.95, bottom=0.05, wspace=0.3, hspace=0.6) + + # Save the figure with the specified path and filename + save_dir = os.path.dirname(save_path) + if not os.path.exists(save_dir): + os.makedirs(save_dir) + plt.savefig(save_path) + + fig.clear() + plt.close() + plt.cla() + plt.clf() \ No newline at end of file diff --git a/camera_ws/src/realsense-ros/realsense2_camera/debian/udev.em b/camera_ws/src/realsense-ros/realsense2_camera/debian/udev.em new file mode 100644 index 0000000000000000000000000000000000000000..ca9dd16dcf4cc2c19d10095cb1fbaf58851b35c2 --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_camera/debian/udev.em @@ -0,0 +1,81 @@ +##Version=1.1## +# Device rules for Intel RealSense devices (R200, F200, SR300 LR200, ZR300, D400, L500, T200) +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0a80", MODE:="0666", GROUP:="plugdev", RUN+="/usr/local/bin/usb-R200-in_udev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0a66", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0aa3", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0aa2", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0aa5", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0abf", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0acb", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0ad0", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="04b4", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0ad1", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0ad2", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0ad3", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0ad4", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0ad5", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0ad6", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0af2", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0af6", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0afe", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0aff", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b00", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b01", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b03", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b07", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b0c", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b0d", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b3a", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b3d", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b48", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b49", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b4b", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b4d", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b52", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b5b", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b5c", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b64", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b68", MODE:="0666", GROUP:="plugdev" + +# Intel RealSense recovery devices (DFU) +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0ab3", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0adb", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0adc", MODE:="0666", GROUP:="plugdev" +SUBSYSTEMS=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b55", MODE:="0666", GROUP:="plugdev" + +# Intel RealSense devices (Movidius, T265) +SUBSYSTEMS=="usb", ENV{DEVTYPE}=="usb_device", ATTRS{idVendor}=="8087", ATTRS{idProduct}=="0af3", MODE="0666", GROUP="plugdev" +SUBSYSTEMS=="usb", ENV{DEVTYPE}=="usb_device", ATTRS{idVendor}=="8087", ATTRS{idProduct}=="0b37", MODE="0666", GROUP="plugdev" +SUBSYSTEMS=="usb", ENV{DEVTYPE}=="usb_device", ATTRS{idVendor}=="03e7", ATTRS{idProduct}=="2150", MODE="0666", GROUP="plugdev" + +KERNEL=="iio*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0ad5", MODE:="0777", GROUP:="plugdev", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p'" +DRIVER=="hid_sensor_custom", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0ad5", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p && chmod 0777 /dev/%k'" +KERNEL=="iio*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0af2", MODE:="0777", GROUP:="plugdev", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p'" +DRIVER=="hid_sensor*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0af2", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p && chmod 0777 /dev/%k'" +KERNEL=="iio*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0afe", MODE:="0777", GROUP:="plugdev", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p'" +DRIVER=="hid_sensor_custom", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0afe", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p && chmod 0777 /dev/%k'" +KERNEL=="iio*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0aff", MODE:="0777", GROUP:="plugdev", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p'" +DRIVER=="hid_sensor_custom", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0aff", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p && chmod 0777 /dev/%k'" +KERNEL=="iio*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b00", MODE:="0777", GROUP:="plugdev", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p'" +DRIVER=="hid_sensor_custom", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b00", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p && chmod 0777 /dev/%k'" +KERNEL=="iio*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b01", MODE:="0777", GROUP:="plugdev", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p'" +DRIVER=="hid_sensor_custom", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b01", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p && chmod 0777 /dev/%k'" +KERNEL=="iio*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b3a", MODE:="0777", GROUP:="plugdev", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p'" +DRIVER=="hid_sensor*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b3a", RUN+="/bin/sh -c ' chmod -R 0777 /sys/%p && chmod 0777 /dev/%k'" +KERNEL=="iio*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b3d", MODE:="0777", GROUP:="plugdev", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p'" +DRIVER=="hid_sensor*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b3d", RUN+="/bin/sh -c ' chmod -R 0777 /sys/%p && chmod 0777 /dev/%k'" +KERNEL=="iio*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b4b", MODE:="0777", GROUP:="plugdev", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p'" +DRIVER=="hid_sensor*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b4b", RUN+="/bin/sh -c ' chmod -R 0777 /sys/%p && chmod 0777 /dev/%k'" +KERNEL=="iio*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b4d", MODE:="0777", GROUP:="plugdev", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p'" +DRIVER=="hid_sensor*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b4d", RUN+="/bin/sh -c ' chmod -R 0777 /sys/%p && chmod 0777 /dev/%k'" +KERNEL=="iio*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b5b", MODE:="0777", GROUP:="plugdev", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p'" +DRIVER=="hid_sensor*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b5b", RUN+="/bin/sh -c ' chmod -R 0777 /sys/%p && chmod 0777 /dev/%k'" +KERNEL=="iio*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b5c", MODE:="0777", GROUP:="plugdev", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p'" +DRIVER=="hid_sensor*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b5c", RUN+="/bin/sh -c ' chmod -R 0777 /sys/%p && chmod 0777 /dev/%k'" +KERNEL=="iio*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b64", MODE:="0777", GROUP:="plugdev", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p'" +DRIVER=="hid_sensor*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b64", RUN+="/bin/sh -c ' chmod -R 0777 /sys/%p && chmod 0777 /dev/%k'" +KERNEL=="iio*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b68", MODE:="0777", GROUP:="plugdev", RUN+="/bin/sh -c 'chmod -R 0777 /sys/%p'" +DRIVER=="hid_sensor*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b68", RUN+="/bin/sh -c ' chmod -R 0777 /sys/%p && chmod 0777 /dev/%k'" + +# For products with motion_module, if (kernels is 4.15 and up) and (device name is "accel_3d") wait, in another process, until (enable flag is set to 1 or 200 mSec passed) and then set it to 0. +KERNEL=="iio*", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0ad5|0afe|0aff|0b00|0b01|0b3a|0b3d|0b64|0b68", RUN+="/bin/sh -c '(major=`uname -r | cut -d \".\" -f1` && minor=`uname -r | cut -d \".\" -f2` && (([ $major -eq 4 ] && [ $minor -ge 15 ]) || [ $major -ge 5 ])) && (enamefile=/sys/%p/name && [ `cat $enamefile` = \"accel_3d\" ]) && enfile=/sys/%p/buffer/enable && echo \"COUNTER=0; while [ \$COUNTER -lt 20 ] && grep -q 0 $enfile; do sleep 0.01; COUNTER=\$((COUNTER+1)); done && echo 0 > $enfile\" | at now'" diff --git a/camera_ws/src/realsense-ros/realsense2_camera/scripts/rs2_listener.py b/camera_ws/src/realsense-ros/realsense2_camera/scripts/rs2_listener.py new file mode 100644 index 0000000000000000000000000000000000000000..3ac35723a5d54bd0ba7d6bf13db04d6dbde19e65 --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_camera/scripts/rs2_listener.py @@ -0,0 +1,293 @@ +import sys +import time +import rospy +from sensor_msgs.msg import Image as msg_Image +from sensor_msgs.msg import CompressedImage as msg_CompressedImage +from sensor_msgs.msg import PointCloud2 as msg_PointCloud2 +import sensor_msgs.point_cloud2 as pc2 +from sensor_msgs.msg import Imu as msg_Imu +import numpy as np +from cv_bridge import CvBridge, CvBridgeError +import inspect +import ctypes +import struct +import tf +try: + from theora_image_transport.msg import Packet as msg_theora +except Exception: + pass + + +def pc2_to_xyzrgb(point): + # Thanks to Panos for his code used in this function. + x, y, z = point[:3] + rgb = point[3] + + # cast float32 to int so that bitwise operations are possible + s = struct.pack('>f', rgb) + i = struct.unpack('>l', s)[0] + # you can get back the float value by the inverse operations + pack = ctypes.c_uint32(i).value + r = (pack & 0x00FF0000) >> 16 + g = (pack & 0x0000FF00) >> 8 + b = (pack & 0x000000FF) + return x, y, z, r, g, b + + +class CWaitForMessage: + def __init__(self, params={}): + self.result = None + + self.break_timeout = False + self.timeout = params.get('timeout_secs', -1) * 1e-3 + self.seq = params.get('seq', -1) + self.time = params.get('time', None) + self.node_name = params.get('node_name', 'rs2_listener') + self.bridge = CvBridge() + self.listener = None + self.prev_msg_time = 0 + self.fout = None + + + self.themes = {'depthStream': {'topic': '/camera/depth/image_rect_raw', 'callback': self.imageColorCallback, 'msg_type': msg_Image}, + 'colorStream': {'topic': '/camera/color/image_raw', 'callback': self.imageColorCallback, 'msg_type': msg_Image}, + 'pointscloud': {'topic': '/camera/depth/color/points', 'callback': self.pointscloudCallback, 'msg_type': msg_PointCloud2}, + 'alignedDepthInfra1': {'topic': '/camera/aligned_depth_to_infra1/image_raw', 'callback': self.imageColorCallback, 'msg_type': msg_Image}, + 'alignedDepthColor': {'topic': '/camera/aligned_depth_to_color/image_raw', 'callback': self.imageColorCallback, 'msg_type': msg_Image}, + 'static_tf': {'topic': '/camera/color/image_raw', 'callback': self.imageColorCallback, 'msg_type': msg_Image}, + 'accelStream': {'topic': '/camera/accel/sample', 'callback': self.imuCallback, 'msg_type': msg_Imu}, + } + + self.func_data = dict() + + def imuCallback(self, theme_name): + def _imuCallback(data): + if self.listener is None: + self.listener = tf.TransformListener() + self.prev_time = time.time() + self.func_data[theme_name].setdefault('value', []) + self.func_data[theme_name].setdefault('ros_value', []) + try: + frame_id = data.header.frame_id + value = data.linear_acceleration + + (trans,rot) = self.listener.lookupTransform('/camera_link', frame_id, rospy.Time(0)) + quat = tf.transformations.quaternion_matrix(rot) + point = np.matrix([value.x, value.y, value.z, 1], dtype='float32') + point.resize((4, 1)) + rotated = quat*point + rotated.resize(1,4) + rotated = np.array(rotated)[0][:3] + except Exception as e: + print(e) + return + self.func_data[theme_name]['value'].append(value) + self.func_data[theme_name]['ros_value'].append(rotated) + return _imuCallback + + def imageColorCallback(self, theme_name): + def _imageColorCallback(data): + self.prev_time = time.time() + self.func_data[theme_name].setdefault('avg', []) + self.func_data[theme_name].setdefault('ok_percent', []) + self.func_data[theme_name].setdefault('num_channels', []) + self.func_data[theme_name].setdefault('shape', []) + self.func_data[theme_name].setdefault('reported_size', []) + + try: + cv_image = self.bridge.imgmsg_to_cv2(data, data.encoding) + except CvBridgeError as e: + print(e) + return + channels = cv_image.shape[2] if len(cv_image.shape) > 2 else 1 + pyimg = np.asarray(cv_image) + + ok_number = (pyimg != 0).sum() + + self.func_data[theme_name]['avg'].append(pyimg.sum() / ok_number) + self.func_data[theme_name]['ok_percent'].append(float(ok_number) / (pyimg.shape[0] * pyimg.shape[1]) / channels) + self.func_data[theme_name]['num_channels'].append(channels) + self.func_data[theme_name]['shape'].append(cv_image.shape) + self.func_data[theme_name]['reported_size'].append((data.width, data.height, data.step)) + return _imageColorCallback + + def imageDepthCallback(self, data): + pass + + def pointscloudCallback(self, theme_name): + def _pointscloudCallback(data): + self.prev_time = time.time() + print ('Got pointcloud: %d, %d' % (data.width, data.height)) + + self.func_data[theme_name].setdefault('frame_counter', 0) + self.func_data[theme_name].setdefault('avg', []) + self.func_data[theme_name].setdefault('size', []) + self.func_data[theme_name].setdefault('width', []) + self.func_data[theme_name].setdefault('height', []) + # until parsing pointcloud is done in real time, I'll use only the first frame. + self.func_data[theme_name]['frame_counter'] += 1 + + if self.func_data[theme_name]['frame_counter'] == 1: + # Known issue - 1st pointcloud published has invalid texture. Skip 1st frame. + return + + try: + points = np.array([pc2_to_xyzrgb(pp) for pp in pc2.read_points(data, skip_nans=True, field_names=("x", "y", "z", "rgb")) if pp[0] > 0]) + except Exception as e: + print(e) + return + self.func_data[theme_name]['avg'].append(points.mean(0)) + self.func_data[theme_name]['size'].append(len(points)) + self.func_data[theme_name]['width'].append(data.width) + self.func_data[theme_name]['height'].append(data.height) + return _pointscloudCallback + + def wait_for_message(self, params, msg_type=msg_Image): + topic = params['topic'] + print ('connect to ROS with name: %s' % self.node_name) + rospy.init_node(self.node_name, anonymous=True) + + out_filename = params.get('filename', None) + if (out_filename): + self.fout = open(out_filename, 'w') + if msg_type is msg_Imu: + col_w = 20 + print ('Writing to file: %s' % out_filename) + columns = ['frame_number', 'frame_time(sec)', 'accel.x', 'accel.y', 'accel.z', 'gyro.x', 'gyro.y', 'gyro.z'] + line = ('{:<%d}'*len(columns) % (col_w, col_w, col_w, col_w, col_w, col_w, col_w, col_w)).format(*columns) + '\n' + sys.stdout.write(line) + self.fout.write(line) + + rospy.loginfo('Subscribing on topic: %s' % topic) + self.sub = rospy.Subscriber(topic, msg_type, self.callback) + + self.prev_time = time.time() + break_timeout = False + while not any([rospy.core.is_shutdown(), break_timeout, self.result]): + rospy.rostime.wallsleep(0.5) + if self.timeout > 0 and time.time() - self.prev_time > self.timeout: + break_timeout = True + self.sub.unregister() + + return self.result + + @staticmethod + def unregister_all(registers): + for test_name in registers: + rospy.loginfo('Un-Subscribing test %s' % test_name) + registers[test_name]['sub'].unregister() + + def wait_for_messages(self, themes): + # tests_params = {: {'callback', 'topic', 'msg_type', 'internal_params'}} + self.func_data = dict([[theme_name, {}] for theme_name in themes]) + + print ('connect to ROS with name: %s' % self.node_name) + rospy.init_node(self.node_name, anonymous=True) + for theme_name in themes: + theme = self.themes[theme_name] + rospy.loginfo('Subscribing %s on topic: %s' % (theme_name, theme['topic'])) + self.func_data[theme_name]['sub'] = rospy.Subscriber(theme['topic'], theme['msg_type'], theme['callback'](theme_name)) + + self.prev_time = time.time() + break_timeout = False + while not any([rospy.core.is_shutdown(), break_timeout]): + rospy.rostime.wallsleep(0.5) + if self.timeout > 0 and time.time() - self.prev_time > self.timeout: + break_timeout = True + self.unregister_all(self.func_data) + + return self.func_data + + def callback(self, data): + msg_time = data.header.stamp.secs + 1e-9 * data.header.stamp.nsecs + + if (self.prev_msg_time > msg_time): + rospy.loginfo('Out of order: %.9f > %.9f' % (self.prev_msg_time, msg_time)) + if type(data) == msg_Imu: + col_w = 20 + frame_number = data.header.seq + accel = data.linear_acceleration + gyro = data.angular_velocity + line = ('\n{:<%d}{:<%d.6f}{:<%d.4f}{:<%d.4f}{:<%d.4f}{:<%d.4f}{:<%d.4f}{:<%d.4f}' % (col_w, col_w, col_w, col_w, col_w, col_w, col_w, col_w)).format(frame_number, msg_time, accel.x, accel.y, accel.z, gyro.x, gyro.y, gyro.z) + sys.stdout.write(line) + if self.fout: + self.fout.write(line) + + self.prev_msg_time = msg_time + self.prev_msg_data = data + + self.prev_time = time.time() + if any([self.seq < 0 and self.time is None, + self.seq > 0 and data.header.seq >= self.seq, + self.time and data.header.stamp.secs == self.time['secs'] and data.header.stamp.nsecs == self.time['nsecs']]): + self.result = data + self.sub.unregister() + + + +def main(): + if len(sys.argv) < 2 or '--help' in sys.argv or '/?' in sys.argv: + print ('USAGE:') + print ('------') + print ('rs2_listener.py [Options]') + print ('example: rs2_listener.py /camera/color/image_raw --time 1532423022.044515610 --timeout 3') + print ('example: rs2_listener.py pointscloud') + print ('') + print ('Application subscribes on , wait for the first message matching [Options].') + print ('When found, prints the timestamp.') + print + print ('[Options:]') + print ('-s ') + print ('--time ') + print ('--timeout ') + print ('--filename : write output to file') + exit(-1) + + # wanted_topic = '/device_0/sensor_0/Depth_0/image/data' + # wanted_seq = 58250 + + wanted_topic = sys.argv[1] + msg_params = {} + if 'points' in wanted_topic: + msg_type = msg_PointCloud2 + elif ('imu' in wanted_topic) or ('gyro' in wanted_topic) or ('accel' in wanted_topic): + msg_type = msg_Imu + elif 'theora' in wanted_topic: + try: + msg_type = msg_theora + except NameError as e: + print ('theora_image_transport is not installed. \nType "sudo apt-get install ros-kinetic-theora-image-transport" to enable registering on messages of type theora.') + raise + elif 'compressed' in wanted_topic: + msg_type = msg_CompressedImage + else: + msg_type = msg_Image + + for idx in range(2, len(sys.argv)): + if sys.argv[idx] == '-s': + msg_params['seq'] = int(sys.argv[idx + 1]) + if sys.argv[idx] == '--time': + msg_params['time'] = dict(zip(['secs', 'nsecs'], [int(part) for part in sys.argv[idx + 1].split('.')])) + if sys.argv[idx] == '--timeout': + msg_params['timeout_secs'] = int(sys.argv[idx + 1]) + if sys.argv[idx] == '--filename': + msg_params['filename'] = sys.argv[idx+1] + + msg_retriever = CWaitForMessage(msg_params) + if '/' in wanted_topic: + msg_params.setdefault('topic', wanted_topic) + res = msg_retriever.wait_for_message(msg_params, msg_type) + rospy.loginfo('Got message: %s' % res.header) + if (hasattr(res, 'encoding')): + print ('res.encoding:', res.encoding) + if (hasattr(res, 'format')): + print ('res.format:', res.format) + else: + themes = [wanted_topic] + res = msg_retriever.wait_for_messages(themes) + print (res) + + +if __name__ == '__main__': + main() + diff --git a/camera_ws/src/realsense-ros/realsense2_description/launch/view_d415_model.launch b/camera_ws/src/realsense-ros/realsense2_description/launch/view_d415_model.launch new file mode 100644 index 0000000000000000000000000000000000000000..670994a780770057940669e5d292b6fdb4288725 --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/launch/view_d415_model.launch @@ -0,0 +1,8 @@ + + + + + + + + diff --git a/camera_ws/src/realsense-ros/realsense2_description/launch/view_d435_model.launch b/camera_ws/src/realsense-ros/realsense2_description/launch/view_d435_model.launch new file mode 100644 index 0000000000000000000000000000000000000000..14f18670727a89eb40705551492057c4f1f5e57c --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/launch/view_d435_model.launch @@ -0,0 +1,8 @@ + + + + + + + + diff --git a/camera_ws/src/realsense-ros/realsense2_description/launch/view_d435i_model.launch b/camera_ws/src/realsense-ros/realsense2_description/launch/view_d435i_model.launch new file mode 100644 index 0000000000000000000000000000000000000000..8c8411af91edab4d820a220a53a0f89a68f8728a --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/launch/view_d435i_model.launch @@ -0,0 +1,8 @@ + + + + + + + + diff --git a/camera_ws/src/realsense-ros/realsense2_description/launch/view_d455_model.launch b/camera_ws/src/realsense-ros/realsense2_description/launch/view_d455_model.launch new file mode 100644 index 0000000000000000000000000000000000000000..1f6a16121414e3b9c5fc3d02d72b529a7af8790e --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/launch/view_d455_model.launch @@ -0,0 +1,8 @@ + + + + + + + + diff --git a/camera_ws/src/realsense-ros/realsense2_description/launch/view_l515_model.launch b/camera_ws/src/realsense-ros/realsense2_description/launch/view_l515_model.launch new file mode 100644 index 0000000000000000000000000000000000000000..1ba9c6d498203abbffc6c6cbec8686bf913b0f5d --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/launch/view_l515_model.launch @@ -0,0 +1,8 @@ + + + + + + + + diff --git a/camera_ws/src/realsense-ros/realsense2_description/launch/view_r410_model.launch b/camera_ws/src/realsense-ros/realsense2_description/launch/view_r410_model.launch new file mode 100644 index 0000000000000000000000000000000000000000..923569b651a4f629d3f6c44d37c9294c67b59d82 --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/launch/view_r410_model.launch @@ -0,0 +1,8 @@ + + + + + + + + diff --git a/camera_ws/src/realsense-ros/realsense2_description/package.xml b/camera_ws/src/realsense-ros/realsense2_description/package.xml new file mode 100644 index 0000000000000000000000000000000000000000..96709be019060df2a4fee12d1da1a59a4f7fe2cf --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/package.xml @@ -0,0 +1,17 @@ + + + realsense2_description + 2.3.2 + RealSense Camera description package for Intel 3D D400 cameras + Doron Hirshberg + Apache 2.0 + + http://www.ros.org/wiki/RealSense + https://github.com/intel-ros/realsense/issues + + Sergey Dorodnicov + Doron Hirshberg + catkin + xacro + rosunit + diff --git a/camera_ws/src/realsense-ros/realsense2_description/rviz/urdf.rviz b/camera_ws/src/realsense-ros/realsense2_description/rviz/urdf.rviz new file mode 100644 index 0000000000000000000000000000000000000000..18372e149d43f4980d5775da22d304a0a56c40a2 --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/rviz/urdf.rviz @@ -0,0 +1,227 @@ +Panels: + - Class: rviz/Displays + Help Height: 78 + Name: Displays + Property Tree Widget: + Expanded: + - /Global Options1 + - /Status1 + - /RobotModel1 + - /TF1 + - /TF1/Frames1 + Splitter Ratio: 0.636029422 + Tree Height: 595 + - Class: rviz/Selection + Name: Selection + - Class: rviz/Tool Properties + Expanded: + - /2D Pose Estimate1 + - /2D Nav Goal1 + - /Publish Point1 + Name: Tool Properties + Splitter Ratio: 0.588679016 + - Class: rviz/Views + Expanded: + - /Current View1 + Name: Views + Splitter Ratio: 0.5 + - Class: rviz/Time + Experimental: false + Name: Time + SyncMode: 0 + SyncSource: PointCloud2 +Toolbars: + toolButtonStyle: 2 +Visualization Manager: + Class: "" + Displays: + - Alpha: 0.5 + Cell Size: 1 + Class: rviz/Grid + Color: 160; 160; 164 + Enabled: true + Line Style: + Line Width: 0.0299999993 + Value: Lines + Name: Grid + Normal Cell Count: 0 + Offset: + X: 0 + Y: 0 + Z: 0 + Plane: XY + Plane Cell Count: 10 + Reference Frame: + Value: true + - Alpha: 1 + Class: rviz/RobotModel + Collision Enabled: false + Enabled: true + Links: + All Links Enabled: true + Expand Joint Details: false + Expand Link Details: false + Expand Tree: false + Link Tree Style: Links in Alphabetic Order + base_link: + Alpha: 1 + Show Axes: false + Show Trail: false + camera_bottom_screw_frame: + Alpha: 1 + Show Axes: false + Show Trail: false + camera_link: + Alpha: 1 + Show Axes: false + Show Trail: false + Value: true + Name: RobotModel + Robot Description: robot_description + TF Prefix: "" + Update Interval: 0 + Value: true + Visual Enabled: true + - Class: rviz/TF + Enabled: true + Frame Timeout: 15 + Frames: + All Enabled: false + base_link: + Value: true + camera_bottom_screw_frame: + Value: false + camera_color_frame: + Value: false + camera_color_optical_frame: + Value: false + camera_depth_frame: + Value: false + camera_depth_optical_frame: + Value: false + camera_infra1_frame: + Value: true + camera_infra1_optical_frame: + Value: true + camera_infra2_frame: + Value: true + camera_infra2_optical_frame: + Value: true + camera_link: + Value: true + Marker Scale: 0.100000001 + Name: TF + Show Arrows: true + Show Axes: true + Show Names: true + Tree: + base_link: + camera_bottom_screw_frame: + camera_link: + camera_color_frame: + camera_color_optical_frame: + {} + camera_depth_frame: + camera_depth_optical_frame: + {} + camera_infra1_frame: + camera_infra1_optical_frame: + {} + camera_infra2_frame: + camera_infra2_optical_frame: + {} + Update Interval: 0 + Value: true + - Alpha: 1 + Autocompute Intensity Bounds: true + Autocompute Value Bounds: + Max Value: 10 + Min Value: -10 + Value: true + Axis: Z + Channel Name: intensity + Class: rviz/PointCloud2 + Color: 255; 255; 255 + Color Transformer: RGB8 + Decay Time: 0 + Enabled: true + Invert Rainbow: false + Max Color: 255; 255; 255 + Max Intensity: 4096 + Min Color: 0; 0; 0 + Min Intensity: 0 + Name: PointCloud2 + Position Transformer: XYZ + Queue Size: 10 + Selectable: true + Size (Pixels): 3 + Size (m): 0.00999999978 + Style: Flat Squares + Topic: /camera/depth/color/points + Unreliable: false + Use Fixed Frame: true + Use rainbow: true + Value: true + Enabled: true + Global Options: + Background Color: 48; 48; 48 + Default Light: true + Fixed Frame: base_link + Frame Rate: 30 + Name: root + Tools: + - Class: rviz/Interact + Hide Inactive Objects: true + - Class: rviz/MoveCamera + - Class: rviz/Select + - Class: rviz/FocusCamera + - Class: rviz/Measure + - Class: rviz/SetInitialPose + Topic: /initialpose + - Class: rviz/SetGoal + Topic: /move_base_simple/goal + - Class: rviz/PublishPoint + Single click: true + Topic: /clicked_point + Value: true + Views: + Current: + Class: rviz/ThirdPersonFollower + Distance: 0.409323573 + Enable Stereo Rendering: + Stereo Eye Separation: 0.0599999987 + Stereo Focal Distance: 1 + Swap Stereo Eyes: false + Value: false + Focal Point: + X: 0 + Y: 0 + Z: 0 + Focal Shape Fixed Size: true + Focal Shape Size: 0.0500000007 + Invert Z Axis: false + Name: Current View + Near Clip Distance: 0.00999999978 + Pitch: -0.00499998499 + Target Frame: + Value: ThirdPersonFollower (rviz) + Yaw: 0.685000181 + Saved: ~ +Window Geometry: + Displays: + collapsed: false + Height: 876 + Hide Left Dock: false + Hide Right Dock: false + QMainWindow State: 000000ff00000000fd000000040000000000000232000002e2fc0200000008fb0000001200530065006c0065006300740069006f006e00000001e10000009b0000006100fffffffb0000001e0054006f006f006c002000500072006f007000650072007400690065007302000001ed000001df00000185000000a3fb000000120056006900650077007300200054006f006f02000001df000002110000018500000122fb000000200054006f006f006c002000500072006f0070006500720074006900650073003203000002880000011d000002210000017afb000000100044006900730070006c0061007900730100000028000002e2000000d700fffffffb0000002000730065006c0065006300740069006f006e00200062007500660066006500720200000138000000aa0000023a00000294fb00000014005700690064006500530074006500720065006f02000000e6000000d2000003ee0000030bfb0000000c004b0069006e0065006300740200000186000001060000030c00000261000000010000010f000002e2fc0200000003fb0000001e0054006f006f006c002000500072006f00700065007200740069006500730100000041000000780000000000000000fb0000000a005600690065007700730100000028000002e2000000ad00fffffffb0000001200530065006c0065006300740069006f006e010000025a000000b200000000000000000000000200000490000000a9fc0100000001fb0000000a00560069006500770073030000004e00000080000002e100000197000000030000055f0000003efc0100000002fb0000000800540069006d006501000000000000055f0000030000fffffffb0000000800540069006d0065010000000000000450000000000000000000000212000002e200000004000000040000000800000008fc0000000100000002000000010000000a0054006f006f006c00730100000000ffffffff0000000000000000 + Selection: + collapsed: false + Time: + collapsed: false + Tool Properties: + collapsed: false + Views: + collapsed: false + Width: 1375 + X: 65 + Y: 24 diff --git a/camera_ws/src/realsense-ros/realsense2_description/tests/dual_d415.xacro b/camera_ws/src/realsense-ros/realsense2_description/tests/dual_d415.xacro new file mode 100644 index 0000000000000000000000000000000000000000..a2c5870fb2e66d0cd7ea0f01942891f2faa9fdbc --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/tests/dual_d415.xacro @@ -0,0 +1,13 @@ + + + + + + + + + + + + + diff --git a/camera_ws/src/realsense-ros/realsense2_description/tests/dual_r410.xacro b/camera_ws/src/realsense-ros/realsense2_description/tests/dual_r410.xacro new file mode 100644 index 0000000000000000000000000000000000000000..686c52a5e3031ff359f22e5d73baa89b8d75d5f8 --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/tests/dual_r410.xacro @@ -0,0 +1,13 @@ + + + + + + + + + + + + + diff --git a/camera_ws/src/realsense-ros/realsense2_description/urdf/_d415.urdf.xacro b/camera_ws/src/realsense-ros/realsense2_description/urdf/_d415.urdf.xacro new file mode 100644 index 0000000000000000000000000000000000000000..07306bfa0be9f383718f1ea1a875f164df97813c --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/urdf/_d415.urdf.xacro @@ -0,0 +1,152 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/camera_ws/src/realsense-ros/realsense2_description/urdf/_d435.urdf.xacro b/camera_ws/src/realsense-ros/realsense2_description/urdf/_d435.urdf.xacro new file mode 100644 index 0000000000000000000000000000000000000000..5593cbc33461f15d3ab54c03604356b512620e3c --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/urdf/_d435.urdf.xacro @@ -0,0 +1,159 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/camera_ws/src/realsense-ros/realsense2_description/urdf/_d435i.urdf.xacro b/camera_ws/src/realsense-ros/realsense2_description/urdf/_d435i.urdf.xacro new file mode 100644 index 0000000000000000000000000000000000000000..b62aef2b83ff15b6594b83908b17946fe0f7873f --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/urdf/_d435i.urdf.xacro @@ -0,0 +1,21 @@ + + + + + + + + + + + + + + + diff --git a/camera_ws/src/realsense-ros/realsense2_description/urdf/test_d435_multiple_cameras.urdf.xacro b/camera_ws/src/realsense-ros/realsense2_description/urdf/test_d435_multiple_cameras.urdf.xacro new file mode 100644 index 0000000000000000000000000000000000000000..5ee8586dfdc9b8a16898ef5bce9d342eb4f84e16 --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/urdf/test_d435_multiple_cameras.urdf.xacro @@ -0,0 +1,14 @@ + + + + + + + + + + + + + + diff --git a/camera_ws/src/realsense-ros/realsense2_description/urdf/test_l515_camera.urdf.xacro b/camera_ws/src/realsense-ros/realsense2_description/urdf/test_l515_camera.urdf.xacro new file mode 100644 index 0000000000000000000000000000000000000000..6b1c3354a073412cfcbe57533672f4f1ab865b46 --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/urdf/test_l515_camera.urdf.xacro @@ -0,0 +1,12 @@ + + + + + + + + + + + + diff --git a/camera_ws/src/realsense-ros/realsense2_description/urdf/test_r430_camera.urdf.xacro b/camera_ws/src/realsense-ros/realsense2_description/urdf/test_r430_camera.urdf.xacro new file mode 100644 index 0000000000000000000000000000000000000000..9064208e0910a7d70bdf0f7d3dcb0c1f668d8d3c --- /dev/null +++ b/camera_ws/src/realsense-ros/realsense2_description/urdf/test_r430_camera.urdf.xacro @@ -0,0 +1,10 @@ + + + + + + + + + +