Add files using upload-large-folder tool

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 .

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()

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

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

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

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
+        ])))

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

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
+    

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)

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)

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)

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)

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": {}
+            }
+        }
+    }
+}

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": []
+    }
+}

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": {}
+            }
+        }
+    }
+}

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": []
+    }
+}

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": []
+    }
+}

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": []
+    }
+}

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": []
+    }
+}

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": []
+    }
+}

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": []
+    }
+}

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])
+    )

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))

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

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)

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

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

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

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
+

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)