release

agent/dataset/buffer.py

@@ -0,0 +1,110 @@
+"""
+Pre-training data loader. Modified from https://github.com/jannerm/diffuser/blob/main/diffuser/datasets/buffer.py
+
+"""
+
+import numpy as np
+import torch
+
+
+def atleast_2d(x):
+    if isinstance(x, torch.Tensor):
+        while x.dim() < 2:
+            x = x.unsqueeze(-1)
+        return x
+    else:
+        while x.ndim < 2:
+            x = np.expand_dims(x, axis=-1)
+    return x
+
+
+class StitchedBuffer:
+
+    def __init__(
+        self,
+        sum_of_path_lengths,
+        device="cpu",
+    ):
+        self.sum_of_path_lengths = sum_of_path_lengths
+        if device == "cpu":
+            self._dict = {
+                "path_lengths": np.zeros(sum_of_path_lengths, dtype=int),
+            }
+        else:
+            self._dict = {
+                "path_lengths": torch.zeros(sum_of_path_lengths, dtype=int).to(device),
+            }
+        self._count = 0
+        self.sum_of_path_lengths = sum_of_path_lengths
+        self.device = device
+
+    def __repr__(self):
+        return "Fields:\n" + "\n".join(
+            f"    {key}: {val.shape}" for key, val in self.items()
+        )
+
+    def __getitem__(self, key):
+        return self._dict[key]
+
+    def __setitem__(self, key, val):
+        self._dict[key] = val
+        self._add_attributes()
+
+    @property
+    def n_episodes(self):
+        return self._count
+
+    @property
+    def n_steps(self):
+        return sum(self["path_lengths"])
+
+    def _add_keys(self, path):
+        if hasattr(self, "keys"):
+            return
+        self.keys = list(path.keys())
+
+    def _add_attributes(self):
+        """
+        can access fields with `buffer.observations`
+        instead of `buffer['observations']`
+        """
+        for key, val in self._dict.items():
+            setattr(self, key, val)
+
+    def items(self):
+        return {k: v for k, v in self._dict.items() if k != "path_lengths"}.items()
+
+    def _allocate(self, key, array):
+        assert key not in self._dict
+        dim = array.shape[1:]  # skip batch dimension
+        shape = (self.sum_of_path_lengths, *dim)
+        if self.device == "cpu":
+            self._dict[key] = np.zeros(shape, dtype=np.float32)
+        else:
+            self._dict[key] = torch.zeros(shape, dtype=torch.float32).to(self.device)
+        # print(f'[ utils/mujoco ] Allocated {key} with size {shape}')
+
+    def add_path(self, path):
+        path_length = len(path["observations"])
+        # assert path_length <= self.sum_of_path_lengths
+
+        ## if first path added, set keys based on contents
+        self._add_keys(path)
+
+        ## add tracked keys in path
+        for key in self.keys:
+            array = atleast_2d(path[key])
+            if key not in self._dict:
+                self._allocate(key, array)
+            self._dict[key][self._count : self._count + path_length] = array
+
+        ## record path length
+        self._dict["path_lengths"][
+            self._count : self._count + path_length
+        ] = path_length
+
+        ## increment path counter
+        self._count += path_length
+
+    def finalize(self):
+        self._add_attributes()

agent/dataset/d3il_dataset/__init__.py

@@ -0,0 +1,18 @@
+import os
+
+FRAMEWORK_DIR = os.path.dirname(__file__)
+
+
+def sim_framework_path(*args) -> str:
+    """
+    Abstraction from os.path.join()
+    Builds absolute paths from relative path strings with SIM_FRAMEWORK/ as root.
+    If args already contains an absolute path, it is used as root for the subsequent joins
+    Args:
+        *args:
+
+    Returns:
+        absolute path
+
+    """
+    return os.path.abspath(os.path.join(FRAMEWORK_DIR, *args))

agent/dataset/d3il_dataset/aligning_dataset.py

@@ -0,0 +1,342 @@
+import random
+from typing import Optional, Callable, Any
+import logging
+
+import os
+import glob
+
+try:
+    import cv2  # not included in pyproject.toml
+except:
+    print("Installing cv2")
+    os.system("pip install opencv-python")
+import torch
+import pickle
+import numpy as np
+from tqdm import tqdm
+
+from agent.dataset.d3il_dataset.base_dataset import TrajectoryDataset
+from agent.dataset.d3il_dataset import sim_framework_path
+
+
+class Aligning_Dataset(TrajectoryDataset):
+    def __init__(
+        self,
+        data_directory: os.PathLike,
+        device="cpu",
+        obs_dim: int = 20,
+        action_dim: int = 2,
+        max_len_data: int = 256,
+        window_size: int = 1,
+    ):
+
+        super().__init__(
+            data_directory=data_directory,
+            device=device,
+            obs_dim=obs_dim,
+            action_dim=action_dim,
+            max_len_data=max_len_data,
+            window_size=window_size,
+        )
+
+        logging.info("Loading Robot Push Dataset")
+
+        inputs = []
+        actions = []
+        masks = []
+
+        # data_dir = sim_framework_path(data_directory)
+        # state_files = glob.glob(data_dir + "/env*")
+
+        rp_data_dir = sim_framework_path("data/aligning/all_data/state")
+        state_files = np.load(sim_framework_path(data_directory), allow_pickle=True)
+
+        for file in state_files:
+
+            with open(os.path.join(rp_data_dir, file), "rb") as f:
+                env_state = pickle.load(f)
+
+            # lengths.append(len(env_state['robot']['des_c_pos']))
+
+            zero_obs = np.zeros((1, self.max_len_data, self.obs_dim), dtype=np.float32)
+            zero_action = np.zeros(
+                (1, self.max_len_data, self.action_dim), dtype=np.float32
+            )
+            zero_mask = np.zeros((1, self.max_len_data), dtype=np.float32)
+
+            # robot and box positions
+            robot_des_pos = env_state["robot"]["des_c_pos"]
+
+            robot_c_pos = env_state["robot"]["c_pos"]
+
+            push_box_pos = env_state["push-box"]["pos"]
+            push_box_quat = env_state["push-box"]["quat"]
+
+            target_box_pos = env_state["target-box"]["pos"]
+            target_box_quat = env_state["target-box"]["quat"]
+
+            # target_box_pos = np.zeros(push_box_pos.shape)
+            # target_box_quat = np.zeros(push_box_quat.shape)
+            # target_box_pos[:] = push_box_pos[-1:]
+            # target_box_quat[:] = push_box_quat[-1:]
+
+            input_state = np.concatenate(
+                (
+                    robot_des_pos,
+                    robot_c_pos,
+                    push_box_pos,
+                    push_box_quat,
+                    target_box_pos,
+                    target_box_quat,
+                ),
+                axis=-1,
+            )
+
+            vel_state = robot_des_pos[1:] - robot_des_pos[:-1]
+
+            valid_len = len(input_state) - 1
+
+            zero_obs[0, :valid_len, :] = input_state[:-1]
+            zero_action[0, :valid_len, :] = vel_state
+            zero_mask[0, :valid_len] = 1
+
+            inputs.append(zero_obs)
+            actions.append(zero_action)
+            masks.append(zero_mask)
+
+        # shape: B, T, n
+        self.observations = torch.from_numpy(np.concatenate(inputs)).to(device).float()
+        self.actions = torch.from_numpy(np.concatenate(actions)).to(device).float()
+        self.masks = torch.from_numpy(np.concatenate(masks)).to(device).float()
+
+        self.num_data = len(self.observations)
+
+        self.slices = self.get_slices()
+
+    def get_slices(self):
+        slices = []
+
+        min_seq_length = np.inf
+        for i in range(self.num_data):
+            T = self.get_seq_length(i)
+            min_seq_length = min(T, min_seq_length)
+
+            if T - self.window_size < 0:
+                print(
+                    f"Ignored short sequence #{i}: len={T}, window={self.window_size}"
+                )
+            else:
+                slices += [
+                    (i, start, start + self.window_size)
+                    for start in range(T - self.window_size + 1)
+                ]  # slice indices follow convention [start, end)
+
+        return slices
+
+    def get_seq_length(self, idx):
+        return int(self.masks[idx].sum().item())
+
+    def get_all_actions(self):
+        result = []
+        # mask out invalid actions
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.actions[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def get_all_observations(self):
+        result = []
+        # mask out invalid observations
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.observations[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def __len__(self):
+        return len(self.slices)
+
+    def __getitem__(self, idx):
+
+        i, start, end = self.slices[idx]
+
+        obs = self.observations[i, start:end]
+        act = self.actions[i, start:end]
+        mask = self.masks[i, start:end]
+
+        return obs, act, mask
+
+
+class Aligning_Img_Dataset(TrajectoryDataset):
+    def __init__(
+        self,
+        data_directory: os.PathLike,
+        device="cpu",
+        obs_dim: int = 20,
+        action_dim: int = 2,
+        max_len_data: int = 256,
+        window_size: int = 1,
+    ):
+
+        super().__init__(
+            data_directory=data_directory,
+            device=device,
+            obs_dim=obs_dim,
+            action_dim=action_dim,
+            max_len_data=max_len_data,
+            window_size=window_size,
+        )
+
+        logging.info("Loading Robot Push Dataset")
+
+        inputs = []
+        actions = []
+        masks = []
+
+        data_dir = sim_framework_path("environments/dataset/data/aligning/all_data")
+
+        state_files = np.load(sim_framework_path(data_directory), allow_pickle=True)
+
+        bp_cam_imgs = []
+        inhand_cam_imgs = []
+
+        for file in tqdm(state_files[:3]):
+
+            with open(os.path.join(data_dir, "state", file), "rb") as f:
+                env_state = pickle.load(f)
+
+            # lengths.append(len(env_state['robot']['des_c_pos']))
+
+            zero_obs = np.zeros((1, self.max_len_data, self.obs_dim), dtype=np.float32)
+            zero_action = np.zeros(
+                (1, self.max_len_data, self.action_dim), dtype=np.float32
+            )
+            zero_mask = np.zeros((1, self.max_len_data), dtype=np.float32)
+
+            # robot and box positions
+            robot_des_pos = env_state["robot"]["des_c_pos"]
+            robot_c_pos = env_state["robot"]["c_pos"]
+
+            file_name = os.path.basename(file).split(".")[0]
+
+            ###############################################################
+            bp_images = []
+            bp_imgs = glob.glob(data_dir + "/images/bp-cam/" + file_name + "/*")
+            bp_imgs.sort(key=lambda x: int(os.path.basename(x).split(".")[0]))
+
+            for img in bp_imgs:
+                image = cv2.imread(img).astype(np.float32)
+                image = image.transpose((2, 0, 1)) / 255.0
+
+                image = torch.from_numpy(image).to(self.device).float().unsqueeze(0)
+
+                bp_images.append(image)
+
+            bp_images = torch.concatenate(bp_images, dim=0)
+            ################################################################
+            inhand_imgs = glob.glob(data_dir + "/images/inhand-cam/" + file_name + "/*")
+            inhand_imgs.sort(key=lambda x: int(os.path.basename(x).split(".")[0]))
+            inhand_images = []
+            for img in inhand_imgs:
+                image = cv2.imread(img).astype(np.float32)
+                image = image.transpose((2, 0, 1)) / 255.0
+
+                image = torch.from_numpy(image).to(self.device).float().unsqueeze(0)
+
+                inhand_images.append(image)
+            inhand_images = torch.concatenate(inhand_images, dim=0)
+            ##################################################################
+
+            # push_box_pos = env_state['push-box']['pos']
+            # push_box_quat = env_state['push-box']['quat']
+            #
+            # target_box_pos = env_state['target-box']['pos']
+            # target_box_quat = env_state['target-box']['quat']
+
+            # target_box_pos = np.zeros(push_box_pos.shape)
+            # target_box_quat = np.zeros(push_box_quat.shape)
+            # target_box_pos[:] = push_box_pos[-1:]
+            # target_box_quat[:] = push_box_quat[-1:]
+
+            # input_state = np.concatenate((robot_des_pos), axis=-1)
+
+            vel_state = robot_des_pos[1:] - robot_des_pos[:-1]
+
+            valid_len = len(vel_state)
+
+            zero_obs[0, :valid_len, :] = robot_des_pos[:-1]
+            zero_action[0, :valid_len, :] = vel_state
+            zero_mask[0, :valid_len] = 1
+
+            bp_cam_imgs.append(bp_images)
+            inhand_cam_imgs.append(inhand_images)
+
+            inputs.append(zero_obs)
+            actions.append(zero_action)
+            masks.append(zero_mask)
+
+        self.bp_cam_imgs = bp_cam_imgs
+        self.inhand_cam_imgs = inhand_cam_imgs
+
+        # shape: B, T, n
+        self.observations = torch.from_numpy(np.concatenate(inputs)).to(device).float()
+        self.actions = torch.from_numpy(np.concatenate(actions)).to(device).float()
+        self.masks = torch.from_numpy(np.concatenate(masks)).to(device).float()
+
+        self.num_data = len(self.observations)
+
+        self.slices = self.get_slices()
+
+    def get_slices(self):
+        slices = []
+
+        min_seq_length = np.inf
+        for i in range(self.num_data):
+            T = self.get_seq_length(i)
+            min_seq_length = min(T, min_seq_length)
+
+            if T - self.window_size < 0:
+                print(
+                    f"Ignored short sequence #{i}: len={T}, window={self.window_size}"
+                )
+            else:
+                slices += [
+                    (i, start, start + self.window_size)
+                    for start in range(T - self.window_size + 1)
+                ]  # slice indices follow convention [start, end)
+
+        return slices
+
+    def get_seq_length(self, idx):
+        return int(self.masks[idx].sum().item())
+
+    def get_all_actions(self):
+        result = []
+        # mask out invalid actions
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.actions[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def get_all_observations(self):
+        result = []
+        # mask out invalid observations
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.observations[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def __len__(self):
+        return len(self.slices)
+
+    def __getitem__(self, idx):
+
+        i, start, end = self.slices[idx]
+
+        obs = self.observations[i, start:end]
+        act = self.actions[i, start:end]
+        mask = self.masks[i, start:end]
+
+        bp_imgs = self.bp_cam_imgs[i][start:end]
+        inhand_imgs = self.inhand_cam_imgs[i][start:end]
+
+        return bp_imgs, inhand_imgs, obs, act, mask

agent/dataset/d3il_dataset/avoiding_dataset.py

@@ -0,0 +1,126 @@
+import logging
+
+import os
+import torch
+import pickle
+import numpy as np
+
+from agent.dataset.d3il_dataset.base_dataset import TrajectoryDataset
+
+
+class Avoiding_Dataset(TrajectoryDataset):
+    def __init__(
+        self,
+        data_directory: os.PathLike,
+        device="cpu",
+        obs_dim: int = 20,
+        action_dim: int = 2,
+        max_len_data: int = 256,
+        window_size: int = 1,
+    ):
+
+        super().__init__(
+            data_directory=data_directory,
+            device=device,
+            obs_dim=obs_dim,
+            action_dim=action_dim,
+            max_len_data=max_len_data,
+            window_size=window_size,
+        )
+
+        logging.info("Loading Sorting Dataset")
+
+        inputs = []
+        actions = []
+        masks = []
+
+        data_dir = data_directory
+        state_files = os.listdir(data_dir)
+
+        for file in state_files:
+            with open(os.path.join(data_dir, file), "rb") as f:
+                env_state = pickle.load(f)
+
+            zero_obs = np.zeros((1, self.max_len_data, self.obs_dim), dtype=np.float32)
+            zero_action = np.zeros(
+                (1, self.max_len_data, self.action_dim), dtype=np.float32
+            )
+            zero_mask = np.zeros((1, self.max_len_data), dtype=np.float32)
+
+            # robot and box posistion
+            robot_des_pos = env_state["robot"]["des_c_pos"][:, :2]
+            robot_c_pos = env_state["robot"]["c_pos"][:, :2]
+
+            input_state = np.concatenate((robot_des_pos, robot_c_pos), axis=-1)
+
+            vel_state = robot_des_pos[1:] - robot_des_pos[:-1]
+            valid_len = len(vel_state)
+
+            zero_obs[0, :valid_len, :] = input_state[:-1]
+            zero_action[0, :valid_len, :] = vel_state
+            zero_mask[0, :valid_len] = 1
+
+            inputs.append(zero_obs)
+            actions.append(zero_action)
+            masks.append(zero_mask)
+
+        # shape: B, T, n
+        self.observations = torch.from_numpy(np.concatenate(inputs)).to(device).float()
+        self.actions = torch.from_numpy(np.concatenate(actions)).to(device).float()
+        self.masks = torch.from_numpy(np.concatenate(masks)).to(device).float()
+
+        self.num_data = len(self.observations)
+
+        self.slices = self.get_slices()
+
+    def get_slices(self):
+        slices = []
+
+        min_seq_length = np.inf
+        for i in range(self.num_data):
+            T = self.get_seq_length(i)
+            min_seq_length = min(T, min_seq_length)
+
+            if T - self.window_size < 0:
+                print(
+                    f"Ignored short sequence #{i}: len={T}, window={self.window_size}"
+                )
+            else:
+                slices += [
+                    (i, start, start + self.window_size)
+                    for start in range(T - self.window_size + 1)
+                ]  # slice indices follow convention [start, end)
+
+        return slices
+
+    def get_seq_length(self, idx):
+        return int(self.masks[idx].sum().item())
+
+    def get_all_actions(self):
+        result = []
+        # mask out invalid actions
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.actions[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def get_all_observations(self):
+        result = []
+        # mask out invalid observations
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.observations[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def __len__(self):
+        return len(self.slices)
+
+    def __getitem__(self, idx):
+
+        i, start, end = self.slices[idx]
+
+        obs = self.observations[i, start:end]
+        act = self.actions[i, start:end]
+        mask = self.masks[i, start:end]
+
+        return obs, act, mask

agent/dataset/d3il_dataset/base_dataset.py

@@ -0,0 +1,54 @@
+import abc
+import os
+
+from torch.utils.data import Dataset
+
+
+class TrajectoryDataset(Dataset, abc.ABC):
+    """
+    A dataset containing trajectories.
+    TrajectoryDataset[i] returns: (observations, actions, mask)
+        observations: Tensor[T, ...], T frames of observations
+        actions: Tensor[T, ...], T frames of actions
+        mask: Tensor[T]: 0: invalid; 1: valid
+    """
+
+    def __init__(
+        self,
+        data_directory: os.PathLike,
+        device="cpu",
+        obs_dim: int = 20,
+        action_dim: int = 2,
+        max_len_data: int = 256,
+        window_size: int = 1,
+    ):
+
+        self.data_directory = data_directory
+        self.device = device
+
+        self.max_len_data = max_len_data
+        self.action_dim = action_dim
+        self.obs_dim = obs_dim
+
+        self.window_size = window_size
+
+    @abc.abstractmethod
+    def get_seq_length(self, idx):
+        """
+        Returns the length of the idx-th trajectory.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_all_actions(self):
+        """
+        Returns all actions from all trajectories, concatenated on dim 0 (time).
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_all_observations(self):
+        """
+        Returns all actions from all trajectories, concatenated on dim 0 (time).
+        """
+        raise NotImplementedError

agent/dataset/d3il_dataset/geo_transform.py

@@ -0,0 +1,350 @@
+import itertools
+
+import numpy as np
+
+"""
+From OpenAIGym Please see there under mujoco/Robots
+"""
+
+# For testing whether a number is close to zero
+_FLOAT_EPS = np.finfo(np.float64).eps
+_EPS4 = _FLOAT_EPS * 4.0
+
+
+def get_quaternion_error(curr_quat, des_quat):
+    """
+    Calculates the difference between the current quaternion and the desired quaternion.
+    See Siciliano textbook page 140 Eq 3.91
+
+    :param curr_quat: current quaternion
+    :param des_quat: desired quaternion
+    :return: difference between current quaternion and desired quaternion
+    """
+    quatError = np.zeros((3,))
+
+    quatError[0] = (
+        curr_quat[0] * des_quat[1]
+        - des_quat[0] * curr_quat[1]
+        - curr_quat[3] * des_quat[2]
+        + curr_quat[2] * des_quat[3]
+    )
+
+    quatError[1] = (
+        curr_quat[0] * des_quat[2]
+        - des_quat[0] * curr_quat[2]
+        + curr_quat[3] * des_quat[1]
+        - curr_quat[1] * des_quat[3]
+    )
+
+    quatError[2] = (
+        curr_quat[0] * des_quat[3]
+        - des_quat[0] * curr_quat[3]
+        - curr_quat[2] * des_quat[1]
+        + curr_quat[1] * des_quat[2]
+    )
+
+    return quatError
+
+
+def euler2mat(euler):
+    """Convert Euler Angles to Rotation Matrix.  See rotation.py for notes"""
+    euler = np.asarray(euler, dtype=np.float64)
+    assert euler.shape[-1] == 3, "Invalid shaped euler {}".format(euler)
+
+    ai, aj, ak = -euler[..., 2], -euler[..., 1], -euler[..., 0]
+    si, sj, sk = np.sin(ai), np.sin(aj), np.sin(ak)
+    ci, cj, ck = np.cos(ai), np.cos(aj), np.cos(ak)
+    cc, cs = ci * ck, ci * sk
+    sc, ss = si * ck, si * sk
+
+    mat = np.empty(euler.shape[:-1] + (3, 3), dtype=np.float64)
+    mat[..., 2, 2] = cj * ck
+    mat[..., 2, 1] = sj * sc - cs
+    mat[..., 2, 0] = sj * cc + ss
+    mat[..., 1, 2] = cj * sk
+    mat[..., 1, 1] = sj * ss + cc
+    mat[..., 1, 0] = sj * cs - sc
+    mat[..., 0, 2] = -sj
+    mat[..., 0, 1] = cj * si
+    mat[..., 0, 0] = cj * ci
+    return mat
+
+
+def euler2quat(euler):
+    """Convert Euler Angles to Quaternions.  See rotation.py for notes"""
+    euler = np.asarray(euler, dtype=np.float64)
+    assert euler.shape[-1] == 3, "Invalid shape euler {}".format(euler)
+
+    ai, aj, ak = euler[..., 2] / 2, -euler[..., 1] / 2, euler[..., 0] / 2
+    si, sj, sk = np.sin(ai), np.sin(aj), np.sin(ak)
+    ci, cj, ck = np.cos(ai), np.cos(aj), np.cos(ak)
+    cc, cs = ci * ck, ci * sk
+    sc, ss = si * ck, si * sk
+
+    quat = np.empty(euler.shape[:-1] + (4,), dtype=np.float64)
+    quat[..., 0] = cj * cc + sj * ss
+    quat[..., 3] = cj * sc - sj * cs
+    quat[..., 2] = -(cj * ss + sj * cc)
+    quat[..., 1] = cj * cs - sj * sc
+    return quat
+
+
+def mat2euler(mat):
+    """Convert Rotation Matrix to Euler Angles.  See rotation.py for notes"""
+    mat = np.asarray(mat, dtype=np.float64)
+    assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat)
+
+    cy = np.sqrt(mat[..., 2, 2] * mat[..., 2, 2] + mat[..., 1, 2] * mat[..., 1, 2])
+    condition = cy > _EPS4
+    euler = np.empty(mat.shape[:-1], dtype=np.float64)
+    euler[..., 2] = np.where(
+        condition,
+        -np.arctan2(mat[..., 0, 1], mat[..., 0, 0]),
+        -np.arctan2(-mat[..., 1, 0], mat[..., 1, 1]),
+    )
+    euler[..., 1] = np.where(
+        condition, -np.arctan2(-mat[..., 0, 2], cy), -np.arctan2(-mat[..., 0, 2], cy)
+    )
+    euler[..., 0] = np.where(
+        condition, -np.arctan2(mat[..., 1, 2], mat[..., 2, 2]), 0.0
+    )
+    return euler
+
+
+def mat2quat(mat):
+    """Convert Rotation Matrix to Quaternion.  See rotation.py for notes"""
+    mat = np.asarray(mat, dtype=np.float64)
+    assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat)
+
+    Qxx, Qyx, Qzx = mat[..., 0, 0], mat[..., 0, 1], mat[..., 0, 2]
+    Qxy, Qyy, Qzy = mat[..., 1, 0], mat[..., 1, 1], mat[..., 1, 2]
+    Qxz, Qyz, Qzz = mat[..., 2, 0], mat[..., 2, 1], mat[..., 2, 2]
+    # Fill only lower half of symmetric matrix
+    K = np.zeros(mat.shape[:-2] + (4, 4), dtype=np.float64)
+    K[..., 0, 0] = Qxx - Qyy - Qzz
+    K[..., 1, 0] = Qyx + Qxy
+    K[..., 1, 1] = Qyy - Qxx - Qzz
+    K[..., 2, 0] = Qzx + Qxz
+    K[..., 2, 1] = Qzy + Qyz
+    K[..., 2, 2] = Qzz - Qxx - Qyy
+    K[..., 3, 0] = Qyz - Qzy
+    K[..., 3, 1] = Qzx - Qxz
+    K[..., 3, 2] = Qxy - Qyx
+    K[..., 3, 3] = Qxx + Qyy + Qzz
+    K /= 3.0
+    # TODO: vectorize this -- probably could be made faster
+    q = np.empty(K.shape[:-2] + (4,))
+    it = np.nditer(q[..., 0], flags=["multi_index"])
+    while not it.finished:
+        # Use Hermitian eigenvectors, values for speed
+        vals, vecs = np.linalg.eigh(K[it.multi_index])
+        # Select largest eigenvector, reorder to w,x,y,z quaternion
+        q[it.multi_index] = vecs[[3, 0, 1, 2], np.argmax(vals)]
+        # Prefer quaternion with positive w
+        # (q * -1 corresponds to same rotation as q)
+        if q[it.multi_index][0] < 0:
+            q[it.multi_index] *= -1
+        it.iternext()
+    return q
+
+
+def quat2euler(quat):
+    """Convert Quaternion to Euler Angles.  See rotation.py for notes"""
+    return mat2euler(quat2mat(quat))
+
+
+def subtract_euler(e1, e2):
+    assert e1.shape == e2.shape
+    assert e1.shape[-1] == 3
+    q1 = euler2quat(e1)
+    q2 = euler2quat(e2)
+    q_diff = quat_mul(q1, quat_conjugate(q2))
+    return quat2euler(q_diff)
+
+
+def quat2mat(quat):
+    """Convert Quaternion to Euler Angles.  See rotation.py for notes"""
+    quat = np.asarray(quat, dtype=np.float64)
+    assert quat.shape[-1] == 4, "Invalid shape quat {}".format(quat)
+
+    w, x, y, z = quat[..., 0], quat[..., 1], quat[..., 2], quat[..., 3]
+    Nq = np.sum(quat * quat, axis=-1)
+    s = 2.0 / Nq
+    X, Y, Z = x * s, y * s, z * s
+    wX, wY, wZ = w * X, w * Y, w * Z
+    xX, xY, xZ = x * X, x * Y, x * Z
+    yY, yZ, zZ = y * Y, y * Z, z * Z
+
+    mat = np.empty(quat.shape[:-1] + (3, 3), dtype=np.float64)
+    mat[..., 0, 0] = 1.0 - (yY + zZ)
+    mat[..., 0, 1] = xY - wZ
+    mat[..., 0, 2] = xZ + wY
+    mat[..., 1, 0] = xY + wZ
+    mat[..., 1, 1] = 1.0 - (xX + zZ)
+    mat[..., 1, 2] = yZ - wX
+    mat[..., 2, 0] = xZ - wY
+    mat[..., 2, 1] = yZ + wX
+    mat[..., 2, 2] = 1.0 - (xX + yY)
+    return np.where((Nq > _FLOAT_EPS)[..., np.newaxis, np.newaxis], mat, np.eye(3))
+
+
+def quat_conjugate(q):
+    inv_q = -q
+    inv_q[..., 0] *= -1
+    return inv_q
+
+
+def quat_mul(q0, q1):
+    assert q0.shape == q1.shape
+    assert q0.shape[-1] == 4
+    assert q1.shape[-1] == 4
+
+    w0 = q0[..., 0]
+    x0 = q0[..., 1]
+    y0 = q0[..., 2]
+    z0 = q0[..., 3]
+
+    w1 = q1[..., 0]
+    x1 = q1[..., 1]
+    y1 = q1[..., 2]
+    z1 = q1[..., 3]
+
+    w = w0 * w1 - x0 * x1 - y0 * y1 - z0 * z1
+    x = w0 * x1 + x0 * w1 + y0 * z1 - z0 * y1
+    y = w0 * y1 + y0 * w1 + z0 * x1 - x0 * z1
+    z = w0 * z1 + z0 * w1 + x0 * y1 - y0 * x1
+    q = np.array([w, x, y, z])
+    if q.ndim == 2:
+        q = q.swapaxes(0, 1)
+    assert q.shape == q0.shape
+    return q
+
+
+def quat_rot_vec(q, v0):
+    q_v0 = np.array([0, v0[0], v0[1], v0[2]])
+    q_v = quat_mul(q, quat_mul(q_v0, quat_conjugate(q)))
+    v = q_v[1:]
+    return v
+
+
+def quat_identity():
+    return np.array([1, 0, 0, 0])
+
+
+def quat2axisangle(quat):
+    theta = 0
+    axis = np.array([0, 0, 1])
+    sin_theta = np.linalg.norm(quat[1:])
+
+    if sin_theta > 0.0001:
+        theta = 2 * np.arcsin(sin_theta)
+        theta *= 1 if quat[0] >= 0 else -1
+        axis = quat[1:] / sin_theta
+
+    return axis, theta
+
+
+def euler2point_euler(euler):
+    _euler = euler.copy()
+    if len(_euler.shape) < 2:
+        _euler = np.expand_dims(_euler, 0)
+    assert _euler.shape[1] == 3
+    _euler_sin = np.sin(_euler)
+    _euler_cos = np.cos(_euler)
+    return np.concatenate([_euler_sin, _euler_cos], axis=-1)
+
+
+def point_euler2euler(euler):
+    _euler = euler.copy()
+    if len(_euler.shape) < 2:
+        _euler = np.expand_dims(_euler, 0)
+    assert _euler.shape[1] == 6
+    angle = np.arctan(_euler[..., :3] / _euler[..., 3:])
+    angle[_euler[..., 3:] < 0] += np.pi
+    return angle
+
+
+def quat2point_quat(quat):
+    # Should be in qw, qx, qy, qz
+    _quat = quat.copy()
+    if len(_quat.shape) < 2:
+        _quat = np.expand_dims(_quat, 0)
+    assert _quat.shape[1] == 4
+    angle = np.arccos(_quat[:, [0]]) * 2
+    xyz = _quat[:, 1:]
+    xyz[np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5] = (xyz / np.sin(angle / 2))[
+        np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5
+    ]
+    return np.concatenate([np.sin(angle), np.cos(angle), xyz], axis=-1)
+
+
+def point_quat2quat(quat):
+    _quat = quat.copy()
+    if len(_quat.shape) < 2:
+        _quat = np.expand_dims(_quat, 0)
+    assert _quat.shape[1] == 5
+    angle = np.arctan(_quat[:, [0]] / _quat[:, [1]])
+    qw = np.cos(angle / 2)
+
+    qxyz = _quat[:, 2:]
+    qxyz[np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5] = (qxyz * np.sin(angle / 2))[
+        np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5
+    ]
+    return np.concatenate([qw, qxyz], axis=-1)
+
+
+def normalize_angles(angles):
+    """Puts angles in [-pi, pi] range."""
+    angles = angles.copy()
+    if angles.size > 0:
+        angles = (angles + np.pi) % (2 * np.pi) - np.pi
+        assert -np.pi - 1e-6 <= angles.min() and angles.max() <= np.pi + 1e-6
+    return angles
+
+
+def round_to_straight_angles(angles):
+    """Returns closest angle modulo 90 degrees"""
+    angles = np.round(angles / (np.pi / 2)) * (np.pi / 2)
+    return normalize_angles(angles)
+
+
+def get_parallel_rotations():
+    mult90 = [0, np.pi / 2, -np.pi / 2, np.pi]
+    parallel_rotations = []
+    for euler in itertools.product(mult90, repeat=3):
+        canonical = mat2euler(euler2mat(euler))
+        canonical = np.round(canonical / (np.pi / 2))
+        if canonical[0] == -2:
+            canonical[0] = 2
+        if canonical[2] == -2:
+            canonical[2] = 2
+        canonical *= np.pi / 2
+        if all([(canonical != rot).any() for rot in parallel_rotations]):
+            parallel_rotations += [canonical]
+    assert len(parallel_rotations) == 24
+    return parallel_rotations
+
+
+def posRotMat2TFMat(pos, rot_mat):
+    """Converts a position and a 3x3 rotation matrix to a 4x4 transformation matrix"""
+    t_mat = np.eye(4)
+    t_mat[:3, :3] = rot_mat
+    t_mat[:3, 3] = np.array(pos)
+    return t_mat
+
+
+def mat2posQuat(mat):
+    """Converts a 4x4 rotation matrix to a position and a quaternion"""
+    pos = mat[:3, 3]
+    quat = mat2quat(mat[:3, :3])
+    return pos, quat
+
+
+def wxyz_to_xyzw(quat):
+    """Converts WXYZ Quaternions to XYZW Quaternions"""
+    return np.roll(quat, -1)
+
+
+def xyzw_to_wxyz(quat):
+    """Converts XYZW Quaternions to WXYZ Quaternions"""
+    return np.roll(quat, 1)

agent/dataset/d3il_dataset/pushing_dataset.py

@@ -0,0 +1,161 @@
+import logging
+
+import os
+import torch
+import pickle
+import numpy as np
+
+from agent.dataset.d3il_dataset.base_dataset import TrajectoryDataset
+from agent.dataset.d3il_dataset import sim_framework_path
+from .geo_transform import quat2euler
+
+
+class Pushing_Dataset(TrajectoryDataset):
+    def __init__(
+        self,
+        data_directory: os.PathLike,
+        device="cpu",
+        obs_dim: int = 20,
+        action_dim: int = 2,
+        max_len_data: int = 256,
+        window_size: int = 1,
+    ):
+
+        super().__init__(
+            data_directory=data_directory,
+            device=device,
+            obs_dim=obs_dim,
+            action_dim=action_dim,
+            max_len_data=max_len_data,
+            window_size=window_size,
+        )
+
+        logging.info("Loading Block Push Dataset")
+
+        inputs = []
+        actions = []
+        masks = []
+
+        # for root, dirs, files in os.walk(self.data_directory):
+        #
+        #     for mode_dir in dirs:
+
+        # state_files = glob.glob(os.path.join(root, mode_dir) + "/env*")
+        # data_dir = os.path.join(sim_framework_path(data_directory), "local")
+        # data_dir = sim_framework_path(data_directory)
+        # state_files = glob.glob(data_dir + "/env*")
+
+        bp_data_dir = sim_framework_path("data/pushing/all_data")
+
+        state_files = np.load(sim_framework_path(data_directory), allow_pickle=True)
+
+        for file in state_files:
+            with open(os.path.join(bp_data_dir, file), "rb") as f:
+                env_state = pickle.load(f)
+
+            # lengths.append(len(env_state['robot']['des_c_pos']))
+
+            zero_obs = np.zeros((1, self.max_len_data, self.obs_dim), dtype=np.float32)
+            zero_action = np.zeros(
+                (1, self.max_len_data, self.action_dim), dtype=np.float32
+            )
+            zero_mask = np.zeros((1, self.max_len_data), dtype=np.float32)
+
+            # robot and box positions
+            robot_des_pos = env_state["robot"]["des_c_pos"][:, :2]
+
+            robot_c_pos = env_state["robot"]["c_pos"][:, :2]
+
+            red_box_pos = env_state["red-box"]["pos"][:, :2]
+            red_box_quat = np.tan(quat2euler(env_state["red-box"]["quat"])[:, -1:])
+
+            green_box_pos = env_state["green-box"]["pos"][:, :2]
+            green_box_quat = np.tan(quat2euler(env_state["green-box"]["quat"])[:, -1:])
+
+            red_target_pos = env_state["red-target"]["pos"][:, :2]
+            green_target_pos = env_state["green-target"]["pos"][:, :2]
+
+            input_state = np.concatenate(
+                (
+                    robot_des_pos,
+                    robot_c_pos,
+                    red_box_pos,
+                    red_box_quat,
+                    green_box_pos,
+                    green_box_quat,
+                ),
+                axis=-1,
+            )
+
+            vel_state = robot_des_pos[1:] - robot_des_pos[:-1]
+
+            valid_len = len(input_state) - 1
+
+            zero_obs[0, :valid_len, :] = input_state[:-1]
+            zero_action[0, :valid_len, :] = vel_state
+            zero_mask[0, :valid_len] = 1
+
+            inputs.append(zero_obs)
+            actions.append(zero_action)
+            masks.append(zero_mask)
+
+        # shape: B, T, n
+        self.observations = torch.from_numpy(np.concatenate(inputs)).to(device).float()
+        self.actions = torch.from_numpy(np.concatenate(actions)).to(device).float()
+        self.masks = torch.from_numpy(np.concatenate(masks)).to(device).float()
+
+        self.num_data = len(self.observations)
+
+        self.slices = self.get_slices()
+
+    def get_slices(self):
+        slices = []
+
+        min_seq_length = np.inf
+        for i in range(self.num_data):
+            T = self.get_seq_length(i)
+            min_seq_length = min(T, min_seq_length)
+
+            if T - self.window_size < 0:
+                print(
+                    f"Ignored short sequence #{i}: len={T}, window={self.window_size}"
+                )
+            else:
+                slices += [
+                    (i, start, start + self.window_size)
+                    for start in range(T - self.window_size + 1)
+                ]  # slice indices follow convention [start, end)
+
+        return slices
+
+    def get_seq_length(self, idx):
+        return int(self.masks[idx].sum().item())
+
+    def get_all_actions(self):
+        result = []
+        # mask out invalid actions
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.actions[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def get_all_observations(self):
+        result = []
+        # mask out invalid observations
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.observations[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def __len__(self):
+        return len(self.slices)
+
+    def __getitem__(self, idx):
+
+        i, start, end = self.slices[idx]
+
+        obs = self.observations[i, start:end]
+        act = self.actions[i, start:end]
+        mask = self.masks[i, start:end]
+
+        return obs, act, mask

agent/dataset/d3il_dataset/sorting_dataset.py

@@ -0,0 +1,491 @@
+import logging
+
+import os
+import glob
+import torch
+import pickle
+import numpy as np
+
+try:
+    import cv2  # not included in pyproject.toml
+except:
+    print("Installing cv2")
+    os.system("pip install opencv-python")
+from tqdm import tqdm
+
+from agent.dataset.d3il_dataset.base_dataset import TrajectoryDataset
+from agent.dataset.d3il_dataset import sim_framework_path
+from .geo_transform import quat2euler
+
+
+class Sorting_Dataset(TrajectoryDataset):
+    def __init__(
+        self,
+        data_directory: os.PathLike,
+        device="cpu",
+        obs_dim: int = 20,
+        action_dim: int = 2,
+        max_len_data: int = 256,
+        window_size: int = 1,
+        num_boxes: int = 2,
+    ):
+
+        super().__init__(
+            data_directory=data_directory,
+            device=device,
+            obs_dim=obs_dim,
+            action_dim=action_dim,
+            max_len_data=max_len_data,
+            window_size=window_size,
+        )
+
+        logging.info("Loading Sorting Dataset")
+
+        inputs = []
+        actions = []
+        masks = []
+
+        # for root, dirs, files in os.walk(self.data_directory):
+        #
+        #     for mode_dir in dirs:
+
+        # state_files = glob.glob(os.path.join(root, mode_dir) + "/env*")
+        # data_dir = os.path.join(sim_framework_path(data_directory), "local")
+        # data_dir = sim_framework_path(data_directory)
+        # state_files = glob.glob(data_dir + "/env*")
+
+        # random.seed(0)
+
+        # data_dir = sim_framework_path(data_directory)
+        # state_files = os.listdir(data_dir)
+        #
+        # random.shuffle(state_files)
+        #
+        # if data == "train":
+        #     env_state_files = state_files[50:]
+        # elif data == "eval":
+        #     env_state_files = state_files[:50]
+        # else:
+        #     assert False, "wrong data type"
+
+        if num_boxes == 2:
+            data_dir = sim_framework_path("data/sorting/2_boxes/state")
+        elif num_boxes == 4:
+            data_dir = sim_framework_path("data/sorting/4_boxes/state")
+        elif num_boxes == 6:
+            data_dir = sim_framework_path("data/sorting/6_boxes/state")
+        else:
+            assert False, "check num boxes"
+
+        state_files = np.load(sim_framework_path(data_directory), allow_pickle=True)
+
+        for file in state_files:
+            with open(os.path.join(data_dir, file), "rb") as f:
+                env_state = pickle.load(f)
+
+            # lengths.append(len(env_state['robot']['des_c_pos']))
+
+            zero_obs = np.zeros((1, self.max_len_data, self.obs_dim), dtype=np.float32)
+            zero_action = np.zeros(
+                (1, self.max_len_data, self.action_dim), dtype=np.float32
+            )
+            zero_mask = np.zeros((1, self.max_len_data), dtype=np.float32)
+
+            # robot and box posistion
+            robot_des_pos = env_state["robot"]["des_c_pos"][:, :2]
+            robot_c_pos = env_state["robot"]["c_pos"][:, :2]
+
+            if num_boxes == 2:
+                red_box1_pos = env_state["red-box1"]["pos"][:, :2]
+                red_box1_quat = np.tan(
+                    quat2euler(env_state["red-box1"]["quat"])[:, -1:]
+                )
+
+                blue_box1_pos = env_state["blue-box1"]["pos"][:, :2]
+                blue_box1_quat = np.tan(
+                    quat2euler(env_state["blue-box1"]["quat"])[:, -1:]
+                )
+
+                input_state = np.concatenate(
+                    (
+                        robot_des_pos,
+                        robot_c_pos,
+                        red_box1_pos,
+                        red_box1_quat,
+                        blue_box1_pos,
+                        blue_box1_quat,
+                    ),
+                    axis=-1,
+                )
+
+            elif num_boxes == 4:
+
+                red_box1_pos = env_state["red-box1"]["pos"][:, :2]
+                red_box1_quat = np.tan(
+                    quat2euler(env_state["red-box1"]["quat"])[:, -1:]
+                )
+
+                red_box2_pos = env_state["red-box2"]["pos"][:, :2]
+                red_box2_quat = np.tan(
+                    quat2euler(env_state["red-box2"]["quat"])[:, -1:]
+                )
+
+                blue_box1_pos = env_state["blue-box1"]["pos"][:, :2]
+                blue_box1_quat = np.tan(
+                    quat2euler(env_state["blue-box1"]["quat"])[:, -1:]
+                )
+
+                blue_box2_pos = env_state["blue-box2"]["pos"][:, :2]
+                blue_box2_quat = np.tan(
+                    quat2euler(env_state["blue-box2"]["quat"])[:, -1:]
+                )
+
+                input_state = np.concatenate(
+                    (
+                        robot_des_pos,
+                        robot_c_pos,
+                        red_box1_pos,
+                        red_box1_quat,
+                        red_box2_pos,
+                        red_box2_quat,
+                        blue_box1_pos,
+                        blue_box1_quat,
+                        blue_box2_pos,
+                        blue_box2_quat,
+                    ),
+                    axis=-1,
+                )
+
+            elif num_boxes == 6:
+
+                red_box1_pos = env_state["red-box1"]["pos"][:, :2]
+                red_box1_quat = np.tan(
+                    quat2euler(env_state["red-box1"]["quat"])[:, -1:]
+                )
+
+                red_box2_pos = env_state["red-box2"]["pos"][:, :2]
+                red_box2_quat = np.tan(
+                    quat2euler(env_state["red-box2"]["quat"])[:, -1:]
+                )
+
+                red_box3_pos = env_state["red-box3"]["pos"][:, :2]
+                red_box3_quat = np.tan(
+                    quat2euler(env_state["red-box3"]["quat"])[:, -1:]
+                )
+
+                blue_box1_pos = env_state["blue-box1"]["pos"][:, :2]
+                blue_box1_quat = np.tan(
+                    quat2euler(env_state["blue-box1"]["quat"])[:, -1:]
+                )
+
+                blue_box2_pos = env_state["blue-box2"]["pos"][:, :2]
+                blue_box2_quat = np.tan(
+                    quat2euler(env_state["blue-box2"]["quat"])[:, -1:]
+                )
+
+                blue_box3_pos = env_state["blue-box3"]["pos"][:, :2]
+                blue_box3_quat = np.tan(
+                    quat2euler(env_state["blue-box3"]["quat"])[:, -1:]
+                )
+
+                input_state = np.concatenate(
+                    (
+                        robot_des_pos,
+                        robot_c_pos,
+                        red_box1_pos,
+                        red_box1_quat,
+                        red_box2_pos,
+                        red_box2_quat,
+                        red_box3_pos,
+                        red_box3_quat,
+                        blue_box1_pos,
+                        blue_box1_quat,
+                        blue_box2_pos,
+                        blue_box2_quat,
+                        blue_box3_pos,
+                        blue_box3_quat,
+                    ),
+                    axis=-1,
+                )
+
+            else:
+                assert False, "check num boxes"
+
+            vel_state = robot_des_pos[1:] - robot_des_pos[:-1]
+            valid_len = len(vel_state)
+
+            zero_obs[0, :valid_len, :] = input_state[:-1]
+            zero_action[0, :valid_len, :] = vel_state
+            zero_mask[0, :valid_len] = 1
+
+            inputs.append(zero_obs)
+            actions.append(zero_action)
+            masks.append(zero_mask)
+
+        # shape: B, T, n
+        self.observations = torch.from_numpy(np.concatenate(inputs)).to(device).float()
+        self.actions = torch.from_numpy(np.concatenate(actions)).to(device).float()
+        self.masks = torch.from_numpy(np.concatenate(masks)).to(device).float()
+
+        self.num_data = len(self.observations)
+
+        self.slices = self.get_slices()
+
+    def get_slices(self):
+        slices = []
+
+        min_seq_length = np.inf
+        for i in range(self.num_data):
+            T = self.get_seq_length(i)
+            min_seq_length = min(T, min_seq_length)
+
+            if T - self.window_size < 0:
+                print(
+                    f"Ignored short sequence #{i}: len={T}, window={self.window_size}"
+                )
+            else:
+                slices += [
+                    (i, start, start + self.window_size)
+                    for start in range(T - self.window_size + 1)
+                ]  # slice indices follow convention [start, end)
+
+        return slices
+
+    def get_seq_length(self, idx):
+        return int(self.masks[idx].sum().item())
+
+    def get_all_actions(self):
+        result = []
+        # mask out invalid actions
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.actions[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def get_all_observations(self):
+        result = []
+        # mask out invalid observations
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.observations[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def __len__(self):
+        return len(self.slices)
+
+    def __getitem__(self, idx):
+
+        i, start, end = self.slices[idx]
+
+        obs = self.observations[i, start:end]
+        act = self.actions[i, start:end]
+        mask = self.masks[i, start:end]
+
+        return obs, act, mask
+
+
+class Sorting_Img_Dataset(TrajectoryDataset):
+    def __init__(
+        self,
+        data_directory: os.PathLike,
+        device="cpu",
+        obs_dim: int = 20,
+        action_dim: int = 2,
+        max_len_data: int = 256,
+        window_size: int = 1,
+        num_boxes: int = 2,
+    ):
+
+        super().__init__(
+            data_directory=data_directory,
+            device=device,
+            obs_dim=obs_dim,
+            action_dim=action_dim,
+            max_len_data=max_len_data,
+            window_size=window_size,
+        )
+
+        logging.info("Loading Sorting Dataset")
+
+        inputs = []
+        actions = []
+        masks = []
+
+        # for root, dirs, files in os.walk(self.data_directory):
+        #
+        #     for mode_dir in dirs:
+
+        # state_files = glob.glob(os.path.join(root, mode_dir) + "/env*")
+        # data_dir = os.path.join(sim_framework_path(data_directory), "local")
+        # data_dir = sim_framework_path(data_directory)
+        # state_files = glob.glob(data_dir + "/env*")
+
+        # random.seed(0)
+        #
+        # data_dir = sim_framework_path(data_directory)
+        # state_files = glob.glob(data_dir + '/state/*')
+        #
+        # random.shuffle(state_files)
+        #
+        # if data == "train":
+        #     env_state_files = state_files[30:]
+        # elif data == "eval":
+        #     env_state_files = state_files[:30]
+        # else:
+        #     assert False, "wrong data type"
+
+        if num_boxes == 2:
+            data_dir = sim_framework_path("environments/dataset/data/sorting/2_boxes/")
+        elif num_boxes == 4:
+            data_dir = sim_framework_path("environments/dataset/data/sorting/4_boxes/")
+        elif num_boxes == 6:
+            data_dir = sim_framework_path("environments/dataset/data/sorting/6_boxes/")
+        else:
+            assert False, "check num boxes"
+
+        state_files = np.load(sim_framework_path(data_directory), allow_pickle=True)
+
+        bp_cam_imgs = []
+        inhand_cam_imgs = []
+
+        for file in tqdm(state_files[:100]):
+            with open(os.path.join(data_dir, "state", file), "rb") as f:
+                env_state = pickle.load(f)
+
+            # lengths.append(len(env_state['robot']['des_c_pos']))
+            zero_obs = np.zeros((1, self.max_len_data, self.obs_dim), dtype=np.float32)
+            zero_action = np.zeros(
+                (1, self.max_len_data, self.action_dim), dtype=np.float32
+            )
+            zero_mask = np.zeros((1, self.max_len_data), dtype=np.float32)
+
+            # robot and box posistion
+            robot_des_pos = env_state["robot"]["des_c_pos"][:, :2]
+            robot_c_pos = env_state["robot"]["c_pos"][:, :2]
+
+            file_name = os.path.basename(file).split(".")[0]
+
+            ###############################################################
+            bp_images = []
+            bp_imgs = glob.glob(data_dir + "/images/bp-cam/" + file_name + "/*")
+            bp_imgs.sort(key=lambda x: int(os.path.basename(x).split(".")[0]))
+
+            for img in bp_imgs:
+                image = cv2.imread(img).astype(np.float32)
+                image = image.transpose((2, 0, 1)) / 255.0
+
+                image = torch.from_numpy(image).to(self.device).float().unsqueeze(0)
+
+                bp_images.append(image)
+
+            bp_images = torch.concatenate(bp_images, dim=0)
+            ################################################################
+            inhand_imgs = glob.glob(data_dir + "/images/inhand-cam/" + file_name + "/*")
+            inhand_imgs.sort(key=lambda x: int(os.path.basename(x).split(".")[0]))
+            inhand_images = []
+            for img in inhand_imgs:
+                image = cv2.imread(img).astype(np.float32)
+                image = image.transpose((2, 0, 1)) / 255.0
+
+                image = torch.from_numpy(image).to(self.device).float().unsqueeze(0)
+
+                inhand_images.append(image)
+            inhand_images = torch.concatenate(inhand_images, dim=0)
+            ##################################################################
+            # input_state = np.concatenate((robot_des_pos, robot_c_pos), axis=-1)
+
+            vel_state = robot_des_pos[1:] - robot_des_pos[:-1]
+
+            valid_len = len(vel_state)
+
+            zero_obs[0, :valid_len, :] = robot_des_pos[:-1]
+            zero_action[0, :valid_len, :] = vel_state
+            zero_mask[0, :valid_len] = 1
+
+            bp_cam_imgs.append(bp_images)
+            inhand_cam_imgs.append(inhand_images)
+
+            inputs.append(zero_obs)
+            actions.append(zero_action)
+            masks.append(zero_mask)
+
+        self.bp_cam_imgs = bp_cam_imgs
+        self.inhand_cam_imgs = inhand_cam_imgs
+
+        # shape: B, T, n
+        self.observations = torch.from_numpy(np.concatenate(inputs)).to(device).float()
+        self.actions = torch.from_numpy(np.concatenate(actions)).to(device).float()
+        self.masks = torch.from_numpy(np.concatenate(masks)).to(device).float()
+
+        self.num_data = len(self.actions)
+
+        self.slices = self.get_slices()
+
+    def get_slices(self):
+        slices = []
+
+        min_seq_length = np.inf
+        for i in range(self.num_data):
+            T = self.get_seq_length(i)
+            min_seq_length = min(T, min_seq_length)
+
+            if T - self.window_size < 0:
+                print(
+                    f"Ignored short sequence #{i}: len={T}, window={self.window_size}"
+                )
+            else:
+                slices += [
+                    (i, start, start + self.window_size)
+                    for start in range(T - self.window_size + 1)
+                ]  # slice indices follow convention [start, end)
+
+        return slices
+
+    def get_seq_length(self, idx):
+        return int(self.masks[idx].sum().item())
+
+    def get_all_actions(self):
+        result = []
+        # mask out invalid actions
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.actions[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def get_all_observations(self):
+        result = []
+        # mask out invalid observations
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.observations[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def __len__(self):
+        return len(self.slices)
+
+    def __getitem__(self, idx):
+
+        i, start, end = self.slices[idx]
+
+        obs = self.observations[i, start:end]
+        act = self.actions[i, start:end]
+        mask = self.masks[i, start:end]
+
+        bp_imgs = self.bp_cam_imgs[i][start:end]
+        inhand_imgs = self.inhand_cam_imgs[i][start:end]
+
+        # bp_imgs = np.zeros((self.window_size, 3, 96, 96), dtype=np.float32)
+        # inhand_imgs = np.zeros((self.window_size, 3, 96, 96), dtype=np.float32)
+        #
+        # for num_frame, img_file in enumerate(bp_img_files):
+        #     image = cv2.imread(img_file).astype(np.float32)
+        #     bp_imgs[num_frame] = image.transpose((2, 0, 1)) / 255.
+        #
+        # for num_frame, img_file in enumerate(inhand_img_files):
+        #     image = cv2.imread(img_file).astype(np.float32)
+        #     inhand_imgs[num_frame] = image.transpose((2, 0, 1)) / 255.
+        #
+        # bp_imgs = torch.from_numpy(bp_imgs).to(self.device).float()
+        # inhand_imgs = torch.from_numpy(inhand_imgs).to(self.device).float()
+
+        return bp_imgs, inhand_imgs, obs, act, mask

agent/dataset/d3il_dataset/stacking_dataset.py

@@ -0,0 +1,398 @@
+import logging
+
+import os
+import glob
+
+try:
+    import cv2  # not included in pyproject.toml
+except:
+    print("Installing cv2")
+    os.system("pip install opencv-python")
+import torch
+import pickle
+import numpy as np
+from tqdm import tqdm
+
+from agent.dataset.d3il_dataset.base_dataset import TrajectoryDataset
+from agent.dataset.d3il_dataset import sim_framework_path
+
+from .geo_transform import quat2euler
+
+
+class Stacking_Dataset(TrajectoryDataset):
+    def __init__(
+        self,
+        data_directory: os.PathLike,
+        # data='train',
+        device="cpu",
+        obs_dim: int = 20,
+        action_dim: int = 2,
+        max_len_data: int = 256,
+        window_size: int = 1,
+    ):
+
+        super().__init__(
+            data_directory=data_directory,
+            device=device,
+            obs_dim=obs_dim,
+            action_dim=action_dim,
+            max_len_data=max_len_data,
+            window_size=window_size,
+        )
+
+        logging.info("Loading CubeStacking Dataset")
+
+        inputs = []
+        actions = []
+        masks = []
+
+        # for root, dirs, files in os.walk(self.data_directory):
+        #
+        #     for mode_dir in dirs:
+
+        # state_files = glob.glob(os.path.join(root, mode_dir) + "/env*")
+        # data_dir = os.path.join(sim_framework_path(data_directory), "local")
+        # data_dir = sim_framework_path(data_directory)
+        # state_files = glob.glob(data_dir + "/env*")
+
+        # bp_data_dir = sim_framework_path("environments/dataset/data/stacking/all_data_new")
+        # state_files = np.load(sim_framework_path(data_directory), allow_pickle=True)
+
+        # bp_data_dir = sim_framework_path("environments/dataset/data/stacking/single_test")
+        # state_files = os.listdir(bp_data_dir)
+
+        # random.seed(0)
+        #
+        # data_dir = sim_framework_path(data_directory)
+        # state_files = os.listdir(data_dir)
+        #
+        # random.shuffle(state_files)
+        #
+        # if data == "train":
+        #     env_state_files = state_files[20:]
+        # elif data == "eval":
+        #     env_state_files = state_files[:20]
+        # else:
+        #     assert False, "wrong data type"
+
+        data_dir = sim_framework_path("data/stacking/all_data")
+        state_files = np.load(sim_framework_path(data_directory), allow_pickle=True)
+
+        for file in state_files:
+            with open(os.path.join(data_dir, file), "rb") as f:
+                env_state = pickle.load(f)
+
+            # lengths.append(len(env_state['robot']['des_c_pos']))
+
+            zero_obs = np.zeros((1, self.max_len_data, self.obs_dim), dtype=np.float32)
+            zero_action = np.zeros(
+                (1, self.max_len_data, self.action_dim), dtype=np.float32
+            )
+            zero_mask = np.zeros((1, self.max_len_data), dtype=np.float32)
+
+            # robot and box positions
+            robot_des_j_pos = env_state["robot"]["des_j_pos"]
+            robot_des_j_vel = env_state["robot"]["des_j_vel"]
+
+            robot_des_c_pos = env_state["robot"]["des_c_pos"]
+            robot_des_quat = env_state["robot"]["des_c_quat"]
+
+            robot_c_pos = env_state["robot"]["c_pos"]
+            robot_c_quat = env_state["robot"]["c_quat"]
+
+            robot_j_pos = env_state["robot"]["j_pos"]
+            robot_j_vel = env_state["robot"]["j_vel"]
+
+            robot_gripper = np.expand_dims(env_state["robot"]["gripper_width"], -1)
+            # pred_gripper = np.zeros(robot_gripper.shape, dtype=np.float32)
+            # pred_gripper[robot_gripper > 0.075] = 1
+
+            sim_steps = np.expand_dims(np.arange(len(robot_des_j_pos)), -1)
+
+            red_box_pos = env_state["red-box"]["pos"]
+            red_box_quat = np.tan(quat2euler(env_state["red-box"]["quat"])[:, -1:])
+            # red_box_quat = np.concatenate((np.sin(red_box_quat), np.cos(red_box_quat)), axis=-1)
+
+            green_box_pos = env_state["green-box"]["pos"]
+            green_box_quat = np.tan(quat2euler(env_state["green-box"]["quat"])[:, -1:])
+            # green_box_quat = np.concatenate((np.sin(green_box_quat), np.cos(green_box_quat)), axis=-1)
+
+            blue_box_pos = env_state["blue-box"]["pos"]
+            blue_box_quat = np.tan(quat2euler(env_state["blue-box"]["quat"])[:, -1:])
+            # blue_box_quat = np.concatenate((np.sin(blue_box_quat), np.cos(blue_box_quat)), axis=-1)
+
+            # target_box_pos = env_state['target-box']['pos'] #- robot_c_pos
+
+            # input_state = np.concatenate((robot_des_c_pos, robot_des_quat, pred_gripper, red_box_pos, red_box_quat), axis=-1)
+
+            # input_state = np.concatenate((robot_des_j_pos, robot_gripper, blue_box_pos, blue_box_quat), axis=-1)
+
+            input_state = np.concatenate(
+                (
+                    robot_des_j_pos,
+                    robot_gripper,
+                    red_box_pos,
+                    red_box_quat,
+                    green_box_pos,
+                    green_box_quat,
+                    blue_box_pos,
+                    blue_box_quat,
+                ),
+                axis=-1,
+            )
+
+            # input_state = np.concatenate((robot_des_j_pos, robot_des_j_vel, robot_c_pos, robot_c_quat, green_box_pos, green_box_quat,
+            #                               target_box_pos), axis=-1)
+
+            vel_state = robot_des_j_pos[1:] - robot_des_j_pos[:-1]
+
+            valid_len = len(input_state) - 1
+
+            zero_obs[0, :valid_len, :] = input_state[:-1]
+            zero_action[0, :valid_len, :] = np.concatenate(
+                (vel_state, robot_gripper[1:]), axis=-1
+            )
+            zero_mask[0, :valid_len] = 1
+
+            inputs.append(zero_obs)
+            actions.append(zero_action)
+            masks.append(zero_mask)
+
+        # shape: B, T, n
+        self.observations = torch.from_numpy(np.concatenate(inputs)).to(device).float()
+        self.actions = torch.from_numpy(np.concatenate(actions)).to(device).float()
+        self.masks = torch.from_numpy(np.concatenate(masks)).to(device).float()
+
+        self.num_data = len(self.observations)
+
+        self.slices = self.get_slices()
+
+    def get_slices(self):
+        slices = []
+
+        min_seq_length = np.inf
+        for i in range(self.num_data):
+            T = self.get_seq_length(i)
+            min_seq_length = min(T, min_seq_length)
+
+            if T - self.window_size < 0:
+                print(
+                    f"Ignored short sequence #{i}: len={T}, window={self.window_size}"
+                )
+            else:
+                slices += [
+                    (i, start, start + self.window_size)
+                    for start in range(T - self.window_size + 1)
+                ]  # slice indices follow convention [start, end)
+
+        return slices
+
+    def get_seq_length(self, idx):
+        return int(self.masks[idx].sum().item())
+
+    def get_all_actions(self):
+        result = []
+        # mask out invalid actions
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.actions[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def get_all_observations(self):
+        result = []
+        # mask out invalid observations
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.observations[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def __len__(self):
+        return len(self.slices)
+
+    def __getitem__(self, idx):
+
+        i, start, end = self.slices[idx]
+
+        obs = self.observations[i, start:end]
+        act = self.actions[i, start:end]
+        mask = self.masks[i, start:end]
+
+        return obs, act, mask
+
+
+class Stacking_Img_Dataset(TrajectoryDataset):
+    def __init__(
+        self,
+        data_directory: os.PathLike,
+        # data='train',
+        device="cpu",
+        obs_dim: int = 20,
+        action_dim: int = 2,
+        max_len_data: int = 256,
+        window_size: int = 1,
+    ):
+
+        super().__init__(
+            data_directory=data_directory,
+            device=device,
+            obs_dim=obs_dim,
+            action_dim=action_dim,
+            max_len_data=max_len_data,
+            window_size=window_size,
+        )
+
+        logging.info("Loading CubeStacking Dataset")
+
+        inputs = []
+        actions = []
+        masks = []
+
+        # TODO: insert data_dir here
+        state_files = np.load(sim_framework_path(data_directory), allow_pickle=True)
+
+        bp_cam_imgs = []
+        inhand_cam_imgs = []
+
+        for file in tqdm(state_files):
+            with open(os.path.join(data_dir, "state", file), "rb") as f:
+                env_state = pickle.load(f)
+
+            # lengths.append(len(env_state['robot']['des_c_pos']))
+
+            zero_obs = np.zeros((1, self.max_len_data, self.obs_dim), dtype=np.float32)
+            zero_action = np.zeros(
+                (1, self.max_len_data, self.action_dim), dtype=np.float32
+            )
+            zero_mask = np.zeros((1, self.max_len_data), dtype=np.float32)
+
+            # robot and box positions
+            robot_des_j_pos = env_state["robot"]["des_j_pos"]
+            robot_des_j_vel = env_state["robot"]["des_j_vel"]
+
+            robot_des_c_pos = env_state["robot"]["des_c_pos"]
+            robot_des_quat = env_state["robot"]["des_c_quat"]
+
+            robot_c_pos = env_state["robot"]["c_pos"]
+            robot_c_quat = env_state["robot"]["c_quat"]
+
+            robot_j_pos = env_state["robot"]["j_pos"]
+            robot_j_vel = env_state["robot"]["j_vel"]
+
+            robot_gripper = np.expand_dims(env_state["robot"]["gripper_width"], -1)
+            # pred_gripper = np.zeros(robot_gripper.shape, dtype=np.float32)
+            # pred_gripper[robot_gripper > 0.075] = 1
+
+            file_name = os.path.basename(file).split(".")[0]
+
+            ###############################################################
+            bp_images = []
+            bp_imgs = glob.glob(data_dir + "/images/bp-cam/" + file_name + "/*")
+            bp_imgs.sort(key=lambda x: int(os.path.basename(x).split(".")[0]))
+
+            for img in bp_imgs:
+                image = cv2.imread(img).astype(np.float32)
+                image = image.transpose((2, 0, 1)) / 255.0
+
+                image = torch.from_numpy(image).to(self.device).float().unsqueeze(0)
+
+                bp_images.append(image)
+
+            bp_images = torch.concatenate(bp_images, dim=0)
+            ################################################################
+            inhand_imgs = glob.glob(data_dir + "/images/inhand-cam/" + file_name + "/*")
+            inhand_imgs.sort(key=lambda x: int(os.path.basename(x).split(".")[0]))
+            inhand_images = []
+            for img in inhand_imgs:
+                image = cv2.imread(img).astype(np.float32)
+                image = image.transpose((2, 0, 1)) / 255.0
+
+                image = torch.from_numpy(image).to(self.device).float().unsqueeze(0)
+
+                inhand_images.append(image)
+            inhand_images = torch.concatenate(inhand_images, dim=0)
+            ##################################################################
+
+            input_state = np.concatenate((robot_des_j_pos, robot_gripper), axis=-1)
+            vel_state = robot_des_j_pos[1:] - robot_des_j_pos[:-1]
+
+            valid_len = len(input_state) - 1
+
+            zero_obs[0, :valid_len, :] = input_state[:-1]
+            zero_action[0, :valid_len, :] = np.concatenate(
+                (vel_state, robot_gripper[1:]), axis=-1
+            )
+            zero_mask[0, :valid_len] = 1
+
+            bp_cam_imgs.append(bp_images)
+            inhand_cam_imgs.append(inhand_images)
+
+            inputs.append(zero_obs)
+            actions.append(zero_action)
+            masks.append(zero_mask)
+
+        self.bp_cam_imgs = bp_cam_imgs
+        self.inhand_cam_imgs = inhand_cam_imgs
+
+        # shape: B, T, n
+        self.observations = torch.from_numpy(np.concatenate(inputs)).to(device).float()
+        self.actions = torch.from_numpy(np.concatenate(actions)).to(device).float()
+        self.masks = torch.from_numpy(np.concatenate(masks)).to(device).float()
+
+        self.num_data = len(self.observations)
+
+        self.slices = self.get_slices()
+
+    def get_slices(self):
+        slices = []
+
+        min_seq_length = np.inf
+        for i in range(self.num_data):
+            T = self.get_seq_length(i)
+            min_seq_length = min(T, min_seq_length)
+
+            if T - self.window_size < 0:
+                print(
+                    f"Ignored short sequence #{i}: len={T}, window={self.window_size}"
+                )
+            else:
+                slices += [
+                    (i, start, start + self.window_size)
+                    for start in range(T - self.window_size + 1)
+                ]  # slice indices follow convention [start, end)
+
+        return slices
+
+    def get_seq_length(self, idx):
+        return int(self.masks[idx].sum().item())
+
+    def get_all_actions(self):
+        result = []
+        # mask out invalid actions
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.actions[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def get_all_observations(self):
+        result = []
+        # mask out invalid observations
+        for i in range(len(self.masks)):
+            T = int(self.masks[i].sum().item())
+            result.append(self.observations[i, :T, :])
+        return torch.cat(result, dim=0)
+
+    def __len__(self):
+        return len(self.slices)
+
+    def __getitem__(self, idx):
+
+        i, start, end = self.slices[idx]
+
+        obs = self.observations[i, start:end]
+        act = self.actions[i, start:end]
+        mask = self.masks[i, start:end]
+
+        bp_imgs = self.bp_cam_imgs[i][start:end]
+        inhand_imgs = self.inhand_cam_imgs[i][start:end]
+
+        return bp_imgs, inhand_imgs, obs, act, mask

agent/dataset/sequence.py

@@ -0,0 +1,162 @@
+"""
+Pre-training data loader. Modified from https://github.com/jannerm/diffuser/blob/main/diffuser/datasets/sequence.py
+
+TODO: implement history observation
+
+No normalization is applied here --- we always normalize the data when pre-processing it with a different script, and the normalization info is also used in RL fine-tuning.
+
+"""
+
+from collections import namedtuple
+from tqdm import tqdm
+import numpy as np
+import torch
+import logging
+import pickle
+import random
+
+log = logging.getLogger(__name__)
+
+from .buffer import StitchedBuffer
+
+
+Batch = namedtuple("Batch", "trajectories conditions")
+ValueBatch = namedtuple("ValueBatch", "trajectories conditions values")
+
+
+class StitchedSequenceDataset(torch.utils.data.Dataset):
+    """
+    Dataset to efficiently load and sample trajectories. Stitches episodes together in the time dimension to avoid excessive zero padding. Episode ID's are used to index unique trajectories.
+
+    Returns a dictionary with values of shape: [sum_e(T_e), *D] where T_e is traj length of episode e and D is
+    (tuple of) dimension of observation, action, images, etc.
+
+    Example:
+        Observations: [----------traj 1----------][---------traj 2----------] ... [---------traj N----------]
+        Episode IDs:  [----------   1  ----------][----------   2  ---------] ... [----------   N  ---------]
+    """
+
+    def __init__(
+        self,
+        dataset_path,
+        horizon_steps=64,
+        cond_steps=1,
+        max_n_episodes=10000,
+        use_img=False,
+        device="cpu",
+    ):
+        self.horizon_steps = horizon_steps
+        self.cond_steps = cond_steps
+        self.device = device
+
+        # Load dataset to device specified
+        if dataset_path.endswith(".npz"):
+            dataset = np.load(dataset_path, allow_pickle=True)
+        else:
+            with open(dataset_path, "rb") as f:
+                dataset = pickle.load(f)
+        num_episodes = dataset["observations"].shape[0]
+
+        # Get the sum total of the valid trajectories' lengths
+        traj_lengths = dataset["traj_length"]
+        sum_of_path_lengths = np.sum(traj_lengths)
+        self.sum_of_path_lengths = sum_of_path_lengths
+
+        fields = StitchedBuffer(sum_of_path_lengths, device)
+        for i in tqdm(
+            range(min(max_n_episodes, num_episodes)), desc="Loading trajectories"
+        ):
+            traj_length = traj_lengths[i]
+            episode = {
+                "observations": dataset["observations"][i][:traj_length],
+                "actions": dataset["actions"][i][:traj_length],
+                "episode_ids": i * np.ones(traj_length),
+            }
+            if use_img:
+                episode["images"] = dataset["images"][i][:traj_length]
+            for key, val in episode.items():
+                if device == "cpu":
+                    episode[key] = val
+                else:
+                    # if None array, save as empty tensor
+                    if np.all(np.equal(episode[key], None)):
+                        episode[key] = torch.empty(episode[key].shape).to(device)
+                    else:
+                        if key == "images":
+                            episode[key] = torch.tensor(val, dtype=torch.uint8).to(
+                                device
+                            )
+                            # (, H, W, C) -> (, C, H, W)
+                            episode[key] = episode[key].permute(0, 3, 1, 2)
+                        else:
+                            episode[key] = torch.tensor(val, dtype=torch.float32).to(
+                                device
+                            )
+            fields.add_path(episode)
+        fields.finalize()
+
+        self.indices = self.make_indices(traj_lengths, horizon_steps)
+        self.obs_dim = fields.observations.shape[-1]
+        self.action_dim = fields.actions.shape[-1]
+        self.fields = fields
+        self.n_episodes = fields.n_episodes
+        self.path_lengths = fields.path_lengths
+        self.traj_lengths = traj_lengths
+        self.use_img = use_img
+        log.info(fields)
+
+    def make_indices(self, traj_lengths, horizon_steps):
+        """
+        makes indices for sampling from dataset;
+        each index maps to a datapoint
+        """
+        indices = []
+        cur_traj_index = 0
+        for traj_length in traj_lengths:
+            max_start = cur_traj_index + traj_length - horizon_steps + 1
+            indices += list(range(cur_traj_index, max_start))
+            cur_traj_index += traj_length
+        return indices
+
+    def set_train_val_split(self, train_split):
+        num_train = int(len(self.indices) * train_split)
+        train_indices = random.sample(self.indices, num_train)
+        val_indices = [i for i in range(len(self.indices)) if i not in train_indices]
+        self.indices = train_indices
+        return val_indices
+
+    def set_indices(self, indices):
+        self.indices = indices
+
+    def get_conditions(self, observations, images=None):
+        """
+        condition on current observation for planning. Take into account the number of conditioning steps.
+        """
+        if images is not None:
+            return {
+                1 - self.cond_steps: {"state": observations[0], "rgb": images[0]}
+            }  # TODO: allow obs history, -1, -2, ...
+        else:
+            return {1 - self.cond_steps: observations[0]}
+
+    def __len__(self):
+        return len(self.indices)
+
+    def __getitem__(self, idx, eps=1e-4):
+        raise NotImplementedError("Get item defined in subclass.")
+
+
+class StitchedActionSequenceDataset(StitchedSequenceDataset):
+    """Only use action trajectory, and then obs_cond for current observation"""
+
+    def __getitem__(self, idx):
+        start = self.indices[idx]
+        end = start + self.horizon_steps
+        observations = self.fields.observations[start:end]
+        actions = self.fields.actions[start:end]
+        images = None
+        if self.use_img:
+            images = self.fields.images[start:end]
+        conditions = self.get_conditions(observations, images)
+        batch = Batch(actions, conditions)
+        return batch

agent/finetune/train_agent.py

@@ -0,0 +1,170 @@
+"""
+Parent fine-tuning agent class.
+
+"""
+
+import os
+import numpy as np
+from omegaconf import OmegaConf
+import torch
+import hydra
+import logging
+import wandb
+import random
+
+log = logging.getLogger(__name__)
+from env.gym_utils import make_async
+
+
+class TrainAgent:
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.device = cfg.device
+        self.seed = cfg.get("seed", 42)
+        random.seed(self.seed)
+        np.random.seed(self.seed)
+        torch.manual_seed(self.seed)
+
+        # Wandb
+        self.use_wandb = cfg.wandb is not None
+        if cfg.wandb is not None:
+            wandb.init(
+                entity=cfg.wandb.entity,
+                project=cfg.wandb.project,
+                name=cfg.wandb.run,
+                config=OmegaConf.to_container(cfg, resolve=True),
+            )
+
+        # Make vectorized env
+        self.env_name = cfg.env.name
+        env_type = cfg.env.get("env_type", None)
+        self.venv = make_async(
+            cfg.env.name,
+            env_type=env_type,
+            num_envs=cfg.env.n_envs,
+            asynchronous=True,
+            max_episode_steps=cfg.env.max_episode_steps,
+            wrappers=cfg.env.get("wrappers", None),
+            robomimic_env_cfg_path=cfg.get("robomimic_env_cfg_path", None),
+            shape_meta=cfg.get("shape_meta", None),
+            use_image_obs=cfg.env.get("use_image_obs", False),
+            render=cfg.env.get("render", False),
+            render_offscreen=cfg.env.get("save_video", False),
+            obs_dim=cfg.obs_dim,
+            action_dim=cfg.action_dim,
+            **cfg.env.specific if "specific" in cfg.env else {},
+        )
+        if not env_type == "furniture":
+            self.venv.seed(
+                [self.seed + i for i in range(cfg.env.n_envs)]
+            )  # otherwise parallel envs might have the same initial states!
+            # isaacgym environments do not need seeding
+        self.n_envs = cfg.env.n_envs
+        self.n_cond_step = cfg.cond_steps
+        self.obs_dim = cfg.obs_dim
+        self.action_dim = cfg.action_dim
+        self.act_steps = cfg.act_steps
+        self.horizon_steps = cfg.horizon_steps
+        self.max_episode_steps = cfg.env.max_episode_steps
+        self.reset_at_iteration = cfg.env.get("reset_at_iteration", True)
+        self.save_full_observations = cfg.env.get("save_full_observations", False)
+        self.furniture_sparse_reward = (
+            cfg.env.specific.get("sparse_reward", False)
+            if "specific" in cfg.env
+            else False
+        )  # furniture specific, for best reward calculation
+
+        # Batch size for gradient update
+        self.batch_size: int = cfg.train.batch_size
+
+        # Build model and load checkpoint
+        self.model = hydra.utils.instantiate(cfg.model)
+
+        # Training params
+        self.itr = 0
+        self.n_train_itr = cfg.train.n_train_itr
+        self.val_freq = cfg.train.val_freq
+        self.force_train = cfg.train.get("force_train", False)
+        self.n_steps = cfg.train.n_steps
+        self.best_reward_threshold_for_success = (
+            len(self.venv.pairs_to_assemble)
+            if env_type == "furniture"
+            else cfg.env.best_reward_threshold_for_success
+        )
+        self.max_grad_norm = cfg.train.get("max_grad_norm", None)
+
+        # Logging, rendering, checkpoints
+        self.logdir = cfg.logdir
+        self.render_dir = os.path.join(self.logdir, "render")
+        self.checkpoint_dir = os.path.join(self.logdir, "checkpoint")
+        self.result_path = os.path.join(self.logdir, "result.pkl")
+        os.makedirs(self.render_dir, exist_ok=True)
+        os.makedirs(self.checkpoint_dir, exist_ok=True)
+        self.save_trajs = cfg.train.get("save_trajs", False)
+        self.log_freq = cfg.train.get("log_freq", 1)
+        self.save_model_freq = cfg.train.save_model_freq
+        self.render_freq = cfg.train.render.freq
+        self.n_render = cfg.train.render.num
+        self.render_video = cfg.env.get("save_video", False)
+        assert self.n_render <= self.n_envs, "n_render must be <= n_envs"
+        assert not (
+            self.n_render <= 0 and self.render_video
+        ), "Need to set n_render > 0 if saving video"
+        self.traj_plotter = (
+            hydra.utils.instantiate(cfg.train.plotter)
+            if "plotter" in cfg.train
+            else None
+        )
+
+    def run(self):
+        pass
+
+    def save_model(self):
+        """
+        saves model to disk; no ema
+        """
+        data = {
+            "itr": self.itr,
+            "model": self.model.state_dict(),
+        }
+        savepath = os.path.join(self.checkpoint_dir, f"state_{self.itr}.pt")
+        torch.save(data, savepath)
+        log.info(f"Saved model to {savepath}")
+
+    def load(self, itr):
+        """
+        loads model from disk
+        """
+        loadpath = os.path.join(self.checkpoint_dir, f"state_{itr}.pt")
+        data = torch.load(loadpath, weights_only=True)
+
+        self.itr = data["itr"]
+        self.model.load_state_dict(data["model"])
+
+    def reset_env_all(self, verbose=False, options_venv=None, **kwargs):
+        if options_venv is None:
+            options_venv = [
+                {k: v for k, v in kwargs.items()} for _ in range(self.n_envs)
+            ]
+        obs_venv = self.venv.reset_arg(options_list=options_venv)
+        # convert to OrderedDict if obs_venv is a list of dict
+        if isinstance(obs_venv, list):
+            obs_venv = {
+                key: np.stack([obs_venv[i][key] for i in range(self.n_envs)])
+                for key in obs_venv[0].keys()
+            }
+        if verbose:
+            for index in range(self.n_envs):
+                logging.info(
+                    f"<-- Reset environment {index} with options {options_venv[index]}"
+                )
+        return obs_venv
+
+    def reset_env(self, env_ind, verbose=False):
+        task = {}
+        obs = self.venv.reset_one_arg(env_ind=env_ind, options=task)
+        if verbose:
+            logging.info(f"<-- Reset environment {env_ind} with task {task}")
+        return obs

agent/finetune/train_awr_diffusion_agent.py

@@ -0,0 +1,389 @@
+"""
+Advantage-weighted regression (AWR) for diffusion policy.
+
+Advantage = discounted-reward-to-go - V(s)
+
+"""
+
+import os
+import pickle
+import einops
+import numpy as np
+import torch
+import logging
+import wandb
+from copy import deepcopy
+
+log = logging.getLogger(__name__)
+from util.timer import Timer
+from collections import deque
+from agent.finetune.train_agent import TrainAgent
+from util.scheduler import CosineAnnealingWarmupRestarts
+
+
+def td_values(
+    states,
+    rewards,
+    dones,
+    state_values,
+    gamma=0.99,
+    alpha=0.95,
+    lam=0.95,
+):
+    """
+    Gives a list of TD estimates for a given list of samples from an RL environment.
+    The TD(λ) estimator is used for this computation.
+
+    :param replay_buffers: The replay buffers filled by exploring the RL environment.
+    Includes: states, rewards, "final state?"s.
+    :param state_values: The currently estimated state values.
+    :return: The TD estimates.
+    """
+    sample_count = len(states)
+    tds = np.zeros_like(state_values, dtype=np.float32)
+    dones[-1] = 1
+    next_value = 1 - dones[-1]
+
+    val = 0.0
+    for i in range(sample_count - 1, -1, -1):
+        # next_value = 0.0 if dones[i] else state_values[i + 1]
+
+        # get next_value for vectorized
+        if i < sample_count - 1:
+            next_value = state_values[i + 1]
+            next_value = next_value * (1 - dones[i])
+
+        state_value = state_values[i]
+        error = rewards[i] + gamma * next_value - state_value
+        val = alpha * error + gamma * lam * (1 - dones[i]) * val
+
+        tds[i] = val + state_value
+    return tds
+
+
+class TrainAWRDiffusionAgent(TrainAgent):
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+        self.logprob_batch_size = cfg.train.get("logprob_batch_size", 10000)
+
+        # note the discount factor gamma here is applied to reward every act_steps, instead of every env step
+        self.gamma = cfg.train.gamma
+
+        # Wwarm up period for critic before actor updates
+        self.n_critic_warmup_itr = cfg.train.n_critic_warmup_itr
+
+        # Optimizer
+        self.actor_optimizer = torch.optim.AdamW(
+            self.model.actor.parameters(),
+            lr=cfg.train.actor_lr,
+            weight_decay=cfg.train.actor_weight_decay,
+        )
+        self.actor_lr_scheduler = CosineAnnealingWarmupRestarts(
+            self.actor_optimizer,
+            first_cycle_steps=cfg.train.actor_lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=cfg.train.actor_lr,
+            min_lr=cfg.train.actor_lr_scheduler.min_lr,
+            warmup_steps=cfg.train.actor_lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+        self.critic_optimizer = torch.optim.AdamW(
+            self.model.critic.parameters(),
+            lr=cfg.train.critic_lr,
+            weight_decay=cfg.train.critic_weight_decay,
+        )
+        self.critic_lr_scheduler = CosineAnnealingWarmupRestarts(
+            self.critic_optimizer,
+            first_cycle_steps=cfg.train.critic_lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=cfg.train.critic_lr,
+            min_lr=cfg.train.critic_lr_scheduler.min_lr,
+            warmup_steps=cfg.train.critic_lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+
+        # Buffer size
+        self.buffer_size = cfg.train.buffer_size
+
+        # Reward exponential
+        self.beta = cfg.train.beta
+
+        # Max weight for AWR
+        self.max_adv_weight = cfg.train.max_adv_weight
+
+        # Scaling reward
+        self.scale_reward_factor = cfg.train.scale_reward_factor
+
+        # Updates
+        self.replay_ratio = cfg.train.replay_ratio
+        self.critic_update_ratio = cfg.train.critic_update_ratio
+
+    def run(self):
+
+        # make a FIFO replay buffer for obs, action, and reward
+        obs_buffer = deque(maxlen=self.buffer_size)
+        action_buffer = deque(maxlen=self.buffer_size)
+        reward_buffer = deque(maxlen=self.buffer_size)
+        done_buffer = deque(maxlen=self.buffer_size)
+        first_buffer = deque(maxlen=self.buffer_size)
+
+        # Start training loop
+        timer = Timer()
+        run_results = []
+        done_venv = np.zeros((1, self.n_envs))
+        while self.itr < self.n_train_itr:
+
+            # Prepare video paths for each envs --- only applies for the first set of episodes if allowing reset within iteration and each iteration has multiple episodes from one env
+            options_venv = [{} for _ in range(self.n_envs)]
+            if self.itr % self.render_freq == 0 and self.render_video:
+                for env_ind in range(self.n_render):
+                    options_venv[env_ind]["video_path"] = os.path.join(
+                        self.render_dir, f"itr-{self.itr}_trial-{env_ind}.mp4"
+                    )
+
+            # Define train or eval - all envs restart
+            eval_mode = self.itr % self.val_freq == 0 and not self.force_train
+            self.model.eval() if eval_mode else self.model.train()
+            firsts_trajs = np.zeros((self.n_steps + 1, self.n_envs))
+
+            # Reset env at the beginning of an iteration
+            if self.reset_at_iteration or eval_mode or last_itr_eval:
+                prev_obs_venv = self.reset_env_all(options_venv=options_venv)
+                firsts_trajs[0] = 1
+            else:
+                firsts_trajs[0] = (
+                    done_venv  # if done at the end of last iteration, then the envs are just reset
+                )
+            last_itr_eval = eval_mode
+            reward_trajs = np.empty((0, self.n_envs))
+
+            # Collect a set of trajectories from env
+            for step in range(self.n_steps):
+                if step % 10 == 0:
+                    print(f"Processed step {step} of {self.n_steps}")
+
+                # Select action
+                with torch.no_grad():
+                    samples = (
+                        self.model(
+                            cond=torch.from_numpy(prev_obs_venv)
+                            .float()
+                            .to(self.device),
+                            deterministic=eval_mode,
+                        )
+                        .cpu()
+                        .numpy()
+                    )  # n_env x horizon x act
+                action_venv = samples[:, : self.act_steps]
+
+                # Apply multi-step action
+                obs_venv, reward_venv, done_venv, info_venv = self.venv.step(
+                    action_venv
+                )
+                reward_trajs = np.vstack((reward_trajs, reward_venv[None]))
+
+                # add to buffer
+                obs_buffer.append(prev_obs_venv)
+                action_buffer.append(action_venv)
+                reward_buffer.append(reward_venv * self.scale_reward_factor)
+                done_buffer.append(done_venv)
+                first_buffer.append(firsts_trajs[step])
+
+                firsts_trajs[step + 1] = done_venv
+                prev_obs_venv = obs_venv
+
+            # Summarize episode reward --- this needs to be handled differently depending on whether the environment is reset after each iteration. Only count episodes that finish within the iteration.
+            episodes_start_end = []
+            for env_ind in range(self.n_envs):
+                env_steps = np.where(firsts_trajs[:, env_ind] == 1)[0]
+                for i in range(len(env_steps) - 1):
+                    start = env_steps[i]
+                    end = env_steps[i + 1]
+                    if end - start > 1:
+                        episodes_start_end.append((env_ind, start, end - 1))
+            if len(episodes_start_end) > 0:
+                reward_trajs_split = [
+                    reward_trajs[start : end + 1, env_ind]
+                    for env_ind, start, end in episodes_start_end
+                ]
+                num_episode_finished = len(reward_trajs_split)
+                episode_reward = np.array(
+                    [np.sum(reward_traj) for reward_traj in reward_trajs_split]
+                )
+                episode_best_reward = np.array(
+                    [
+                        np.max(reward_traj) / self.act_steps
+                        for reward_traj in reward_trajs_split
+                    ]
+                )
+                avg_episode_reward = np.mean(episode_reward)
+                avg_best_reward = np.mean(episode_best_reward)
+                success_rate = np.mean(
+                    episode_best_reward >= self.best_reward_threshold_for_success
+                )
+            else:
+                episode_reward = np.array([])
+                num_episode_finished = 0
+                avg_episode_reward = 0
+                avg_best_reward = 0
+                success_rate = 0
+                log.info("[WARNING] No episode completed within the iteration!")
+
+            # Update
+            if not eval_mode:
+
+                obs_trajs = np.array(deepcopy(obs_buffer))
+                reward_trajs = np.array(deepcopy(reward_buffer))
+                dones_trajs = np.array(deepcopy(done_buffer))
+
+                obs_t = einops.rearrange(
+                    torch.from_numpy(obs_trajs).float().to(self.device),
+                    "s e h d -> (s e) h d",
+                )
+                values_t = np.array(self.model.critic(obs_t).detach().cpu().numpy())
+                values_trajs = values_t.reshape(-1, self.n_envs)
+                td_trajs = td_values(obs_trajs, reward_trajs, dones_trajs, values_trajs)
+
+                # flatten
+                obs_trajs = einops.rearrange(
+                    obs_trajs,
+                    "s e h d -> (s e) h d",
+                )
+                td_trajs = einops.rearrange(
+                    td_trajs,
+                    "s e -> (s e)",
+                )
+
+                # Update policy and critic
+                num_batch = int(
+                    self.n_steps * self.n_envs / self.batch_size * self.replay_ratio
+                )
+                for _ in range(num_batch // self.critic_update_ratio):
+
+                    # Sample batch
+                    inds = np.random.choice(len(obs_trajs), self.batch_size)
+                    obs_b = torch.from_numpy(obs_trajs[inds]).float().to(self.device)
+                    td_b = torch.from_numpy(td_trajs[inds]).float().to(self.device)
+
+                    # Update critic
+                    loss_critic = self.model.loss_critic(obs_b, td_b)
+                    self.critic_optimizer.zero_grad()
+                    loss_critic.backward()
+                    self.critic_optimizer.step()
+
+                obs_trajs = np.array(deepcopy(obs_buffer))
+                samples_trajs = np.array(deepcopy(action_buffer))
+                reward_trajs = np.array(deepcopy(reward_buffer))
+                dones_trajs = np.array(deepcopy(done_buffer))
+
+                obs_t = einops.rearrange(
+                    torch.from_numpy(obs_trajs).float().to(self.device),
+                    "s e h d -> (s e) h d",
+                )
+                values_t = np.array(self.model.critic(obs_t).detach().cpu().numpy())
+                values_trajs = values_t.reshape(-1, self.n_envs)
+                td_trajs = td_values(obs_trajs, reward_trajs, dones_trajs, values_trajs)
+                advantages_trajs = td_trajs - values_trajs
+
+                # flatten
+                obs_trajs = einops.rearrange(
+                    obs_trajs,
+                    "s e h d -> (s e) h d",
+                )
+                samples_trajs = einops.rearrange(
+                    samples_trajs,
+                    "s e h d -> (s e) h d",
+                )
+                advantages_trajs = einops.rearrange(
+                    advantages_trajs,
+                    "s e -> (s e)",
+                )
+
+                for _ in range(num_batch):
+
+                    # Sample batch
+                    inds = np.random.choice(len(obs_trajs), self.batch_size)
+                    obs_b = torch.from_numpy(obs_trajs[inds]).float().to(self.device)
+                    actions_b = (
+                        torch.from_numpy(samples_trajs[inds]).float().to(self.device)
+                    )
+                    advantages_b = (
+                        torch.from_numpy(advantages_trajs[inds]).float().to(self.device)
+                    )
+                    advantages_b = (advantages_b - advantages_b.mean()) / (
+                        advantages_b.std() + 1e-6
+                    )
+                    advantages_b_scaled = torch.exp(self.beta * advantages_b)
+                    advantages_b_scaled.clamp_(max=self.max_adv_weight)
+
+                    # Update policy with collected trajectories
+                    loss = self.model.loss(
+                        actions_b,
+                        obs_b,
+                        advantages_b_scaled.detach(),
+                    )
+                    self.actor_optimizer.zero_grad()
+                    loss.backward()
+                    if self.itr >= self.n_critic_warmup_itr:
+                        if self.max_grad_norm is not None:
+                            torch.nn.utils.clip_grad_norm_(
+                                self.model.actor.parameters(), self.max_grad_norm
+                            )
+                        self.actor_optimizer.step()
+
+            # Update lr
+            self.actor_lr_scheduler.step()
+            self.critic_lr_scheduler.step()
+
+            # Save model
+            if self.itr % self.save_model_freq == 0 or self.itr == self.n_train_itr - 1:
+                self.save_model()
+
+            # Log loss and save metrics
+            run_results.append(
+                {
+                    "itr": self.itr,
+                }
+            )
+            if self.itr % self.log_freq == 0:
+                if eval_mode:
+                    log.info(
+                        f"eval: success rate {success_rate:8.4f} | avg episode reward {avg_episode_reward:8.4f} | avg best reward {avg_best_reward:8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "success rate - eval": success_rate,
+                                "avg episode reward - eval": avg_episode_reward,
+                                "avg best reward - eval": avg_best_reward,
+                                "num episode - eval": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=False,
+                        )
+                    run_results[-1]["eval_success_rate"] = success_rate
+                    run_results[-1]["eval_episode_reward"] = avg_episode_reward
+                    run_results[-1]["eval_best_reward"] = avg_best_reward
+                else:
+                    log.info(
+                        f"{self.itr}: loss {loss:8.4f} | reward {avg_episode_reward:8.4f} |t:{timer():8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "loss": loss,
+                                "loss - critic": loss_critic,
+                                "avg episode reward - train": avg_episode_reward,
+                                "num episode - train": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=True,
+                        )
+                    run_results[-1]["loss"] = loss
+                    run_results[-1]["loss_critic"] = loss_critic
+                    run_results[-1]["train_episode_reward"] = avg_episode_reward
+                run_results[-1]["time"] = timer()
+                with open(self.result_path, "wb") as f:
+                    pickle.dump(run_results, f)
+            self.itr += 1

agent/finetune/train_dipo_diffusion_agent.py

@@ -0,0 +1,358 @@
+"""
+Model-free online RL with DIffusion POlicy (DIPO)
+
+Applies action gradient to perturb actions towards maximizer of Q-function.
+
+a_t <- a_t + \eta * \grad_a Q(s, a)
+"""
+
+import os
+import pickle
+import numpy as np
+import torch
+import logging
+import wandb
+
+log = logging.getLogger(__name__)
+from util.timer import Timer
+from collections import deque
+from agent.finetune.train_agent import TrainAgent
+from util.scheduler import CosineAnnealingWarmupRestarts
+
+
+class TrainDIPODiffusionAgent(TrainAgent):
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+        # note the discount factor gamma here is applied to reward every act_steps, instead of every env step
+        self.gamma = cfg.train.gamma
+
+        # Wwarm up period for critic before actor updates
+        self.n_critic_warmup_itr = cfg.train.n_critic_warmup_itr
+
+        # Optimizer
+        self.actor_optimizer = torch.optim.AdamW(
+            self.model.actor.parameters(),
+            lr=cfg.train.actor_lr,
+            weight_decay=cfg.train.actor_weight_decay,
+        )
+        # use cosine scheduler with linear warmup
+        self.actor_lr_scheduler = CosineAnnealingWarmupRestarts(
+            self.actor_optimizer,
+            first_cycle_steps=cfg.train.actor_lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=cfg.train.actor_lr,
+            min_lr=cfg.train.actor_lr_scheduler.min_lr,
+            warmup_steps=cfg.train.actor_lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+        self.critic_optimizer = torch.optim.AdamW(
+            self.model.critic.parameters(),
+            lr=cfg.train.critic_lr,
+            weight_decay=cfg.train.critic_weight_decay,
+        )
+        self.critic_lr_scheduler = CosineAnnealingWarmupRestarts(
+            self.critic_optimizer,
+            first_cycle_steps=cfg.train.critic_lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=cfg.train.critic_lr,
+            min_lr=cfg.train.critic_lr_scheduler.min_lr,
+            warmup_steps=cfg.train.critic_lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+
+        # Buffer size
+        self.buffer_size = cfg.train.buffer_size
+
+        # Perturbation scale
+        self.eta = cfg.train.eta
+
+        # Updates
+        self.replay_ratio = cfg.train.replay_ratio
+
+        # Scaling reward
+        self.scale_reward_factor = cfg.train.scale_reward_factor
+
+        # Apply action gradient many steps
+        self.action_gradient_steps = cfg.train.action_gradient_steps
+
+    def run(self):
+
+        # make a FIFO replay buffer for obs, action, and reward
+        obs_buffer = deque(maxlen=self.buffer_size)
+        next_obs_buffer = deque(maxlen=self.buffer_size)
+        action_buffer = deque(maxlen=self.buffer_size)
+        reward_buffer = deque(maxlen=self.buffer_size)
+        done_buffer = deque(maxlen=self.buffer_size)
+        first_buffer = deque(maxlen=self.buffer_size)
+
+        # Start training loop
+        timer = Timer()
+        run_results = []
+        done_venv = np.zeros((1, self.n_envs))
+        while self.itr < self.n_train_itr:
+
+            # Prepare video paths for each envs --- only applies for the first set of episodes if allowing reset within iteration and each iteration has multiple episodes from one env
+            options_venv = [{} for _ in range(self.n_envs)]
+            if self.itr % self.render_freq == 0 and self.render_video:
+                for env_ind in range(self.n_render):
+                    options_venv[env_ind]["video_path"] = os.path.join(
+                        self.render_dir, f"itr-{self.itr}_trial-{env_ind}.mp4"
+                    )
+
+            # Define train or eval - all envs restart
+            eval_mode = self.itr % self.val_freq == 0 and not self.force_train
+            self.model.eval() if eval_mode else self.model.train()
+            firsts_trajs = np.zeros((self.n_steps + 1, self.n_envs))
+
+            # Reset env at the beginning of an iteration
+            if self.reset_at_iteration or eval_mode or last_itr_eval:
+                prev_obs_venv = self.reset_env_all(options_venv=options_venv)
+                firsts_trajs[0] = 1
+            else:
+                firsts_trajs[0] = (
+                    done_venv  # if done at the end of last iteration, then the envs are just reset
+                )
+            last_itr_eval = eval_mode
+            reward_trajs = np.empty((0, self.n_envs))
+
+            # Collect a set of trajectories from env
+            for step in range(self.n_steps):
+                if step % 10 == 0:
+                    print(f"Processed step {step} of {self.n_steps}")
+
+                # Select action
+                with torch.no_grad():
+                    samples = (
+                        self.model(
+                            cond=torch.from_numpy(prev_obs_venv)
+                            .float()
+                            .to(self.device),
+                            deterministic=eval_mode,
+                        )
+                        .cpu()
+                        .numpy()
+                    )  # n_env x horizon x act
+                action_venv = samples[:, : self.act_steps]
+
+                # Apply multi-step action
+                obs_venv, reward_venv, done_venv, info_venv = self.venv.step(
+                    action_venv
+                )
+                reward_trajs = np.vstack((reward_trajs, reward_venv[None]))
+
+                # add to buffer
+                for i in range(self.n_envs):
+                    obs_buffer.append(prev_obs_venv[i])
+                    next_obs_buffer.append(obs_venv[i])
+                    action_buffer.append(action_venv[i])
+                    reward_buffer.append(reward_venv[i] * self.scale_reward_factor)
+                    done_buffer.append(done_venv[i])
+                first_buffer.append(firsts_trajs[step])
+
+                firsts_trajs[step + 1] = done_venv
+                prev_obs_venv = obs_venv
+
+            # Summarize episode reward --- this needs to be handled differently depending on whether the environment is reset after each iteration. Only count episodes that finish within the iteration.
+            episodes_start_end = []
+            for env_ind in range(self.n_envs):
+                env_steps = np.where(firsts_trajs[:, env_ind] == 1)[0]
+                for i in range(len(env_steps) - 1):
+                    start = env_steps[i]
+                    end = env_steps[i + 1]
+                    if end - start > 1:
+                        episodes_start_end.append((env_ind, start, end - 1))
+            if len(episodes_start_end) > 0:
+                reward_trajs_split = [
+                    reward_trajs[start : end + 1, env_ind]
+                    for env_ind, start, end in episodes_start_end
+                ]
+                num_episode_finished = len(reward_trajs_split)
+                episode_reward = np.array(
+                    [np.sum(reward_traj) for reward_traj in reward_trajs_split]
+                )
+                episode_best_reward = np.array(
+                    [
+                        np.max(reward_traj) / self.act_steps
+                        for reward_traj in reward_trajs_split
+                    ]
+                )
+                avg_episode_reward = np.mean(episode_reward)
+                avg_best_reward = np.mean(episode_best_reward)
+                success_rate = np.mean(
+                    episode_best_reward >= self.best_reward_threshold_for_success
+                )
+            else:
+                episode_reward = np.array([])
+                num_episode_finished = 0
+                avg_episode_reward = 0
+                avg_best_reward = 0
+                success_rate = 0
+                log.info("[WARNING] No episode completed within the iteration!")
+
+            if not eval_mode:
+
+                num_batch = self.replay_ratio
+
+                # Critic learning
+                for _ in range(num_batch):
+                    # Sample batch
+                    inds = np.random.choice(len(obs_buffer), self.batch_size)
+                    obs_b = (
+                        torch.from_numpy(np.vstack([obs_buffer[i][None] for i in inds]))
+                        .float()
+                        .to(self.device)
+                    )
+                    next_obs_b = (
+                        torch.from_numpy(
+                            np.vstack([next_obs_buffer[i][None] for i in inds])
+                        )
+                        .float()
+                        .to(self.device)
+                    )
+                    actions_b = (
+                        torch.from_numpy(
+                            np.vstack([action_buffer[i][None] for i in inds])
+                        )
+                        .float()
+                        .to(self.device)
+                    )
+                    rewards_b = (
+                        torch.from_numpy(np.vstack([reward_buffer[i] for i in inds]))
+                        .float()
+                        .to(self.device)
+                    )
+                    dones_b = (
+                        torch.from_numpy(np.vstack([done_buffer[i] for i in inds]))
+                        .float()
+                        .to(self.device)
+                    )
+
+                    # Update critic
+                    loss_critic = self.model.loss_critic(
+                        obs_b, next_obs_b, actions_b, rewards_b, dones_b, self.gamma
+                    )
+                    self.critic_optimizer.zero_grad()
+                    loss_critic.backward()
+                    self.critic_optimizer.step()
+
+                # Actor learning
+                for _ in range(num_batch):
+
+                    # Sample batch
+                    inds = np.random.choice(len(obs_buffer), self.batch_size)
+                    obs_b = (
+                        torch.from_numpy(np.vstack([obs_buffer[i][None] for i in inds]))
+                        .float()
+                        .to(self.device)
+                    )
+                    actions_b = (
+                        torch.from_numpy(
+                            np.vstack([action_buffer[i][None] for i in inds])
+                        )
+                        .float()
+                        .to(self.device)
+                    )
+
+                    # Replace actions in buffer with guided actions
+                    guided_action_list = []
+
+                    # get Q-perturbed actions by optimizing
+                    actions_flat = actions_b.reshape(actions_b.shape[0], -1)
+                    actions_optim = torch.optim.Adam(
+                        [actions_flat], lr=self.eta, eps=1e-5
+                    )
+                    for _ in range(self.action_gradient_steps):
+                        actions_flat.requires_grad_(True)
+                        q_values_1, q_values_2 = self.model.critic(obs_b, actions_flat)
+                        q_values = torch.min(q_values_1, q_values_2)
+                        action_opt_loss = -q_values.sum()
+
+                        actions_optim.zero_grad()
+                        action_opt_loss.backward(torch.ones_like(action_opt_loss))
+
+                        # get the perturbed action
+                        actions_optim.step()
+
+                        actions_flat.requires_grad_(False)
+                        actions_flat.clamp_(-1.0, 1.0)
+                    guided_action = actions_flat.detach()
+                    guided_action = guided_action.reshape(
+                        guided_action.shape[0], -1, self.action_dim
+                    )
+                    guided_action_list.append(guided_action)
+                    guided_action_stacked = torch.cat(guided_action_list, 0)
+
+                    # Add to buffer (need separate indices since we're working with a limited subset)
+                    for i, i_buf in enumerate(inds):
+                        action_buffer[i_buf] = (
+                            guided_action_stacked[i].detach().cpu().numpy()
+                        )
+
+                    # Update policy with collected trajectories
+                    loss = self.model.loss(guided_action.detach(), {0: obs_b})
+                    self.actor_optimizer.zero_grad()
+                    loss.backward()
+                    if self.itr >= self.n_critic_warmup_itr:
+                        if self.max_grad_norm is not None:
+                            torch.nn.utils.clip_grad_norm_(
+                                self.model.actor.parameters(), self.max_grad_norm
+                            )
+                        self.actor_optimizer.step()
+
+            # Update lr
+            self.actor_lr_scheduler.step()
+            self.critic_lr_scheduler.step()
+
+            # Save model
+            if self.itr % self.save_model_freq == 0 or self.itr == self.n_train_itr - 1:
+                self.save_model()
+
+            # Log loss and save metrics
+            run_results.append(
+                {
+                    "itr": self.itr,
+                }
+            )
+            if self.itr % self.log_freq == 0:
+                if eval_mode:
+                    log.info(
+                        f"eval: success rate {success_rate:8.4f} | avg episode reward {avg_episode_reward:8.4f} | avg best reward {avg_best_reward:8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "success rate - eval": success_rate,
+                                "avg episode reward - eval": avg_episode_reward,
+                                "avg best reward - eval": avg_best_reward,
+                                "num episode - eval": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=False,
+                        )
+                    run_results[-1]["eval_success_rate"] = success_rate
+                    run_results[-1]["eval_episode_reward"] = avg_episode_reward
+                    run_results[-1]["eval_best_reward"] = avg_best_reward
+                else:
+                    log.info(
+                        f"{self.itr}: loss {loss:8.4f} | reward {avg_episode_reward:8.4f} |t:{timer():8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "loss": loss,
+                                "loss - critic": loss_critic,
+                                "avg episode reward - train": avg_episode_reward,
+                                "num episode - train": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=True,
+                        )
+                    run_results[-1]["loss"] = loss
+                    run_results[-1]["loss_critic"] = loss_critic
+                    run_results[-1]["train_episode_reward"] = avg_episode_reward
+                run_results[-1]["time"] = timer()
+                with open(self.result_path, "wb") as f:
+                    pickle.dump(run_results, f)
+            self.itr += 1

agent/finetune/train_dql_diffusion_agent.py

@@ -0,0 +1,317 @@
+"""
+Diffusion Q-Learning (DQL)
+
+Learns a critic Q-function and backprops the expected Q-value to train the actor
+
+pi = argmin L_d(\theta) - \alpha * E[Q(s, a)]
+L_d is demonstration loss for regularization
+"""
+
+import os
+import pickle
+import numpy as np
+import torch
+import logging
+import wandb
+
+log = logging.getLogger(__name__)
+from util.timer import Timer
+from collections import deque
+from agent.finetune.train_agent import TrainAgent
+from util.scheduler import CosineAnnealingWarmupRestarts
+
+
+class TrainDQLDiffusionAgent(TrainAgent):
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+        # note the discount factor gamma here is applied to reward every act_steps, instead of every env step
+        self.gamma = cfg.train.gamma
+
+        # Wwarm up period for critic before actor updates
+        self.n_critic_warmup_itr = cfg.train.n_critic_warmup_itr
+
+        # Optimizer
+        self.actor_optimizer = torch.optim.AdamW(
+            self.model.actor.parameters(),
+            lr=cfg.train.actor_lr,
+            weight_decay=cfg.train.actor_weight_decay,
+        )
+        # use cosine scheduler with linear warmup
+        self.actor_lr_scheduler = CosineAnnealingWarmupRestarts(
+            self.actor_optimizer,
+            first_cycle_steps=cfg.train.actor_lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=cfg.train.actor_lr,
+            min_lr=cfg.train.actor_lr_scheduler.min_lr,
+            warmup_steps=cfg.train.actor_lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+        self.critic_optimizer = torch.optim.AdamW(
+            self.model.critic.parameters(),
+            lr=cfg.train.critic_lr,
+            weight_decay=cfg.train.critic_weight_decay,
+        )
+        self.critic_lr_scheduler = CosineAnnealingWarmupRestarts(
+            self.critic_optimizer,
+            first_cycle_steps=cfg.train.critic_lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=cfg.train.critic_lr,
+            min_lr=cfg.train.critic_lr_scheduler.min_lr,
+            warmup_steps=cfg.train.critic_lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+
+        # Buffer size
+        self.buffer_size = cfg.train.buffer_size
+
+        # Perturbation scale
+        self.eta = cfg.train.eta
+
+        # Reward factor - scale down mujoco reward for better critic training
+        self.scale_reward_factor = cfg.train.scale_reward_factor
+
+        # Updates
+        self.replay_ratio = cfg.train.replay_ratio
+
+    def run(self):
+
+        # make a FIFO replay buffer for obs, action, and reward
+        obs_buffer = deque(maxlen=self.buffer_size)
+        next_obs_buffer = deque(maxlen=self.buffer_size)
+        action_buffer = deque(maxlen=self.buffer_size)
+        reward_buffer = deque(maxlen=self.buffer_size)
+        done_buffer = deque(maxlen=self.buffer_size)
+        first_buffer = deque(maxlen=self.buffer_size)
+
+        # Start training loop
+        timer = Timer()
+        run_results = []
+        done_venv = np.zeros((1, self.n_envs))
+        while self.itr < self.n_train_itr:
+
+            # Prepare video paths for each envs --- only applies for the first set of episodes if allowing reset within iteration and each iteration has multiple episodes from one env
+            options_venv = [{} for _ in range(self.n_envs)]
+            if self.itr % self.render_freq == 0 and self.render_video:
+                for env_ind in range(self.n_render):
+                    options_venv[env_ind]["video_path"] = os.path.join(
+                        self.render_dir, f"itr-{self.itr}_trial-{env_ind}.mp4"
+                    )
+
+            # Define train or eval - all envs restart
+            eval_mode = self.itr % self.val_freq == 0 and not self.force_train
+            self.model.eval() if eval_mode else self.model.train()
+            firsts_trajs = np.zeros((self.n_steps + 1, self.n_envs))
+
+            # Reset env at the beginning of an iteration
+            if self.reset_at_iteration or eval_mode or last_itr_eval:
+                prev_obs_venv = self.reset_env_all(options_venv=options_venv)
+                firsts_trajs[0] = 1
+            else:
+                firsts_trajs[0] = (
+                    done_venv  # if done at the end of last iteration, then the envs are just reset
+                )
+            last_itr_eval = eval_mode
+            reward_trajs = np.empty((0, self.n_envs))
+
+            # Collect a set of trajectories from env
+            for step in range(self.n_steps):
+                if step % 10 == 0:
+                    print(f"Processed step {step} of {self.n_steps}")
+
+                # Select action
+                with torch.no_grad():
+                    samples = (
+                        self.model(
+                            cond=torch.from_numpy(prev_obs_venv)
+                            .float()
+                            .to(self.device),
+                            deterministic=eval_mode,
+                        )
+                        .cpu()
+                        .numpy()
+                    )  # n_env x horizon x act
+                action_venv = samples[:, : self.act_steps]
+
+                # Apply multi-step action
+                obs_venv, reward_venv, done_venv, info_venv = self.venv.step(
+                    action_venv
+                )
+                reward_trajs = np.vstack((reward_trajs, reward_venv[None]))
+
+                # add to buffer
+                for i in range(self.n_envs):
+                    obs_buffer.append(prev_obs_venv[i])
+                    next_obs_buffer.append(obs_venv[i])
+                    action_buffer.append(action_venv[i])
+                    reward_buffer.append(reward_venv[i] * self.scale_reward_factor)
+                    done_buffer.append(done_venv[i])
+                first_buffer.append(firsts_trajs[step])
+
+                firsts_trajs[step + 1] = done_venv
+                prev_obs_venv = obs_venv
+
+            # Summarize episode reward --- this needs to be handled differently depending on whether the environment is reset after each iteration. Only count episodes that finish within the iteration.
+            episodes_start_end = []
+            for env_ind in range(self.n_envs):
+                env_steps = np.where(firsts_trajs[:, env_ind] == 1)[0]
+                for i in range(len(env_steps) - 1):
+                    start = env_steps[i]
+                    end = env_steps[i + 1]
+                    if end - start > 1:
+                        episodes_start_end.append((env_ind, start, end - 1))
+            if len(episodes_start_end) > 0:
+                reward_trajs_split = [
+                    reward_trajs[start : end + 1, env_ind]
+                    for env_ind, start, end in episodes_start_end
+                ]
+                num_episode_finished = len(reward_trajs_split)
+                episode_reward = np.array(
+                    [np.sum(reward_traj) for reward_traj in reward_trajs_split]
+                )
+                episode_best_reward = np.array(
+                    [
+                        np.max(reward_traj) / self.act_steps
+                        for reward_traj in reward_trajs_split
+                    ]
+                )
+                avg_episode_reward = np.mean(episode_reward)
+                avg_best_reward = np.mean(episode_best_reward)
+                success_rate = np.mean(
+                    episode_best_reward >= self.best_reward_threshold_for_success
+                )
+            else:
+                episode_reward = np.array([])
+                num_episode_finished = 0
+                avg_episode_reward = 0
+                avg_best_reward = 0
+                success_rate = 0
+                log.info("[WARNING] No episode completed within the iteration!")
+
+            if not eval_mode:
+
+                num_batch = self.replay_ratio
+
+                # Critic learning
+                for _ in range(num_batch):
+                    # Sample batch
+                    inds = np.random.choice(len(obs_buffer), self.batch_size)
+                    obs_b = (
+                        torch.from_numpy(np.vstack([obs_buffer[i][None] for i in inds]))
+                        .float()
+                        .to(self.device)
+                    )
+                    next_obs_b = (
+                        torch.from_numpy(
+                            np.vstack([next_obs_buffer[i][None] for i in inds])
+                        )
+                        .float()
+                        .to(self.device)
+                    )
+                    actions_b = (
+                        torch.from_numpy(
+                            np.vstack([action_buffer[i][None] for i in inds])
+                        )
+                        .float()
+                        .to(self.device)
+                    )
+                    rewards_b = (
+                        torch.from_numpy(np.vstack([reward_buffer[i] for i in inds]))
+                        .float()
+                        .to(self.device)
+                    )
+                    dones_b = (
+                        torch.from_numpy(np.vstack([done_buffer[i] for i in inds]))
+                        .float()
+                        .to(self.device)
+                    )
+
+                    # Update critic
+                    loss_critic = self.model.loss_critic(
+                        obs_b, next_obs_b, actions_b, rewards_b, dones_b, self.gamma
+                    )
+                    self.critic_optimizer.zero_grad()
+                    loss_critic.backward()
+                    self.critic_optimizer.step()
+
+                    # get the new action and q values
+                    samples = self.model.forward_train(
+                        cond=obs_b.to(self.device),
+                        deterministic=eval_mode,
+                    )
+                    output_venv = samples  # n_env x horizon x act
+                    action_venv = output_venv[:, : self.act_steps, : self.action_dim]
+                    actions_flat_b = action_venv.reshape(action_venv.shape[0], -1)
+                    q_values_b = self.model.critic(obs_b, actions_flat_b)
+                    q1_new_action, q2_new_action = q_values_b
+
+                    # Update policy with collected trajectories
+                    self.actor_optimizer.zero_grad()
+                    actor_loss = self.model.loss_actor(
+                        obs_b, actions_b, q1_new_action, q2_new_action, self.eta
+                    )
+                    actor_loss.backward()
+                    if self.itr >= self.n_critic_warmup_itr:
+                        if self.max_grad_norm is not None:
+                            torch.nn.utils.clip_grad_norm_(
+                                self.model.actor.parameters(), self.max_grad_norm
+                            )
+                        self.actor_optimizer.step()
+                    loss = actor_loss
+
+            # Update lr
+            self.actor_lr_scheduler.step()
+            self.critic_lr_scheduler.step()
+
+            # Save model
+            if self.itr % self.save_model_freq == 0 or self.itr == self.n_train_itr - 1:
+                self.save_model()
+
+            # Log loss and save metrics
+            run_results.append(
+                {
+                    "itr": self.itr,
+                }
+            )
+            if self.itr % self.log_freq == 0:
+                if eval_mode:
+                    log.info(
+                        f"eval: success rate {success_rate:8.4f} | avg episode reward {avg_episode_reward:8.4f} | avg best reward {avg_best_reward:8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "success rate - eval": success_rate,
+                                "avg episode reward - eval": avg_episode_reward,
+                                "avg best reward - eval": avg_best_reward,
+                                "num episode - eval": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=False,
+                        )
+                    run_results[-1]["eval_success_rate"] = success_rate
+                    run_results[-1]["eval_episode_reward"] = avg_episode_reward
+                    run_results[-1]["eval_best_reward"] = avg_best_reward
+                else:
+                    log.info(
+                        f"{self.itr}: loss {loss:8.4f} | reward {avg_episode_reward:8.4f} |t:{timer():8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "loss": loss,
+                                "loss - critic": loss_critic,
+                                "avg episode reward - train": avg_episode_reward,
+                                "num episode - train": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=True,
+                        )
+                    run_results[-1]["loss"] = loss
+                    run_results[-1]["loss_critic"] = loss_critic
+                    run_results[-1]["train_episode_reward"] = avg_episode_reward
+                run_results[-1]["time"] = timer()
+                with open(self.result_path, "wb") as f:
+                    pickle.dump(run_results, f)
+            self.itr += 1

agent/finetune/train_idql_diffusion_agent.py

@@ -0,0 +1,347 @@
+"""
+Implicit diffusion Q-learning (IDQL) trainer for diffusion policy.
+
+"""
+
+import os
+import pickle
+import einops
+import numpy as np
+import torch
+import logging
+import wandb
+from copy import deepcopy
+import random
+
+log = logging.getLogger(__name__)
+from util.timer import Timer
+from collections import deque
+from agent.finetune.train_agent import TrainAgent
+from util.scheduler import CosineAnnealingWarmupRestarts
+
+
+class TrainIDQLDiffusionAgent(TrainAgent):
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+        # note the discount factor gamma here is applied to reward every act_steps, instead of every env step
+        self.gamma = cfg.train.gamma
+
+        # Wwarm up period for critic before actor updates
+        self.n_critic_warmup_itr = cfg.train.n_critic_warmup_itr
+
+        # Optimizer
+        self.actor_optimizer = torch.optim.AdamW(
+            self.model.actor.parameters(),
+            lr=cfg.train.actor_lr,
+            weight_decay=cfg.train.actor_weight_decay,
+        )
+        self.actor_lr_scheduler = CosineAnnealingWarmupRestarts(
+            self.actor_optimizer,
+            first_cycle_steps=cfg.train.actor_lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=cfg.train.actor_lr,
+            min_lr=cfg.train.actor_lr_scheduler.min_lr,
+            warmup_steps=cfg.train.actor_lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+        self.critic_q_optimizer = torch.optim.AdamW(
+            self.model.critic_q.parameters(),
+            lr=cfg.train.critic_lr,
+            weight_decay=cfg.train.critic_weight_decay,
+        )
+        self.critic_v_optimizer = torch.optim.AdamW(
+            self.model.critic_v.parameters(),
+            lr=cfg.train.critic_lr,
+            weight_decay=cfg.train.critic_weight_decay,
+        )
+        self.critic_v_lr_scheduler = CosineAnnealingWarmupRestarts(
+            self.critic_v_optimizer,
+            first_cycle_steps=cfg.train.critic_lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=cfg.train.critic_lr,
+            min_lr=cfg.train.critic_lr_scheduler.min_lr,
+            warmup_steps=cfg.train.critic_lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+        self.critic_q_lr_scheduler = CosineAnnealingWarmupRestarts(
+            self.critic_q_optimizer,
+            first_cycle_steps=cfg.train.critic_lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=cfg.train.critic_lr,
+            min_lr=cfg.train.critic_lr_scheduler.min_lr,
+            warmup_steps=cfg.train.critic_lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+
+        # Buffer size
+        self.buffer_size = cfg.train.buffer_size
+
+        # Actor params
+        self.use_expectile_exploration = cfg.train.use_expectile_exploration
+
+        # Scaling reward
+        self.scale_reward_factor = cfg.train.scale_reward_factor
+
+        # Updates
+        self.replay_ratio = cfg.train.replay_ratio
+        self.critic_tau = cfg.train.critic_tau
+
+        # Whether to use deterministic mode when sampling at eval
+        self.eval_deterministic = cfg.train.get("eval_deterministic", False)
+
+        # Sampling
+        self.num_sample = cfg.train.eval_sample_num
+
+    def run(self):
+
+        # make a FIFO replay buffer for obs, action, and reward
+        obs_buffer = deque(maxlen=self.buffer_size)
+        action_buffer = deque(maxlen=self.buffer_size)
+        next_obs_buffer = deque(maxlen=self.buffer_size)
+        reward_buffer = deque(maxlen=self.buffer_size)
+        done_buffer = deque(maxlen=self.buffer_size)
+        first_buffer = deque(maxlen=self.buffer_size)
+
+        # Start training loop
+        timer = Timer()
+        run_results = []
+        done_venv = np.zeros((1, self.n_envs))
+        while self.itr < self.n_train_itr:
+
+            # Prepare video paths for each envs --- only applies for the first set of episodes if allowing reset within iteration and each iteration has multiple episodes from one env
+            options_venv = [{} for _ in range(self.n_envs)]
+            if self.itr % self.render_freq == 0 and self.render_video:
+                for env_ind in range(self.n_render):
+                    options_venv[env_ind]["video_path"] = os.path.join(
+                        self.render_dir, f"itr-{self.itr}_trial-{env_ind}.mp4"
+                    )
+
+            # Define train or eval - all envs restart
+            eval_mode = self.itr % self.val_freq == 0 and not self.force_train
+            self.model.eval() if eval_mode else self.model.train()
+            firsts_trajs = np.zeros((self.n_steps + 1, self.n_envs))
+
+            # Reset env at the beginning of an iteration
+            if self.reset_at_iteration or eval_mode or last_itr_eval:
+                prev_obs_venv = self.reset_env_all(options_venv=options_venv)
+                firsts_trajs[0] = 1
+            else:
+                firsts_trajs[0] = (
+                    done_venv  # if done at the end of last iteration, then the envs are just reset
+                )
+            last_itr_eval = eval_mode
+            reward_trajs = np.empty((0, self.n_envs))
+
+            # Collect a set of trajectories from env
+            for step in range(self.n_steps):
+                if step % 10 == 0:
+                    print(f"Processed step {step} of {self.n_steps}")
+
+                # Select action
+                with torch.no_grad():
+                    samples = (
+                        self.model(
+                            cond=torch.from_numpy(prev_obs_venv)
+                            .float()
+                            .to(self.device),
+                            deterministic=eval_mode and self.eval_deterministic,
+                            num_sample=self.num_sample,
+                            use_expectile_exploration=self.use_expectile_exploration,
+                        )
+                        .cpu()
+                        .numpy()
+                    )  # n_env x horizon x act
+                action_venv = samples[:, : self.act_steps]
+
+                # Apply multi-step action
+                obs_venv, reward_venv, done_venv, info_venv = self.venv.step(
+                    action_venv
+                )
+                reward_trajs = np.vstack((reward_trajs, reward_venv[None]))
+
+                # add to buffer
+                obs_buffer.append(prev_obs_venv)
+                action_buffer.append(action_venv)
+                next_obs_buffer.append(obs_venv)
+                reward_buffer.append(reward_venv * self.scale_reward_factor)
+                done_buffer.append(done_venv)
+                first_buffer.append(firsts_trajs[step])
+
+                firsts_trajs[step + 1] = done_venv
+                prev_obs_venv = obs_venv
+
+            # Summarize episode reward --- this needs to be handled differently depending on whether the environment is reset after each iteration. Only count episodes that finish within the iteration.
+            episodes_start_end = []
+            for env_ind in range(self.n_envs):
+                env_steps = np.where(firsts_trajs[:, env_ind] == 1)[0]
+                for i in range(len(env_steps) - 1):
+                    start = env_steps[i]
+                    end = env_steps[i + 1]
+                    if end - start > 1:
+                        episodes_start_end.append((env_ind, start, end - 1))
+            if len(episodes_start_end) > 0:
+                reward_trajs_split = [
+                    reward_trajs[start : end + 1, env_ind]
+                    for env_ind, start, end in episodes_start_end
+                ]
+                num_episode_finished = len(reward_trajs_split)
+                episode_reward = np.array(
+                    [np.sum(reward_traj) for reward_traj in reward_trajs_split]
+                )
+                episode_best_reward = np.array(
+                    [
+                        np.max(reward_traj) / self.act_steps
+                        for reward_traj in reward_trajs_split
+                    ]
+                )
+                avg_episode_reward = np.mean(episode_reward)
+                avg_best_reward = np.mean(episode_best_reward)
+                success_rate = np.mean(
+                    episode_best_reward >= self.best_reward_threshold_for_success
+                )
+            else:
+                episode_reward = np.array([])
+                num_episode_finished = 0
+                avg_episode_reward = 0
+                avg_best_reward = 0
+                success_rate = 0
+                log.info("[WARNING] No episode completed within the iteration!")
+
+            # Update
+            if not eval_mode:
+
+                obs_trajs = np.array(deepcopy(obs_buffer))
+                action_trajs = np.array(deepcopy(action_buffer))
+                next_obs_trajs = np.array(deepcopy(next_obs_buffer))
+                reward_trajs = np.array(deepcopy(reward_buffer))
+                done_trajs = np.array(deepcopy(done_buffer))
+                first_trajs = np.array(deepcopy(first_buffer))
+
+                # flatten
+                obs_trajs = einops.rearrange(
+                    obs_trajs,
+                    "s e h d -> (s e) h d",
+                )
+                next_obs_trajs = einops.rearrange(
+                    next_obs_trajs,
+                    "s e h d -> (s e) h d",
+                )
+                action_trajs = einops.rearrange(
+                    action_trajs,
+                    "s e h d -> (s e) h d",
+                )
+                reward_trajs = reward_trajs.reshape(-1)
+                done_trajs = done_trajs.reshape(-1)
+                first_trajs = first_trajs.reshape(-1)
+
+                num_batch = int(
+                    self.n_steps * self.n_envs / self.batch_size * self.replay_ratio
+                )
+
+                for _ in range(num_batch):
+
+                    # Sample batch
+                    inds = np.random.choice(len(obs_trajs), self.batch_size)
+                    obs_b = torch.from_numpy(obs_trajs[inds]).float().to(self.device)
+                    next_obs_b = (
+                        torch.from_numpy(next_obs_trajs[inds]).float().to(self.device)
+                    )
+                    actions_b = (
+                        torch.from_numpy(action_trajs[inds]).float().to(self.device)
+                    )
+                    reward_b = (
+                        torch.from_numpy(reward_trajs[inds]).float().to(self.device)
+                    )
+                    done_b = torch.from_numpy(done_trajs[inds]).float().to(self.device)
+
+                    # update critic value function
+                    critic_loss_v = self.model.loss_critic_v(obs_b, actions_b)
+                    self.critic_v_optimizer.zero_grad()
+                    critic_loss_v.backward()
+                    self.critic_v_optimizer.step()
+
+                    # update critic q function
+                    critic_loss_q = self.model.loss_critic_q(
+                        obs_b, next_obs_b, actions_b, reward_b, done_b, self.gamma
+                    )
+                    self.critic_q_optimizer.zero_grad()
+                    critic_loss_q.backward()
+                    self.critic_q_optimizer.step()
+
+                    # update target q function
+                    self.model.update_target_critic(self.critic_tau)
+
+                    loss_critic = critic_loss_q.detach() + critic_loss_v.detach()
+
+                    # Update policy with collected trajectories - no weighting
+                    loss = self.model.loss(
+                        actions_b,
+                        obs_b,
+                    )
+                    self.actor_optimizer.zero_grad()
+                    loss.backward()
+                    if self.itr >= self.n_critic_warmup_itr:
+                        if self.max_grad_norm is not None:
+                            torch.nn.utils.clip_grad_norm_(
+                                self.model.actor.parameters(), self.max_grad_norm
+                            )
+                        self.actor_optimizer.step()
+
+            # Update lr
+            self.actor_lr_scheduler.step()
+            self.critic_v_lr_scheduler.step()
+            self.critic_q_lr_scheduler.step()
+
+            # Save model
+            if self.itr % self.save_model_freq == 0 or self.itr == self.n_train_itr - 1:
+                self.save_model()
+
+            # Log loss and save metrics
+            run_results.append(
+                {
+                    "itr": self.itr,
+                }
+            )
+            if self.itr % self.log_freq == 0:
+                if eval_mode:
+                    log.info(
+                        f"eval: success rate {success_rate:8.4f} | avg episode reward {avg_episode_reward:8.4f} | avg best reward {avg_best_reward:8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "success rate - eval": success_rate,
+                                "avg episode reward - eval": avg_episode_reward,
+                                "avg best reward - eval": avg_best_reward,
+                                "num episode - eval": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=False,
+                        )
+                    run_results[-1]["eval_success_rate"] = success_rate
+                    run_results[-1]["eval_episode_reward"] = avg_episode_reward
+                    run_results[-1]["eval_best_reward"] = avg_best_reward
+                else:
+                    log.info(
+                        f"{self.itr}: loss {loss:8.4f} | reward {avg_episode_reward:8.4f} |t:{timer():8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "loss": loss,
+                                "loss - critic": loss_critic,
+                                "avg episode reward - train": avg_episode_reward,
+                                "num episode - train": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=True,
+                        )
+                    run_results[-1]["loss"] = loss
+                    run_results[-1]["loss_critic"] = loss_critic
+                    run_results[-1]["train_episode_reward"] = avg_episode_reward
+                run_results[-1]["time"] = timer()
+                with open(self.result_path, "wb") as f:
+                    pickle.dump(run_results, f)
+            self.itr += 1

agent/finetune/train_ppo_agent.py

@@ -0,0 +1,112 @@
+"""
+Parent PPO fine-tuning agent class.
+
+"""
+
+from typing import Optional
+import torch
+import logging
+from util.scheduler import CosineAnnealingWarmupRestarts
+
+log = logging.getLogger(__name__)
+from agent.finetune.train_agent import TrainAgent
+from util.reward_scaling import RunningRewardScaler
+
+
+class TrainPPOAgent(TrainAgent):
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+        # Batch size for logprobs calculations after an iteration --- prevent out of memory if using a single batch
+        self.logprob_batch_size = cfg.train.get("logprob_batch_size", 10000)
+        assert (
+            self.logprob_batch_size % self.n_envs == 0
+        ), "logprob_batch_size must be divisible by n_envs"
+
+        # note the discount factor gamma here is applied to reward every act_steps, instead of every env step
+        self.gamma = cfg.train.gamma
+
+        # Wwarm up period for critic before actor updates
+        self.n_critic_warmup_itr = cfg.train.n_critic_warmup_itr
+
+        # Optimizer
+        self.actor_optimizer = torch.optim.AdamW(
+            self.model.actor_ft.parameters(),
+            lr=cfg.train.actor_lr,
+            weight_decay=cfg.train.actor_weight_decay,
+        )
+        # use cosine scheduler with linear warmup
+        self.actor_lr_scheduler = CosineAnnealingWarmupRestarts(
+            self.actor_optimizer,
+            first_cycle_steps=cfg.train.actor_lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=cfg.train.actor_lr,
+            min_lr=cfg.train.actor_lr_scheduler.min_lr,
+            warmup_steps=cfg.train.actor_lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+        self.critic_optimizer = torch.optim.AdamW(
+            self.model.critic.parameters(),
+            lr=cfg.train.critic_lr,
+            weight_decay=cfg.train.critic_weight_decay,
+        )
+        self.critic_lr_scheduler = CosineAnnealingWarmupRestarts(
+            self.critic_optimizer,
+            first_cycle_steps=cfg.train.critic_lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=cfg.train.critic_lr,
+            min_lr=cfg.train.critic_lr_scheduler.min_lr,
+            warmup_steps=cfg.train.critic_lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+
+        # Generalized advantage estimation
+        self.gae_lambda: float = cfg.train.get("gae_lambda", 0.95)
+
+        # If specified, stop gradient update once KL difference reaches it
+        self.target_kl: Optional[float] = cfg.train.target_kl
+
+        # Number of times the collected data is used in gradient update
+        self.update_epochs: int = cfg.train.update_epochs
+
+        # Entropy loss coefficient
+        self.ent_coef: float = cfg.train.get("ent_coef", 0)
+
+        # Value loss coefficient
+        self.vf_coef: float = cfg.train.get("vf_coef", 0)
+
+        # Whether to use running reward scaling
+        self.reward_scale_running: bool = cfg.train.reward_scale_running
+        if self.reward_scale_running:
+            self.running_reward_scaler = RunningRewardScaler(self.n_envs)
+
+        # Scaling reward with constant
+        self.reward_scale_const: float = cfg.train.get("reward_scale_const", 1)
+
+        # Use base policy
+        self.use_bc_loss: bool = cfg.train.get("use_bc_loss", False)
+        self.bc_loss_coeff: float = cfg.train.get("bc_loss_coeff", 0)
+
+    def reset_actor_optimizer(self):
+        """Not used anywhere currently"""
+        new_optimizer = torch.optim.AdamW(
+            self.model.actor_ft.parameters(),
+            lr=self.cfg.train.actor_lr,
+            weight_decay=self.cfg.train.actor_weight_decay,
+        )
+        new_optimizer.load_state_dict(self.actor_optimizer.state_dict())
+        self.actor_optimizer = new_optimizer
+
+        new_scheduler = CosineAnnealingWarmupRestarts(
+            self.actor_optimizer,
+            first_cycle_steps=self.cfg.train.actor_lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=self.cfg.train.actor_lr,
+            min_lr=self.cfg.train.actor_lr_scheduler.min_lr,
+            warmup_steps=self.cfg.train.actor_lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+        new_scheduler.load_state_dict(self.actor_lr_scheduler.state_dict())
+        self.actor_lr_scheduler = new_scheduler
+        log.info("Reset actor optimizer")

agent/finetune/train_ppo_diffusion_agent.py

@@ -0,0 +1,452 @@
+"""
+DPPO fine-tuning.
+
+"""
+
+import os
+import pickle
+import einops
+import numpy as np
+import torch
+import logging
+import wandb
+
+log = logging.getLogger(__name__)
+from util.timer import Timer
+from agent.finetune.train_ppo_agent import TrainPPOAgent
+from util.scheduler import CosineAnnealingWarmupRestarts
+
+
+class TrainPPODiffusionAgent(TrainPPOAgent):
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+        # Reward horizon --- always set to act_steps for now
+        self.reward_horizon = cfg.get("reward_horizon", self.act_steps)
+
+        # Eta - between DDIM (=0 for eval) and DDPM (=1 for training)
+        self.learn_eta = self.model.learn_eta
+        if self.learn_eta:
+            self.eta_update_interval = cfg.train.eta_update_interval
+            self.eta_optimizer = torch.optim.AdamW(
+                self.model.eta.parameters(),
+                lr=cfg.train.eta_lr,
+                weight_decay=cfg.train.eta_weight_decay,
+            )
+            self.eta_lr_scheduler = CosineAnnealingWarmupRestarts(
+                self.eta_optimizer,
+                first_cycle_steps=cfg.train.eta_lr_scheduler.first_cycle_steps,
+                cycle_mult=1.0,
+                max_lr=cfg.train.eta_lr,
+                min_lr=cfg.train.eta_lr_scheduler.min_lr,
+                warmup_steps=cfg.train.eta_lr_scheduler.warmup_steps,
+                gamma=1.0,
+            )
+
+    def run(self):
+
+        # Start training loop
+        timer = Timer()
+        run_results = []
+        last_itr_eval = False
+        done_venv = np.zeros((1, self.n_envs))
+        while self.itr < self.n_train_itr:
+
+            # Prepare video paths for each envs --- only applies for the first set of episodes if allowing reset within iteration and each iteration has multiple episodes from one env
+            options_venv = [{} for _ in range(self.n_envs)]
+            if self.itr % self.render_freq == 0 and self.render_video:
+                for env_ind in range(self.n_render):
+                    options_venv[env_ind]["video_path"] = os.path.join(
+                        self.render_dir, f"itr-{self.itr}_trial-{env_ind}.mp4"
+                    )
+
+            # Define train or eval - all envs restart
+            eval_mode = self.itr % self.val_freq == 0 and not self.force_train
+            self.model.eval() if eval_mode else self.model.train()
+            last_itr_eval = eval_mode
+
+            # Reset env before iteration starts (1) if specified, (2) at eval mode, or (3) right after eval mode
+            dones_trajs = np.empty((0, self.n_envs))
+            firsts_trajs = np.zeros((self.n_steps + 1, self.n_envs))
+            if self.reset_at_iteration or eval_mode or last_itr_eval:
+                prev_obs_venv = self.reset_env_all(options_venv=options_venv)
+                firsts_trajs[0] = 1
+            else:
+                firsts_trajs[0] = (
+                    done_venv  # if done at the end of last iteration, then the envs are just reset
+                )
+
+            # Holder
+            obs_trajs = np.empty((0, self.n_envs, self.n_cond_step, self.obs_dim))
+            chains_trajs = np.empty(
+                (
+                    0,
+                    self.n_envs,
+                    self.model.ft_denoising_steps + 1,
+                    self.horizon_steps,
+                    self.action_dim,
+                )
+            )
+            reward_trajs = np.empty((0, self.n_envs))
+            obs_full_trajs = np.empty((0, self.n_envs, self.obs_dim))
+            obs_full_trajs = np.vstack(
+                (obs_full_trajs, prev_obs_venv[None].squeeze(2))
+            )  # remove cond_step dim
+
+            # Collect a set of trajectories from env
+            for step in range(self.n_steps):
+                if step % 10 == 0:
+                    print(f"Processed step {step} of {self.n_steps}")
+
+                # Select action
+                with torch.no_grad():
+                    samples = self.model(
+                        cond=torch.from_numpy(prev_obs_venv).float().to(self.device),
+                        deterministic=eval_mode,
+                        return_chain=True,
+                    )
+                    output_venv = (
+                        samples.trajectories.cpu().numpy()
+                    )  # n_env x horizon x act
+                    chains_venv = (
+                        samples.chains.cpu().numpy()
+                    )  # n_env x denoising x horizon x act
+                action_venv = output_venv[:, : self.act_steps]
+
+                # Apply multi-step action
+                obs_venv, reward_venv, done_venv, info_venv = self.venv.step(
+                    action_venv
+                )
+                if self.save_full_observations:
+                    obs_full_venv = np.vstack(
+                        [info["full_obs"][None] for info in info_venv]
+                    )  # n_envs x n_act_steps x obs_dim
+                    obs_full_trajs = np.vstack(
+                        (obs_full_trajs, obs_full_venv.transpose(1, 0, 2))
+                    )
+                obs_trajs = np.vstack((obs_trajs, prev_obs_venv[None]))
+                chains_trajs = np.vstack((chains_trajs, chains_venv[None]))
+                reward_trajs = np.vstack((reward_trajs, reward_venv[None]))
+                dones_trajs = np.vstack((dones_trajs, done_venv[None]))
+                firsts_trajs[step + 1] = done_venv
+                prev_obs_venv = obs_venv
+
+            # Summarize episode reward --- this needs to be handled differently depending on whether the environment is reset after each iteration. Only count episodes that finish within the iteration.
+            episodes_start_end = []
+            for env_ind in range(self.n_envs):
+                env_steps = np.where(firsts_trajs[:, env_ind] == 1)[0]
+                for i in range(len(env_steps) - 1):
+                    start = env_steps[i]
+                    end = env_steps[i + 1]
+                    if end - start > 1:
+                        episodes_start_end.append((env_ind, start, end - 1))
+            if len(episodes_start_end) > 0:
+                reward_trajs_split = [
+                    reward_trajs[start : end + 1, env_ind]
+                    for env_ind, start, end in episodes_start_end
+                ]
+                num_episode_finished = len(reward_trajs_split)
+                episode_reward = np.array(
+                    [np.sum(reward_traj) for reward_traj in reward_trajs_split]
+                )
+                if (
+                    self.furniture_sparse_reward
+                ):  # only for furniture tasks, where reward only occurs in one env step
+                    episode_best_reward = episode_reward
+                else:
+                    episode_best_reward = np.array(
+                        [
+                            np.max(reward_traj) / self.act_steps
+                            for reward_traj in reward_trajs_split
+                        ]
+                    )
+                avg_episode_reward = np.mean(episode_reward)
+                avg_best_reward = np.mean(episode_best_reward)
+                success_rate = np.mean(
+                    episode_best_reward >= self.best_reward_threshold_for_success
+                )
+            else:
+                episode_reward = np.array([])
+                num_episode_finished = 0
+                avg_episode_reward = 0
+                avg_best_reward = 0
+                success_rate = 0
+                log.info("[WARNING] No episode completed within the iteration!")
+
+            # Update models
+            if not eval_mode:
+                with torch.no_grad():
+                    # Calculate value and logprobs - split into batches to prevent out of memory
+                    obs_t = einops.rearrange(
+                        torch.from_numpy(obs_trajs).float().to(self.device),
+                        "s e h d -> (s e) h d",
+                    )
+                    obs_ts = torch.split(obs_t, self.logprob_batch_size, dim=0)
+                    values_trajs = np.empty((0, self.n_envs))
+                    for obs in obs_ts:
+                        values = self.model.critic(obs).cpu().numpy().flatten()
+                        values_trajs = np.vstack(
+                            (values_trajs, values.reshape(-1, self.n_envs))
+                        )
+                    chains_t = einops.rearrange(
+                        torch.from_numpy(chains_trajs).float().to(self.device),
+                        "s e t h d -> (s e) t h d",
+                    )
+                    chains_ts = torch.split(chains_t, self.logprob_batch_size, dim=0)
+                    logprobs_trajs = np.empty(
+                        (
+                            0,
+                            self.model.ft_denoising_steps,
+                            self.horizon_steps,
+                            self.action_dim,
+                        )
+                    )
+                    for obs, chains in zip(obs_ts, chains_ts):
+                        logprobs = self.model.get_logprobs(obs, chains).cpu().numpy()
+                        logprobs_trajs = np.vstack(
+                            (
+                                logprobs_trajs,
+                                logprobs.reshape(-1, *logprobs_trajs.shape[1:]),
+                            )
+                        )
+
+                    # normalize reward with running variance if specified
+                    if self.reward_scale_running:
+                        reward_trajs_transpose = self.running_reward_scaler(
+                            reward=reward_trajs.T, first=firsts_trajs[:-1].T
+                        )
+                        reward_trajs = reward_trajs_transpose.T
+
+                    # bootstrap value with GAE if not done - apply reward scaling with constant if specified
+                    obs_venv_ts = torch.from_numpy(obs_venv).float().to(self.device)
+                    with torch.no_grad():
+                        next_value = (
+                            self.model.critic(obs_venv_ts).reshape(1, -1).cpu().numpy()
+                        )
+                    advantages_trajs = np.zeros_like(reward_trajs)
+                    lastgaelam = 0
+                    for t in reversed(range(self.n_steps)):
+                        if t == self.n_steps - 1:
+                            nextnonterminal = 1.0 - done_venv
+                            nextvalues = next_value
+                        else:
+                            nextnonterminal = 1.0 - dones_trajs[t + 1]
+                            nextvalues = values_trajs[t + 1]
+                        # delta = r + gamma*V(st+1) - V(st)
+                        delta = (
+                            reward_trajs[t] * self.reward_scale_const
+                            + self.gamma * nextvalues * nextnonterminal
+                            - values_trajs[t]
+                        )
+                        # A = delta_t + gamma*lamdba*delta_{t+1} + ...
+                        advantages_trajs[t] = lastgaelam = (
+                            delta
+                            + self.gamma
+                            * self.gae_lambda
+                            * nextnonterminal
+                            * lastgaelam
+                        )
+                    returns_trajs = advantages_trajs + values_trajs
+
+                # k for environment step
+                obs_k = einops.rearrange(
+                    torch.tensor(obs_trajs).float().to(self.device),
+                    "s e h d -> (s e) h d",
+                )
+                chains_k = einops.rearrange(
+                    torch.tensor(chains_trajs).float().to(self.device),
+                    "s e t h d -> (s e) t h d",
+                )
+                returns_k = (
+                    torch.tensor(returns_trajs).float().to(self.device).reshape(-1)
+                )
+                values_k = (
+                    torch.tensor(values_trajs).float().to(self.device).reshape(-1)
+                )
+                advantages_k = (
+                    torch.tensor(advantages_trajs).float().to(self.device).reshape(-1)
+                )
+                logprobs_k = torch.tensor(logprobs_trajs).float().to(self.device)
+
+                # Update policy and critic
+                total_steps = self.n_steps * self.n_envs
+                inds_k = np.arange(total_steps)
+                clipfracs = []
+                for update_epoch in range(self.update_epochs):
+
+                    # for each epoch, go through all data in batches
+                    flag_break = False
+                    np.random.shuffle(inds_k)
+                    num_batch = max(1, total_steps // self.batch_size)  # skip last ones
+                    for batch in range(num_batch):
+                        start = batch * self.batch_size
+                        end = start + self.batch_size
+                        inds_b = inds_k[start:end]  # b for batch
+                        obs_b = obs_k[inds_b]
+                        chains_b = chains_k[inds_b]
+                        returns_b = returns_k[inds_b]
+                        values_b = values_k[inds_b]
+                        advantages_b = advantages_k[inds_b]
+                        logprobs_b = logprobs_k[inds_b]
+
+                        # get loss
+                        (
+                            pg_loss,
+                            entropy_loss,
+                            v_loss,
+                            clipfrac,
+                            approx_kl,
+                            ratio,
+                            bc_loss,
+                            eta,
+                        ) = self.model.loss(
+                            obs_b,
+                            chains_b,
+                            returns_b,
+                            values_b,
+                            advantages_b,
+                            logprobs_b,
+                            use_bc_loss=self.use_bc_loss,
+                            reward_horizon=self.reward_horizon,
+                        )
+                        loss = (
+                            pg_loss
+                            + entropy_loss * self.ent_coef
+                            + v_loss * self.vf_coef
+                            + bc_loss * self.bc_loss_coeff
+                        )
+                        clipfracs += [clipfrac]
+
+                        # update policy and critic
+                        self.actor_optimizer.zero_grad()
+                        self.critic_optimizer.zero_grad()
+                        if self.learn_eta:
+                            self.eta_optimizer.zero_grad()
+                        loss.backward()
+                        if self.itr >= self.n_critic_warmup_itr:
+                            if self.max_grad_norm is not None:
+                                torch.nn.utils.clip_grad_norm_(
+                                    self.model.actor_ft.parameters(), self.max_grad_norm
+                                )
+                            self.actor_optimizer.step()
+                            if self.learn_eta and batch % self.eta_update_interval == 0:
+                                self.eta_optimizer.step()
+                        self.critic_optimizer.step()
+                        log.info(
+                            f"approx_kl: {approx_kl}, update_epoch: {update_epoch}, num_batch: {num_batch}"
+                        )
+
+                        # Stop gradient update if KL difference reaches target
+                        if self.target_kl is not None and approx_kl > self.target_kl:
+                            flag_break = True
+                            break
+                    if flag_break:
+                        break
+
+                # Explained variation of future rewards using value function
+                y_pred, y_true = values_k.cpu().numpy(), returns_k.cpu().numpy()
+                var_y = np.var(y_true)
+                explained_var = (
+                    np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y
+                )
+
+            # Plot state trajectories in D3IL
+            if (
+                self.itr % self.render_freq == 0
+                and self.n_render > 0
+                and self.traj_plotter is not None
+            ):
+                self.traj_plotter(
+                    obs_full_trajs=obs_full_trajs,
+                    n_render=self.n_render,
+                    max_episode_steps=self.max_episode_steps,
+                    render_dir=self.render_dir,
+                    itr=self.itr,
+                )
+
+            # Update lr, min_sampling_std
+            if self.itr >= self.n_critic_warmup_itr:
+                self.actor_lr_scheduler.step()
+                if self.learn_eta:
+                    self.eta_lr_scheduler.step()
+            self.critic_lr_scheduler.step()
+            self.model.step()
+            diffusion_min_sampling_std = self.model.get_min_sampling_denoising_std()
+
+            # Save model
+            if self.itr % self.save_model_freq == 0 or self.itr == self.n_train_itr - 1:
+                self.save_model()
+
+            # Log loss and save metrics
+            run_results.append(
+                {
+                    "itr": self.itr,
+                }
+            )
+            if self.save_trajs:
+                run_results[-1]["obs_full_trajs"] = obs_full_trajs
+                run_results[-1]["obs_trajs"] = obs_trajs
+                run_results[-1]["chains_trajs"] = chains_trajs
+                run_results[-1]["reward_trajs"] = reward_trajs
+            if self.itr % self.log_freq == 0:
+                time = timer()
+                if eval_mode:
+                    log.info(
+                        f"eval: success rate {success_rate:8.4f} | avg episode reward {avg_episode_reward:8.4f} | avg best reward {avg_best_reward:8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "success rate - eval": success_rate,
+                                "avg episode reward - eval": avg_episode_reward,
+                                "avg best reward - eval": avg_best_reward,
+                                "num episode - eval": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=False,
+                        )
+                    run_results[-1]["eval_success_rate"] = success_rate
+                    run_results[-1]["eval_episode_reward"] = avg_episode_reward
+                    run_results[-1]["eval_best_reward"] = avg_best_reward
+                else:
+                    log.info(
+                        f"{self.itr}: loss {loss:8.4f} | pg loss {pg_loss:8.4f} | value loss {v_loss:8.4f} | bc loss {bc_loss:8.4f} | reward {avg_episode_reward:8.4f} | eta {eta:8.4f} | t:{time:8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "loss": loss,
+                                "pg loss": pg_loss,
+                                "value loss": v_loss,
+                                "bc loss": bc_loss,
+                                "eta": eta,
+                                "approx kl": approx_kl,
+                                "ratio": ratio,
+                                "clipfrac": np.mean(clipfracs),
+                                "explained variance": explained_var,
+                                "avg episode reward - train": avg_episode_reward,
+                                "num episode - train": num_episode_finished,
+                                "diffusion - min sampling std": diffusion_min_sampling_std,
+                                "actor lr": self.actor_optimizer.param_groups[0]["lr"],
+                                "critic lr": self.critic_optimizer.param_groups[0][
+                                    "lr"
+                                ],
+                            },
+                            step=self.itr,
+                            commit=True,
+                        )
+                    run_results[-1]["loss"] = loss
+                    run_results[-1]["pg_loss"] = pg_loss
+                    run_results[-1]["value_loss"] = v_loss
+                    run_results[-1]["bc_loss"] = bc_loss
+                    run_results[-1]["eta"] = eta
+                    run_results[-1]["approx_kl"] = approx_kl
+                    run_results[-1]["ratio"] = ratio
+                    run_results[-1]["clip_frac"] = np.mean(clipfracs)
+                    run_results[-1]["explained_variance"] = explained_var
+                    run_results[-1]["train_episode_reward"] = avg_episode_reward
+                run_results[-1]["time"] = time
+                with open(self.result_path, "wb") as f:
+                    pickle.dump(run_results, f)
+            self.itr += 1

agent/finetune/train_ppo_diffusion_img_agent.py

@@ -0,0 +1,468 @@
+"""
+DPPO fine-tuning for pixel observations.
+
+"""
+
+import os
+import pickle
+import einops
+import numpy as np
+import torch
+import logging
+import wandb
+import math
+
+log = logging.getLogger(__name__)
+from util.timer import Timer
+from agent.finetune.train_ppo_diffusion_agent import TrainPPODiffusionAgent
+from model.common.modules import RandomShiftsAug
+
+
+class TrainPPOImgDiffusionAgent(TrainPPODiffusionAgent):
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+        # Image randomization
+        self.augment = cfg.train.augment
+        if self.augment:
+            self.aug = RandomShiftsAug(pad=4)
+
+        # Set obs dim -  we will save the different obs in batch in a dict
+        shape_meta = cfg.shape_meta
+        self.obs_dims = {k: shape_meta.obs[k]["shape"] for k in shape_meta.obs.keys()}
+
+        # Gradient accumulation to deal with large GPU RAM usage
+        self.grad_accumulate = cfg.train.grad_accumulate
+
+    def run(self):
+
+        # Start training loop
+        timer = Timer()
+        run_results = []
+        last_itr_eval = False
+        done_venv = np.zeros((1, self.n_envs))
+        while self.itr < self.n_train_itr:
+
+            # Prepare video paths for each envs --- only applies for the first set of episodes if allowing reset within iteration and each iteration has multiple episodes from one env
+            options_venv = [{} for _ in range(self.n_envs)]
+            if self.itr % self.render_freq == 0 and self.render_video:
+                for env_ind in range(self.n_render):
+                    options_venv[env_ind]["video_path"] = os.path.join(
+                        self.render_dir, f"itr-{self.itr}_trial-{env_ind}.mp4"
+                    )
+
+            # Define train or eval - all envs restart
+            eval_mode = self.itr % self.val_freq == 0 and not self.force_train
+            self.model.eval() if eval_mode else self.model.train()
+            last_itr_eval = eval_mode
+
+            # Reset env before iteration starts (1) if specified, (2) at eval mode, or (3) right after eval mode
+            dones_trajs = np.empty((0, self.n_envs))
+            firsts_trajs = np.zeros((self.n_steps + 1, self.n_envs))
+            if self.reset_at_iteration or eval_mode or last_itr_eval:
+                prev_obs_venv = self.reset_env_all(options_venv=options_venv)
+                firsts_trajs[0] = 1
+            else:
+                firsts_trajs[0] = (
+                    done_venv  # if done at the end of last iteration, then the envs are just reset
+                )
+
+            # Holder
+            obs_trajs = {
+                k: np.empty((0, self.n_envs, self.n_cond_step, *self.obs_dims[k]))
+                for k in self.obs_dims
+            }
+            chains_trajs = np.empty(
+                (
+                    0,
+                    self.n_envs,
+                    self.model.ft_denoising_steps + 1,
+                    self.horizon_steps,
+                    self.action_dim,
+                )
+            )
+            reward_trajs = np.empty((0, self.n_envs))
+
+            # Collect a set of trajectories from env
+            for step in range(self.n_steps):
+                if step % 10 == 0:
+                    print(f"Processed step {step} of {self.n_steps}")
+
+                # Select action
+                with torch.no_grad():
+                    cond = {
+                        key: torch.from_numpy(prev_obs_venv[key])
+                        .float()
+                        .to(self.device)
+                        for key in self.obs_dims.keys()
+                    }  # batch each type of obs and put into dict
+                    samples = self.model(
+                        cond=cond,
+                        deterministic=eval_mode,
+                        return_chain=True,
+                    )
+                    output_venv = (
+                        samples.trajectories.cpu().numpy()
+                    )  # n_env x horizon x act
+                    chains_venv = (
+                        samples.chains.cpu().numpy()
+                    )  # n_env x denoising x horizon x act
+                action_venv = output_venv[:, : self.act_steps]
+
+                # Apply multi-step action
+                obs_venv, reward_venv, done_venv, info_venv = self.venv.step(
+                    action_venv
+                )
+                for k in obs_trajs.keys():
+                    obs_trajs[k] = np.vstack((obs_trajs[k], prev_obs_venv[k][None]))
+                chains_trajs = np.vstack((chains_trajs, chains_venv[None]))
+                reward_trajs = np.vstack((reward_trajs, reward_venv[None]))
+                dones_trajs = np.vstack((dones_trajs, done_venv[None]))
+                firsts_trajs[step + 1] = done_venv
+                prev_obs_venv = obs_venv
+
+            # Summarize episode reward --- this needs to be handled differently depending on whether the environment is reset after each iteration. Only count episodes that finish within the iteration.
+            episodes_start_end = []
+            for env_ind in range(self.n_envs):
+                env_steps = np.where(firsts_trajs[:, env_ind] == 1)[0]
+                for i in range(len(env_steps) - 1):
+                    start = env_steps[i]
+                    end = env_steps[i + 1]
+                    if end - start > 1:
+                        episodes_start_end.append((env_ind, start, end - 1))
+            if len(episodes_start_end) > 0:
+                reward_trajs_split = [
+                    reward_trajs[start : end + 1, env_ind]
+                    for env_ind, start, end in episodes_start_end
+                ]
+                num_episode_finished = len(reward_trajs_split)
+                episode_reward = np.array(
+                    [np.sum(reward_traj) for reward_traj in reward_trajs_split]
+                )
+                episode_best_reward = np.array(
+                    [
+                        np.max(reward_traj) / self.act_steps
+                        for reward_traj in reward_trajs_split
+                    ]
+                )
+                avg_episode_reward = np.mean(episode_reward)
+                avg_best_reward = np.mean(episode_best_reward)
+                success_rate = np.mean(
+                    episode_best_reward >= self.best_reward_threshold_for_success
+                )
+            else:
+                episode_reward = np.array([])
+                num_episode_finished = 0
+                avg_episode_reward = 0
+                avg_best_reward = 0
+                success_rate = 0
+                log.info("[WARNING] No episode completed within the iteration!")
+
+            # Update
+            if not eval_mode:
+                with torch.no_grad():
+                    # apply image randomization
+                    obs_trajs["rgb"] = (
+                        torch.from_numpy(obs_trajs["rgb"]).float().to(self.device)
+                    )
+                    obs_trajs["state"] = (
+                        torch.from_numpy(obs_trajs["state"]).float().to(self.device)
+                    )
+                    if self.augment:
+                        rgb = einops.rearrange(
+                            obs_trajs["rgb"],
+                            "s e t c h w -> (s e t) c h w",
+                        )
+                        rgb = self.aug(rgb)
+                        obs_trajs["rgb"] = einops.rearrange(
+                            rgb,
+                            "(s e t) c h w -> s e t c h w",
+                            s=self.n_steps,
+                            e=self.n_envs,
+                        )
+
+                    # Calculate value and logprobs - split into batches to prevent out of memory
+                    num_split = math.ceil(
+                        self.n_envs * self.n_steps / self.logprob_batch_size
+                    )
+                    obs_ts = [{} for _ in range(num_split)]
+                    for k in obs_trajs.keys():
+                        obs_k = einops.rearrange(
+                            obs_trajs[k],
+                            "s e ... -> (s e) ...",
+                        )
+                        obs_ts_k = torch.split(obs_k, self.logprob_batch_size, dim=0)
+                        for i, obs_t in enumerate(obs_ts_k):
+                            obs_ts[i][k] = obs_t
+                    values_trajs = np.empty((0, self.n_envs))
+                    for obs in obs_ts:
+                        values = (
+                            self.model.critic(obs, no_augment=True)
+                            .cpu()
+                            .numpy()
+                            .flatten()
+                        )
+                        values_trajs = np.vstack(
+                            (values_trajs, values.reshape(-1, self.n_envs))
+                        )
+                    chains_t = einops.rearrange(
+                        torch.from_numpy(chains_trajs).float().to(self.device),
+                        "s e t h d -> (s e) t h d",
+                    )
+                    chains_ts = torch.split(chains_t, self.logprob_batch_size, dim=0)
+                    logprobs_trajs = np.empty(
+                        (
+                            0,
+                            self.model.ft_denoising_steps,
+                            self.horizon_steps,
+                            self.action_dim,
+                        )
+                    )
+                    for obs, chains in zip(obs_ts, chains_ts):
+                        logprobs = self.model.get_logprobs(obs, chains).cpu().numpy()
+                        logprobs_trajs = np.vstack(
+                            (
+                                logprobs_trajs,
+                                logprobs.reshape(-1, *logprobs_trajs.shape[1:]),
+                            )
+                        )
+
+                    # normalize reward with running variance if specified
+                    if self.reward_scale_running:
+                        reward_trajs_transpose = self.running_reward_scaler(
+                            reward=reward_trajs.T, first=firsts_trajs[:-1].T
+                        )
+                        reward_trajs = reward_trajs_transpose.T
+
+                    # bootstrap value with GAE if not done - apply reward scaling with constant if specified
+                    obs_venv_ts = {
+                        key: torch.from_numpy(obs_venv[key]).float().to(self.device)
+                        for key in self.obs_dims.keys()
+                    }
+                    with torch.no_grad():
+                        next_value = (
+                            self.model.critic(obs_venv_ts, no_augment=True)
+                            .reshape(1, -1)
+                            .cpu()
+                            .numpy()
+                        )
+                    advantages_trajs = np.zeros_like(reward_trajs)
+                    lastgaelam = 0
+                    for t in reversed(range(self.n_steps)):
+                        if t == self.n_steps - 1:
+                            nextnonterminal = 1.0 - done_venv
+                            nextvalues = next_value
+                        else:
+                            nextnonterminal = 1.0 - dones_trajs[t + 1]
+                            nextvalues = values_trajs[t + 1]
+                        # delta = r + gamma*V(st+1) - V(st)
+                        delta = (
+                            reward_trajs[t] * self.reward_scale_const
+                            + self.gamma * nextvalues * nextnonterminal
+                            - values_trajs[t]
+                        )
+                        # A = delta_t + gamma*lamdba*delta_{t+1} + ...
+                        advantages_trajs[t] = lastgaelam = (
+                            delta
+                            + self.gamma
+                            * self.gae_lambda
+                            * nextnonterminal
+                            * lastgaelam
+                        )
+                    returns_trajs = advantages_trajs + values_trajs
+
+                # k for environment step
+                obs_k = {
+                    k: einops.rearrange(
+                        obs_trajs[k],
+                        "s e ... -> (s e) ...",
+                    )
+                    for k in obs_trajs.keys()
+                }
+                chains_k = einops.rearrange(
+                    torch.tensor(chains_trajs).float().to(self.device),
+                    "s e t h d -> (s e) t h d",
+                )
+                returns_k = (
+                    torch.tensor(returns_trajs).float().to(self.device).reshape(-1)
+                )
+                values_k = (
+                    torch.tensor(values_trajs).float().to(self.device).reshape(-1)
+                )
+                advantages_k = (
+                    torch.tensor(advantages_trajs).float().to(self.device).reshape(-1)
+                )
+                logprobs_k = torch.tensor(logprobs_trajs).float().to(self.device)
+
+                # Update policy and critic
+                total_steps = self.n_steps * self.n_envs
+                inds_k = np.arange(total_steps)
+                clipfracs = []
+                for update_epoch in range(self.update_epochs):
+
+                    # for each epoch, go through all data in batches
+                    flag_break = False
+                    np.random.shuffle(inds_k)
+                    num_batch = max(1, total_steps // self.batch_size)  # skip last ones
+                    for batch in range(num_batch):
+                        start = batch * self.batch_size
+                        end = start + self.batch_size
+                        inds_b = inds_k[start:end]  # b for batch
+                        obs_b = {k: obs_k[k][inds_b] for k in obs_k.keys()}
+                        chains_b = chains_k[inds_b]
+                        returns_b = returns_k[inds_b]
+                        values_b = values_k[inds_b]
+                        advantages_b = advantages_k[inds_b]
+                        logprobs_b = logprobs_k[inds_b]
+
+                        # get loss
+                        (
+                            pg_loss,
+                            entropy_loss,
+                            v_loss,
+                            clipfrac,
+                            approx_kl,
+                            ratio,
+                            bc_loss,
+                            eta,
+                        ) = self.model.loss(
+                            obs_b,
+                            chains_b,
+                            returns_b,
+                            values_b,
+                            advantages_b,
+                            logprobs_b,
+                            use_bc_loss=self.use_bc_loss,
+                            reward_horizon=self.reward_horizon,
+                        )
+                        loss = (
+                            pg_loss
+                            + entropy_loss * self.ent_coef
+                            + v_loss * self.vf_coef
+                            + bc_loss * self.bc_loss_coeff
+                        )
+                        clipfracs += [clipfrac]
+
+                        # update policy and critic
+                        loss.backward()
+                        if (batch + 1) % self.grad_accumulate == 0:
+                            if self.itr >= self.n_critic_warmup_itr:
+                                if self.max_grad_norm is not None:
+                                    torch.nn.utils.clip_grad_norm_(
+                                        self.model.actor_ft.parameters(),
+                                        self.max_grad_norm,
+                                    )
+                                self.actor_optimizer.step()
+                                if (
+                                    self.learn_eta
+                                    and batch % self.eta_update_interval == 0
+                                ):
+                                    self.eta_optimizer.step()
+                            self.critic_optimizer.step()
+                            self.actor_optimizer.zero_grad()
+                            self.critic_optimizer.zero_grad()
+                            if self.learn_eta:
+                                self.eta_optimizer.zero_grad()
+                            log.info(f"run grad update at batch {batch}")
+                            log.info(
+                                f"approx_kl: {approx_kl}, update_epoch: {update_epoch}, num_batch: {num_batch}"
+                            )
+
+                            # Stop gradient update if KL difference reaches target
+                            if (
+                                self.target_kl is not None
+                                and approx_kl > self.target_kl
+                                and self.itr >= self.n_critic_warmup_itr
+                            ):
+                                flag_break = True
+                                break
+                    if flag_break:
+                        break
+
+                # Explained variation of future rewards using value function
+                y_pred, y_true = values_k.cpu().numpy(), returns_k.cpu().numpy()
+                var_y = np.var(y_true)
+                explained_var = (
+                    np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y
+                )
+
+            # Update lr
+            if self.itr >= self.n_critic_warmup_itr:
+                self.actor_lr_scheduler.step()
+                if self.learn_eta:
+                    self.eta_lr_scheduler.step()
+            self.critic_lr_scheduler.step()
+            self.model.step()
+            diffusion_min_sampling_std = self.model.get_min_sampling_denoising_std()
+
+            # Save model
+            if self.itr % self.save_model_freq == 0 or self.itr == self.n_train_itr - 1:
+                self.save_model()
+
+            # Log loss and save metrics
+            run_results.append(
+                {
+                    "itr": self.itr,
+                }
+            )
+            if self.itr % self.log_freq == 0:
+                if eval_mode:
+                    log.info(
+                        f"eval: success rate {success_rate:8.4f} | avg episode reward {avg_episode_reward:8.4f} | avg best reward {avg_best_reward:8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "success rate - eval": success_rate,
+                                "avg episode reward - eval": avg_episode_reward,
+                                "avg best reward - eval": avg_best_reward,
+                                "num episode - eval": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=False,
+                        )
+                    run_results[-1]["eval_success_rate"] = success_rate
+                    run_results[-1]["eval_episode_reward"] = avg_episode_reward
+                    run_results[-1]["eval_best_reward"] = avg_best_reward
+                else:
+                    log.info(
+                        f"{self.itr}: loss {loss:8.4f} | pg loss {pg_loss:8.4f} | value loss {v_loss:8.4f} | bc loss {bc_loss:8.4f} | reward {avg_episode_reward:8.4f} | eta {eta:8.4f} | t:{timer():8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "loss": loss,
+                                "pg loss": pg_loss,
+                                "value loss": v_loss,
+                                "bc loss": bc_loss,
+                                "eta": eta,
+                                "approx kl": approx_kl,
+                                "ratio": ratio,
+                                "clipfrac": np.mean(clipfracs),
+                                "explained variance": explained_var,
+                                "avg episode reward - train": avg_episode_reward,
+                                "num episode - train": num_episode_finished,
+                                "diffusion - min sampling std": diffusion_min_sampling_std,
+                                "actor lr": self.actor_optimizer.param_groups[0]["lr"],
+                                "critic lr": self.critic_optimizer.param_groups[0][
+                                    "lr"
+                                ],
+                            },
+                            step=self.itr,
+                            commit=True,
+                        )
+                    run_results[-1]["loss"] = loss
+                    run_results[-1]["pg_loss"] = pg_loss
+                    run_results[-1]["value_loss"] = v_loss
+                    run_results[-1]["bc_loss"] = bc_loss
+                    run_results[-1]["eta"] = eta
+                    run_results[-1]["approx_kl"] = approx_kl
+                    run_results[-1]["ratio"] = ratio
+                    run_results[-1]["clip_frac"] = np.mean(clipfracs)
+                    run_results[-1]["explained_variance"] = explained_var
+                    run_results[-1]["train_episode_reward"] = avg_episode_reward
+                run_results[-1]["time"] = timer()
+                with open(self.result_path, "wb") as f:
+                    pickle.dump(run_results, f)
+            self.itr += 1

agent/finetune/train_ppo_exact_diffusion_agent.py

@@ -0,0 +1,405 @@
+"""
+Use diffusion exact likelihood for policy gradient.
+
+"""
+
+import os
+import pickle
+import einops
+import numpy as np
+import torch
+import logging
+import wandb
+
+log = logging.getLogger(__name__)
+from util.timer import Timer
+from agent.finetune.train_ppo_diffusion_agent import TrainPPODiffusionAgent
+
+
+class TrainPPOExactDiffusionAgent(TrainPPODiffusionAgent):
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+    def run(self):
+        """
+        For exact likelihood, we do not need to save the chains.
+        """
+
+        # Start training loop
+        timer = Timer()
+        run_results = []
+        last_itr_eval = False
+        done_venv = np.zeros((1, self.n_envs))
+        while self.itr < self.n_train_itr:
+
+            # Prepare video paths for each envs --- only applies for the first set of episodes if allowing reset within iteration and each iteration has multiple episodes from one env
+            options_venv = [{} for _ in range(self.n_envs)]
+            if self.itr % self.render_freq == 0 and self.render_video:
+                for env_ind in range(self.n_render):
+                    options_venv[env_ind]["video_path"] = os.path.join(
+                        self.render_dir, f"itr-{self.itr}_trial-{env_ind}.mp4"
+                    )
+
+            # Define train or eval - all envs restart
+            eval_mode = self.itr % self.val_freq == 0 and not self.force_train
+            self.model.eval() if eval_mode else self.model.train()
+            last_itr_eval = eval_mode
+
+            # Reset env before iteration starts (1) if specified, (2) at eval mode, or (3) right after eval mode
+            dones_trajs = np.empty((0, self.n_envs))
+            firsts_trajs = np.zeros((self.n_steps + 1, self.n_envs))
+            if self.reset_at_iteration or eval_mode or last_itr_eval:
+                prev_obs_venv = self.reset_env_all(options_venv=options_venv)
+                firsts_trajs[0] = 1
+            else:
+                firsts_trajs[0] = (
+                    done_venv  # if done at the end of last iteration, then the envs are just reset
+                )
+
+            # Holder
+            obs_trajs = np.empty((0, self.n_envs, self.n_cond_step, self.obs_dim))
+            samples_trajs = np.empty(
+                (
+                    0,
+                    self.n_envs,
+                    self.horizon_steps,
+                    self.action_dim,
+                )
+            )
+            chains_trajs = np.empty(
+                (
+                    0,
+                    self.n_envs,
+                    self.model.ft_denoising_steps + 1,
+                    self.horizon_steps,
+                    self.action_dim,
+                )
+            )
+            reward_trajs = np.empty((0, self.n_envs))
+            obs_full_trajs = np.empty((0, self.n_envs, self.obs_dim))
+            obs_full_trajs = np.vstack(
+                (obs_full_trajs, prev_obs_venv[None].squeeze(2))
+            )  # remove cond_step dim
+
+            # Collect a set of trajectories from env
+            for step in range(self.n_steps):
+                if step % 10 == 0:
+                    print(f"Processed step {step} of {self.n_steps}")
+
+                # Select action
+                with torch.no_grad():
+                    samples = self.model(
+                        cond=torch.from_numpy(prev_obs_venv).float().to(self.device),
+                        deterministic=eval_mode,
+                        return_chain=True,
+                    )
+                    output_venv = (
+                        samples.trajectories.cpu().numpy()
+                    )  # n_env x horizon x act
+                    chains_venv = (
+                        samples.chains.cpu().numpy()
+                    )  # n_env x denoising x horizon x act
+                action_venv = output_venv[:, : self.act_steps]
+
+                # Apply multi-step action
+                obs_venv, reward_venv, done_venv, info_venv = self.venv.step(
+                    action_venv
+                )
+                if self.save_full_observations:
+                    obs_full_venv = np.vstack(
+                        [info["full_obs"][None] for info in info_venv]
+                    )  # n_envs x n_act_steps x obs_dim
+                    obs_full_trajs = np.vstack(
+                        (obs_full_trajs, obs_full_venv.transpose(1, 0, 2))
+                    )
+                obs_trajs = np.vstack((obs_trajs, prev_obs_venv[None]))
+                chains_trajs = np.vstack((chains_trajs, chains_venv[None]))
+                samples_trajs = np.vstack((samples_trajs, output_venv[None]))
+                reward_trajs = np.vstack((reward_trajs, reward_venv[None]))
+                dones_trajs = np.vstack((dones_trajs, done_venv[None]))
+                firsts_trajs[step + 1] = done_venv
+                prev_obs_venv = obs_venv
+
+            # Summarize episode reward --- this needs to be handled differently depending on whether the environment is reset after each iteration. Only count episodes that finish within the iteration.
+            episodes_start_end = []
+            for env_ind in range(self.n_envs):
+                env_steps = np.where(firsts_trajs[:, env_ind] == 1)[0]
+                for i in range(len(env_steps) - 1):
+                    start = env_steps[i]
+                    end = env_steps[i + 1]
+                    if end - start > 1:
+                        episodes_start_end.append((env_ind, start, end - 1))
+            if len(episodes_start_end) > 0:
+                reward_trajs_split = [
+                    reward_trajs[start : end + 1, env_ind]
+                    for env_ind, start, end in episodes_start_end
+                ]
+                num_episode_finished = len(reward_trajs_split)
+                episode_reward = np.array(
+                    [np.sum(reward_traj) for reward_traj in reward_trajs_split]
+                )
+                episode_best_reward = np.array(
+                    [
+                        np.max(reward_traj) / self.act_steps
+                        for reward_traj in reward_trajs_split
+                    ]
+                )
+                avg_episode_reward = np.mean(episode_reward)
+                avg_best_reward = np.mean(episode_best_reward)
+                success_rate = np.mean(
+                    episode_best_reward >= self.best_reward_threshold_for_success
+                )
+            else:
+                episode_reward = np.array([])
+                num_episode_finished = 0
+                avg_episode_reward = 0
+                avg_best_reward = 0
+                success_rate = 0
+                log.info("[WARNING] No episode completed within the iteration!")
+
+            # Update
+            if not eval_mode:
+                with torch.no_grad():
+                    # Calculate value and logprobs - split into batches to prevent out of memory
+                    obs_t = einops.rearrange(
+                        torch.from_numpy(obs_trajs).float().to(self.device),
+                        "s e h d -> (s e) h d",
+                    )
+                    obs_ts = torch.split(obs_t, self.logprob_batch_size, dim=0)
+                    values_trajs = np.empty((0, self.n_envs))
+                    for obs in obs_ts:
+                        values = self.model.critic(obs).cpu().numpy().flatten()
+                        values_trajs = np.vstack(
+                            (values_trajs, values.reshape(-1, self.n_envs))
+                        )
+                    samples_t = einops.rearrange(
+                        torch.from_numpy(samples_trajs).float().to(self.device),
+                        "s e h d -> (s e) h d",
+                    )
+                    samples_ts = torch.split(samples_t, self.logprob_batch_size, dim=0)
+                    logprobs_trajs = np.empty((0))
+                    for obs, samples in zip(obs_ts, samples_ts):
+                        logprobs = (
+                            self.model.get_exact_logprobs(obs, samples).cpu().numpy()
+                        )
+                        logprobs_trajs = np.concatenate((logprobs_trajs, logprobs))
+
+                    # normalize reward with running variance if specified
+                    if self.reward_scale_running:
+                        reward_trajs_transpose = self.running_reward_scaler(
+                            reward=reward_trajs.T, first=firsts_trajs[:-1].T
+                        )
+                        reward_trajs = reward_trajs_transpose.T
+
+                    # bootstrap value with GAE if not done - apply reward scaling with constant if specified
+                    obs_venv_ts = torch.from_numpy(obs_venv).float().to(self.device)
+                    with torch.no_grad():
+                        next_value = (
+                            self.model.critic(obs_venv_ts).reshape(1, -1).cpu().numpy()
+                        )
+                    advantages_trajs = np.zeros_like(reward_trajs)
+                    lastgaelam = 0
+                    for t in reversed(range(self.n_steps)):
+                        if t == self.n_steps - 1:
+                            nextnonterminal = 1.0 - done_venv
+                            nextvalues = next_value
+                        else:
+                            nextnonterminal = 1.0 - dones_trajs[t + 1]
+                            nextvalues = values_trajs[t + 1]
+                        # delta = r + gamma*V(st+1) - V(st)
+                        delta = (
+                            reward_trajs[t] * self.reward_scale_const
+                            + self.gamma * nextvalues * nextnonterminal
+                            - values_trajs[t]
+                        )
+                        # A = delta_t + gamma*lamdba*delta_{t+1} + ...
+                        advantages_trajs[t] = lastgaelam = (
+                            delta
+                            + self.gamma
+                            * self.gae_lambda
+                            * nextnonterminal
+                            * lastgaelam
+                        )
+                    returns_trajs = advantages_trajs + values_trajs
+
+                # k for environment step
+                obs_k = einops.rearrange(
+                    torch.tensor(obs_trajs).float().to(self.device),
+                    "s e h d -> (s e) h d",
+                )
+                samples_k = einops.rearrange(
+                    torch.tensor(samples_trajs).float().to(self.device),
+                    "s e h d -> (s e) h d",
+                )
+                returns_k = (
+                    torch.tensor(returns_trajs).float().to(self.device).reshape(-1)
+                )
+                values_k = (
+                    torch.tensor(values_trajs).float().to(self.device).reshape(-1)
+                )
+                advantages_k = (
+                    torch.tensor(advantages_trajs).float().to(self.device).reshape(-1)
+                )
+                logprobs_k = torch.tensor(logprobs_trajs).float().to(self.device)
+
+                # Update policy and critic
+                total_steps = self.n_steps * self.n_envs
+                inds_k = np.arange(total_steps)
+                clipfracs = []
+                for update_epoch in range(self.update_epochs):
+
+                    # for each epoch, go through all data in batches
+                    flag_break = False
+                    np.random.shuffle(inds_k)
+                    num_batch = max(1, total_steps // self.batch_size)  # skip last ones
+                    for batch in range(num_batch):
+                        start = batch * self.batch_size
+                        end = start + self.batch_size
+                        inds_b = inds_k[start:end]  # b for batch
+                        obs_b = obs_k[inds_b]
+                        samples_b = samples_k[inds_b]
+                        returns_b = returns_k[inds_b]
+                        values_b = values_k[inds_b]
+                        advantages_b = advantages_k[inds_b]
+                        logprobs_b = logprobs_k[inds_b]
+
+                        # get loss
+                        (
+                            pg_loss,
+                            v_loss,
+                            clipfrac,
+                            approx_kl,
+                            ratio,
+                            bc_loss,
+                        ) = self.model.loss(
+                            obs_b,
+                            samples_b,
+                            returns_b,
+                            values_b,
+                            advantages_b,
+                            logprobs_b,
+                            use_bc_loss=self.use_bc_loss,
+                            reward_horizon=self.reward_horizon,
+                        )
+                        loss = (
+                            pg_loss
+                            + v_loss * self.vf_coef
+                            + bc_loss * self.bc_loss_coeff
+                        )
+                        clipfracs += [clipfrac]
+
+                        # update policy and critic
+                        self.actor_optimizer.zero_grad()
+                        self.critic_optimizer.zero_grad()
+                        loss.backward()
+                        if self.itr >= self.n_critic_warmup_itr:
+                            if self.max_grad_norm is not None:
+                                torch.nn.utils.clip_grad_norm_(
+                                    self.model.actor_ft.parameters(), self.max_grad_norm
+                                )
+                            self.actor_optimizer.step()
+                        self.critic_optimizer.step()
+                        log.info(
+                            f"approx_kl: {approx_kl}, update_epoch: {update_epoch}, num_batch: {num_batch}"
+                        )
+
+                        # Stop gradient update if KL difference reaches target
+                        if self.target_kl is not None and approx_kl > self.target_kl:
+                            flag_break = True
+                            break
+                    if flag_break:
+                        break
+
+                # Explained variation of future rewards using value function
+                y_pred, y_true = values_k.cpu().numpy(), returns_k.cpu().numpy()
+                var_y = np.var(y_true)
+                explained_var = (
+                    np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y
+                )
+
+            # Plot state trajectories
+            if (
+                self.itr % self.render_freq == 0
+                and self.n_render > 0
+                and self.traj_plotter is not None
+            ):
+                self.traj_plotter(
+                    obs_full_trajs=obs_full_trajs,
+                    n_render=self.n_render,
+                    max_episode_steps=self.max_episode_steps,
+                    render_dir=self.render_dir,
+                    itr=self.itr,
+                )
+
+            # Update lr
+            if self.itr >= self.n_critic_warmup_itr:
+                self.actor_lr_scheduler.step()
+            self.critic_lr_scheduler.step()
+
+            # Save model
+            if self.itr % self.save_model_freq == 0 or self.itr == self.n_train_itr - 1:
+                self.save_model()
+
+            # Log loss and save metrics
+            run_results.append(
+                {
+                    "itr": self.itr,
+                }
+            )
+            if self.save_trajs:
+                run_results[-1]["obs_full_trajs"] = obs_full_trajs
+                run_results[-1]["obs_trajs"] = obs_trajs
+                run_results[-1]["action_trajs"] = samples_trajs
+                run_results[-1]["chains_trajs"] = chains_trajs
+                run_results[-1]["reward_trajs"] = reward_trajs
+            if self.itr % self.log_freq == 0:
+                if eval_mode:
+                    log.info(
+                        f"eval: success rate {success_rate:8.4f} | avg episode reward {avg_episode_reward:8.4f} | avg best reward {avg_best_reward:8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "success rate - eval": success_rate,
+                                "avg episode reward - eval": avg_episode_reward,
+                                "avg best reward - eval": avg_best_reward,
+                                "num episode - eval": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=False,
+                        )
+                    run_results[-1]["eval_success_rate"] = success_rate
+                    run_results[-1]["eval_episode_reward"] = avg_episode_reward
+                    run_results[-1]["eval_best_reward"] = avg_best_reward
+                else:
+                    log.info(
+                        f"{self.itr}: loss {loss:8.4f} | pg loss {pg_loss:8.4f} | value loss {v_loss:8.4f} | reward {avg_episode_reward:8.4f} | t:{timer():8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "loss": loss,
+                                "pg loss": pg_loss,
+                                "value loss": v_loss,
+                                "approx kl": approx_kl,
+                                "ratio": ratio,
+                                "clipfrac": np.mean(clipfracs),
+                                "explained variance": explained_var,
+                                "avg episode reward - train": avg_episode_reward,
+                            },
+                            step=self.itr,
+                            commit=True,
+                        )
+                    run_results[-1]["loss"] = loss
+                    run_results[-1]["pg_loss"] = pg_loss
+                    run_results[-1]["value_loss"] = v_loss
+                    run_results[-1]["approx_kl"] = approx_kl
+                    run_results[-1]["ratio"] = ratio
+                    run_results[-1]["clip_frac"] = np.mean(clipfracs)
+                    run_results[-1]["explained_variance"] = explained_var
+                    run_results[-1]["train_episode_reward"] = avg_episode_reward
+                run_results[-1]["time"] = timer()
+                with open(self.result_path, "wb") as f:
+                    pickle.dump(run_results, f)
+            self.itr += 1

agent/finetune/train_ppo_gaussian_agent.py

@@ -0,0 +1,404 @@
+"""
+PPO training for Gaussian/GMM policy.
+
+"""
+
+import os
+import pickle
+import einops
+import numpy as np
+import torch
+import logging
+import wandb
+
+log = logging.getLogger(__name__)
+from util.timer import Timer
+from agent.finetune.train_ppo_agent import TrainPPOAgent
+
+
+class TrainPPOGaussianAgent(TrainPPOAgent):
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+    def run(self):
+
+        # Start training loop
+        timer = Timer()
+        run_results = []
+        last_itr_eval = False
+        done_venv = np.zeros((1, self.n_envs))
+        while self.itr < self.n_train_itr:
+
+            # Prepare video paths for each envs --- only applies for the first set of episodes if allowing reset within iteration and each iteration has multiple episodes from one env
+            options_venv = [{} for _ in range(self.n_envs)]
+            if self.itr % self.render_freq == 0 and self.render_video:
+                for env_ind in range(self.n_render):
+                    options_venv[env_ind]["video_path"] = os.path.join(
+                        self.render_dir, f"itr-{self.itr}_trial-{env_ind}.mp4"
+                    )
+
+            # Define train or eval - all envs restart
+            eval_mode = self.itr % self.val_freq == 0 and not self.force_train
+            self.model.eval() if eval_mode else self.model.train()
+            last_itr_eval = eval_mode
+
+            # Reset env before iteration starts (1) if specified, (2) at eval mode, or (3) right after eval mode
+            dones_trajs = np.empty((0, self.n_envs))
+            firsts_trajs = np.zeros((self.n_steps + 1, self.n_envs))
+            if self.reset_at_iteration or eval_mode or last_itr_eval:
+                prev_obs_venv = self.reset_env_all(options_venv=options_venv)
+                firsts_trajs[0] = 1
+            else:
+                firsts_trajs[0] = (
+                    done_venv  # if done at the end of last iteration, then the envs are just reset
+                )
+
+            # Holder
+            obs_trajs = np.empty((0, self.n_envs, self.n_cond_step, self.obs_dim))
+            samples_trajs = np.empty(
+                (
+                    0,
+                    self.n_envs,
+                    self.horizon_steps,
+                    self.action_dim,
+                )
+            )
+            reward_trajs = np.empty((0, self.n_envs))
+            obs_full_trajs = np.empty((0, self.n_envs, self.obs_dim))
+            obs_full_trajs = np.vstack(
+                (obs_full_trajs, prev_obs_venv[None].squeeze(2))
+            )  # remove cond_step dim
+
+            # Collect a set of trajectories from env
+            for step in range(self.n_steps):
+                if step % 10 == 0:
+                    print(f"Processed step {step} of {self.n_steps}")
+
+                # Select action
+                with torch.no_grad():
+                    samples = self.model(
+                        cond=torch.from_numpy(prev_obs_venv).float().to(self.device),
+                        deterministic=eval_mode,
+                    )
+                    output_venv = samples.cpu().numpy()
+                action_venv = output_venv[:, : self.act_steps, : self.action_dim]
+                obs_trajs = np.vstack((obs_trajs, prev_obs_venv[None]))
+                samples_trajs = np.vstack((samples_trajs, output_venv[None]))
+
+                # Apply multi-step action
+                obs_venv, reward_venv, done_venv, info_venv = self.venv.step(
+                    action_venv
+                )
+                if self.save_full_observations:
+                    obs_full_venv = np.vstack(
+                        [info["full_obs"][None] for info in info_venv]
+                    )  # n_envs x n_act_steps x obs_dim
+                    obs_full_trajs = np.vstack(
+                        (obs_full_trajs, obs_full_venv.transpose(1, 0, 2))
+                    )
+                reward_trajs = np.vstack((reward_trajs, reward_venv[None]))
+                dones_trajs = np.vstack((dones_trajs, done_venv[None]))
+                firsts_trajs[step + 1] = done_venv
+                prev_obs_venv = obs_venv
+
+            # Summarize episode reward --- this needs to be handled differently depending on whether the environment is reset after each iteration. Only count episodes that finish within the iteration.
+            episodes_start_end = []
+            for env_ind in range(self.n_envs):
+                env_steps = np.where(firsts_trajs[:, env_ind] == 1)[0]
+                for i in range(len(env_steps) - 1):
+                    start = env_steps[i]
+                    end = env_steps[i + 1]
+                    if end - start > 1:
+                        episodes_start_end.append((env_ind, start, end - 1))
+            if len(episodes_start_end) > 0:
+                reward_trajs_split = [
+                    reward_trajs[start : end + 1, env_ind]
+                    for env_ind, start, end in episodes_start_end
+                ]
+                num_episode_finished = len(reward_trajs_split)
+                episode_reward = np.array(
+                    [np.sum(reward_traj) for reward_traj in reward_trajs_split]
+                )
+                if (
+                    self.furniture_sparse_reward
+                ):  # only for furniture tasks, where reward only occurs in one env step
+                    episode_best_reward = episode_reward
+                else:
+                    episode_best_reward = np.array(
+                        [
+                            np.max(reward_traj) / self.act_steps
+                            for reward_traj in reward_trajs_split
+                        ]
+                    )
+                avg_episode_reward = np.mean(episode_reward)
+                avg_best_reward = np.mean(episode_best_reward)
+                success_rate = np.mean(
+                    episode_best_reward >= self.best_reward_threshold_for_success
+                )
+            else:
+                episode_reward = np.array([])
+                num_episode_finished = 0
+                avg_episode_reward = 0
+                avg_best_reward = 0
+                success_rate = 0
+                log.info("[WARNING] No episode completed within the iteration!")
+
+            # Update
+            if not eval_mode:
+                with torch.no_grad():
+                    # Calculate value and logprobs - split into batches to prevent out of memory
+                    obs_t = einops.rearrange(
+                        torch.from_numpy(obs_trajs).float().to(self.device),
+                        "s e h d -> (s e) h d",
+                    )
+                    obs_ts = torch.split(obs_t, self.logprob_batch_size, dim=0)
+                    values_trajs = np.empty((0, self.n_envs))
+                    for obs in obs_ts:
+                        values = self.model.critic(obs).cpu().numpy().flatten()
+                        values_trajs = np.vstack(
+                            (values_trajs, values.reshape(-1, self.n_envs))
+                        )
+                    samples_t = einops.rearrange(
+                        torch.from_numpy(samples_trajs).float().to(self.device),
+                        "s e h d -> (s e) h d",
+                    )
+                    samples_ts = torch.split(samples_t, self.logprob_batch_size, dim=0)
+                    logprobs_trajs = np.empty((0))
+                    for obs_t, samples_t in zip(obs_ts, samples_ts):
+                        logprobs = (
+                            self.model.get_logprobs(obs_t, samples_t)[0].cpu().numpy()
+                        )
+                        logprobs_trajs = np.concatenate(
+                            (
+                                logprobs_trajs,
+                                logprobs.reshape(-1),
+                            )
+                        )
+
+                    # normalize reward with running variance if specified
+                    if self.reward_scale_running:
+                        reward_trajs_transpose = self.running_reward_scaler(
+                            reward=reward_trajs.T, first=firsts_trajs[:-1].T
+                        )
+                        reward_trajs = reward_trajs_transpose.T
+
+                    # bootstrap value with GAE if not done - apply reward scaling with constant if specified
+                    obs_venv_ts = torch.from_numpy(obs_venv).float().to(self.device)
+                    with torch.no_grad():
+                        next_value = (
+                            self.model.critic(obs_venv_ts).reshape(1, -1).cpu().numpy()
+                        )
+                    advantages_trajs = np.zeros_like(reward_trajs)
+                    lastgaelam = 0
+                    for t in reversed(range(self.n_steps)):
+                        if t == self.n_steps - 1:
+                            nextnonterminal = 1.0 - done_venv
+                            nextvalues = next_value
+                        else:
+                            nextnonterminal = 1.0 - dones_trajs[t + 1]
+                            nextvalues = values_trajs[t + 1]
+                        # delta = r + gamma*V(st+1) - V(st)
+                        delta = (
+                            reward_trajs[t] * self.reward_scale_const
+                            + self.gamma * nextvalues * nextnonterminal
+                            - values_trajs[t]
+                        )
+                        # A = delta_t + gamma*lamdba*delta_{t+1} + ...
+                        advantages_trajs[t] = lastgaelam = (
+                            delta
+                            + self.gamma
+                            * self.gae_lambda
+                            * nextnonterminal
+                            * lastgaelam
+                        )
+                    returns_trajs = advantages_trajs + values_trajs
+
+                # k for environment step
+                obs_k = einops.rearrange(
+                    torch.tensor(obs_trajs).float().to(self.device),
+                    "s e h d -> (s e) h d",
+                )
+                samples_k = einops.rearrange(
+                    torch.tensor(samples_trajs).float().to(self.device),
+                    "s e h d -> (s e) h d",
+                )
+                returns_k = (
+                    torch.tensor(returns_trajs).float().to(self.device).reshape(-1)
+                )
+                values_k = (
+                    torch.tensor(values_trajs).float().to(self.device).reshape(-1)
+                )
+                advantages_k = (
+                    torch.tensor(advantages_trajs).float().to(self.device).reshape(-1)
+                )
+                logprobs_k = (
+                    torch.tensor(logprobs_trajs).float().to(self.device).reshape(-1)
+                )
+
+                # Update policy and critic
+                total_steps = self.n_steps * self.n_envs
+                inds_k = np.arange(total_steps)
+                clipfracs = []
+                for update_epoch in range(self.update_epochs):
+
+                    # for each epoch, go through all data in batches
+                    flag_break = False
+                    np.random.shuffle(inds_k)
+                    num_batch = max(1, total_steps // self.batch_size)  # skip last ones
+                    for batch in range(num_batch):
+                        start = batch * self.batch_size
+                        end = start + self.batch_size
+                        inds_b = inds_k[start:end]  # b for batch
+                        obs_b = obs_k[inds_b]
+                        samples_b = samples_k[inds_b]
+                        returns_b = returns_k[inds_b]
+                        values_b = values_k[inds_b]
+                        advantages_b = advantages_k[inds_b]
+                        logprobs_b = logprobs_k[inds_b]
+
+                        # get loss
+                        (
+                            pg_loss,
+                            entropy_loss,
+                            v_loss,
+                            clipfrac,
+                            approx_kl,
+                            ratio,
+                            bc_loss,
+                            std,
+                        ) = self.model.loss(
+                            obs_b,
+                            samples_b,
+                            returns_b,
+                            values_b,
+                            advantages_b,
+                            logprobs_b,
+                            use_bc_loss=self.use_bc_loss,
+                        )
+                        loss = (
+                            pg_loss
+                            + entropy_loss * self.ent_coef
+                            + v_loss * self.vf_coef
+                            + bc_loss * self.bc_loss_coeff
+                        )
+                        clipfracs += [clipfrac]
+
+                        # update policy and critic
+                        self.actor_optimizer.zero_grad()
+                        self.critic_optimizer.zero_grad()
+                        loss.backward()
+                        if self.itr >= self.n_critic_warmup_itr:
+                            if self.max_grad_norm is not None:
+                                torch.nn.utils.clip_grad_norm_(
+                                    self.model.actor_ft.parameters(), self.max_grad_norm
+                                )
+                            self.actor_optimizer.step()
+                        self.critic_optimizer.step()
+                        log.info(
+                            f"approx_kl: {approx_kl}, update_epoch: {update_epoch}, num_batch: {num_batch}"
+                        )
+
+                        # Stop gradient update if KL difference reaches target
+                        if self.target_kl is not None and approx_kl > self.target_kl:
+                            flag_break = True
+                            break
+                    if flag_break:
+                        break
+
+                # Explained variation of future rewards using value function
+                y_pred, y_true = values_k.cpu().numpy(), returns_k.cpu().numpy()
+                var_y = np.var(y_true)
+                explained_var = (
+                    np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y
+                )
+
+            # Plot state trajectories
+            if (
+                self.itr % self.render_freq == 0
+                and self.n_render > 0
+                and self.traj_plotter is not None
+            ):
+                self.traj_plotter(
+                    obs_full_trajs=obs_full_trajs,
+                    n_render=self.n_render,
+                    max_episode_steps=self.max_episode_steps,
+                    render_dir=self.render_dir,
+                    itr=self.itr,
+                )
+
+            # Update lr
+            if self.itr >= self.n_critic_warmup_itr:
+                self.actor_lr_scheduler.step()
+            self.critic_lr_scheduler.step()
+
+            # Save model
+            if self.itr % self.save_model_freq == 0 or self.itr == self.n_train_itr - 1:
+                self.save_model()
+
+            # Log loss and save metrics
+            run_results.append(
+                {
+                    "itr": self.itr,
+                }
+            )
+            if self.save_trajs:
+                run_results[-1]["obs_full_trajs"] = obs_full_trajs
+                run_results[-1]["obs_trajs"] = obs_trajs
+                run_results[-1]["action_trajs"] = samples_trajs
+                run_results[-1]["reward_trajs"] = reward_trajs
+            if self.itr % self.log_freq == 0:
+                time = timer()
+                if eval_mode:
+                    log.info(
+                        f"eval: success rate {success_rate:8.4f} | avg episode reward {avg_episode_reward:8.4f} | avg best reward {avg_best_reward:8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "success rate - eval": success_rate,
+                                "avg episode reward - eval": avg_episode_reward,
+                                "avg best reward - eval": avg_best_reward,
+                                "num episode - eval": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=False,
+                        )
+                    run_results[-1]["eval_success_rate"] = success_rate
+                    run_results[-1]["eval_episode_reward"] = avg_episode_reward
+                    run_results[-1]["eval_best_reward"] = avg_best_reward
+                else:
+                    log.info(
+                        f"{self.itr}: loss {loss:8.4f} | pg loss {pg_loss:8.4f} | value loss {v_loss:8.4f} | ent {-entropy_loss:8.4f} | reward {avg_episode_reward:8.4f} | t:{time:8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "loss": loss,
+                                "pg loss": pg_loss,
+                                "value loss": v_loss,
+                                "entropy": -entropy_loss,
+                                "std": std,
+                                "approx kl": approx_kl,
+                                "ratio": ratio,
+                                "clipfrac": np.mean(clipfracs),
+                                "explained variance": explained_var,
+                                "avg episode reward - train": avg_episode_reward,
+                                "num episode - train": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=True,
+                        )
+                    run_results[-1]["loss"] = loss
+                    run_results[-1]["pg_loss"] = pg_loss
+                    run_results[-1]["value_loss"] = v_loss
+                    run_results[-1]["entropy_loss"] = entropy_loss
+                    run_results[-1]["approx_kl"] = approx_kl
+                    run_results[-1]["ratio"] = ratio
+                    run_results[-1]["clip_frac"] = np.mean(clipfracs)
+                    run_results[-1]["explained_variance"] = explained_var
+                    run_results[-1]["train_episode_reward"] = avg_episode_reward
+                run_results[-1]["time"] = time
+                with open(self.result_path, "wb") as f:
+                    pickle.dump(run_results, f)
+            self.itr += 1

agent/finetune/train_ppo_gaussian_img_agent.py

@@ -0,0 +1,443 @@
+"""
+PPO training for Gaussian/GMM policy with pixel observations.
+
+"""
+
+import os
+import pickle
+import einops
+import numpy as np
+import torch
+import logging
+import wandb
+import math
+
+log = logging.getLogger(__name__)
+from util.timer import Timer
+from agent.finetune.train_ppo_gaussian_agent import TrainPPOGaussianAgent
+from model.common.modules import RandomShiftsAug
+
+
+class TrainPPOImgGaussianAgent(TrainPPOGaussianAgent):
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+        # Image randomization
+        self.augment = cfg.train.augment
+        if self.augment:
+            self.aug = RandomShiftsAug(pad=4)
+
+        # Set obs dim -  we will save the different obs in batch in a dict
+        shape_meta = cfg.shape_meta
+        self.obs_dims = {k: shape_meta.obs[k]["shape"] for k in shape_meta.obs.keys()}
+
+        # Gradient accumulation to deal with large GPU RAM usage
+        self.grad_accumulate = cfg.train.grad_accumulate
+
+    def run(self):
+
+        # Start training loop
+        timer = Timer()
+        run_results = []
+        last_itr_eval = False
+        done_venv = np.zeros((1, self.n_envs))
+        while self.itr < self.n_train_itr:
+
+            # Prepare video paths for each envs --- only applies for the first set of episodes if allowing reset within iteration and each iteration has multiple episodes from one env
+            options_venv = [{} for _ in range(self.n_envs)]
+            if self.itr % self.render_freq == 0 and self.render_video:
+                for env_ind in range(self.n_render):
+                    options_venv[env_ind]["video_path"] = os.path.join(
+                        self.render_dir, f"itr-{self.itr}_trial-{env_ind}.mp4"
+                    )
+
+            # Define train or eval - all envs restart
+            eval_mode = self.itr % self.val_freq == 0 and not self.force_train
+            self.model.eval() if eval_mode else self.model.train()
+            last_itr_eval = eval_mode
+
+            # Reset env before iteration starts (1) if specified, (2) at eval mode, or (3) right after eval mode
+            dones_trajs = np.empty((0, self.n_envs))
+            firsts_trajs = np.zeros((self.n_steps + 1, self.n_envs))
+            if self.reset_at_iteration or eval_mode or last_itr_eval:
+                prev_obs_venv = self.reset_env_all(options_venv=options_venv)
+                firsts_trajs[0] = 1
+            else:
+                firsts_trajs[0] = (
+                    done_venv  # if done at the end of last iteration, then the envs are just reset
+                )
+
+            # Holder
+            obs_trajs = {
+                k: np.empty((0, self.n_envs, self.n_cond_step, *self.obs_dims[k]))
+                for k in self.obs_dims
+            }
+            samples_trajs = np.empty(
+                (
+                    0,
+                    self.n_envs,
+                    self.horizon_steps,
+                    self.action_dim,
+                )
+            )
+            reward_trajs = np.empty((0, self.n_envs))
+
+            # Collect a set of trajectories from env
+            for step in range(self.n_steps):
+                if step % 10 == 0:
+                    print(f"Processed step {step} of {self.n_steps}")
+
+                # Select action
+                with torch.no_grad():
+                    cond = {
+                        key: torch.from_numpy(prev_obs_venv[key])
+                        .float()
+                        .to(self.device)
+                        for key in self.obs_dims.keys()
+                    }  # batch each type of obs and put into dict
+                    samples = self.model(
+                        cond=cond,
+                        deterministic=eval_mode,
+                    )
+                    output_venv = samples.cpu().numpy()
+                action_venv = output_venv[:, : self.act_steps]
+
+                # Apply multi-step action
+                obs_venv, reward_venv, done_venv, info_venv = self.venv.step(
+                    action_venv
+                )
+                for k in obs_trajs.keys():
+                    obs_trajs[k] = np.vstack((obs_trajs[k], prev_obs_venv[k][None]))
+                samples_trajs = np.vstack((samples_trajs, output_venv[None]))
+                reward_trajs = np.vstack((reward_trajs, reward_venv[None]))
+                dones_trajs = np.vstack((dones_trajs, done_venv[None]))
+                firsts_trajs[step + 1] = done_venv
+                prev_obs_venv = obs_venv
+
+            # Summarize episode reward --- this needs to be handled differently depending on whether the environment is reset after each iteration. Only count episodes that finish within the iteration.
+            episodes_start_end = []
+            for env_ind in range(self.n_envs):
+                env_steps = np.where(firsts_trajs[:, env_ind] == 1)[0]
+                for i in range(len(env_steps) - 1):
+                    start = env_steps[i]
+                    end = env_steps[i + 1]
+                    if end - start > 1:
+                        episodes_start_end.append((env_ind, start, end - 1))
+            if len(episodes_start_end) > 0:
+                reward_trajs_split = [
+                    reward_trajs[start : end + 1, env_ind]
+                    for env_ind, start, end in episodes_start_end
+                ]
+                num_episode_finished = len(reward_trajs_split)
+                episode_reward = np.array(
+                    [np.sum(reward_traj) for reward_traj in reward_trajs_split]
+                )
+                episode_best_reward = np.array(
+                    [
+                        np.max(reward_traj) / self.act_steps
+                        for reward_traj in reward_trajs_split
+                    ]
+                )
+                avg_episode_reward = np.mean(episode_reward)
+                avg_best_reward = np.mean(episode_best_reward)
+                success_rate = np.mean(
+                    episode_best_reward >= self.best_reward_threshold_for_success
+                )
+            else:
+                episode_reward = np.array([])
+                num_episode_finished = 0
+                avg_episode_reward = 0
+                avg_best_reward = 0
+                success_rate = 0
+                log.info("[WARNING] No episode completed within the iteration!")
+
+            # Update
+            if not eval_mode:
+                with torch.no_grad():
+                    # apply image randomization
+                    obs_trajs["rgb"] = (
+                        torch.from_numpy(obs_trajs["rgb"]).float().to(self.device)
+                    )
+                    obs_trajs["state"] = (
+                        torch.from_numpy(obs_trajs["state"]).float().to(self.device)
+                    )
+                    if self.augment:
+                        rgb = einops.rearrange(
+                            obs_trajs["rgb"],
+                            "s e t c h w -> (s e t) c h w",
+                        )
+                        rgb = self.aug(rgb)
+                        obs_trajs["rgb"] = einops.rearrange(
+                            rgb,
+                            "(s e t) c h w -> s e t c h w",
+                            s=self.n_steps,
+                            e=self.n_envs,
+                        )
+
+                    # Calculate value and logprobs - split into batches to prevent out of memory
+                    num_split = math.ceil(
+                        self.n_envs * self.n_steps / self.logprob_batch_size
+                    )
+                    obs_ts = [{} for _ in range(num_split)]
+                    for k in obs_trajs.keys():
+                        obs_k = einops.rearrange(
+                            obs_trajs[k],
+                            "s e ... -> (s e) ...",
+                        )
+                        obs_ts_k = torch.split(obs_k, self.logprob_batch_size, dim=0)
+                        for i, obs_t in enumerate(obs_ts_k):
+                            obs_ts[i][k] = obs_t
+                    values_trajs = np.empty((0, self.n_envs))
+                    for obs in obs_ts:
+                        values = (
+                            self.model.critic(obs, no_augment=True)
+                            .cpu()
+                            .numpy()
+                            .flatten()
+                        )
+                        values_trajs = np.vstack(
+                            (values_trajs, values.reshape(-1, self.n_envs))
+                        )
+                    samples_t = einops.rearrange(
+                        torch.from_numpy(samples_trajs).float().to(self.device),
+                        "s e h d -> (s e) h d",
+                    )
+                    samples_ts = torch.split(samples_t, self.logprob_batch_size, dim=0)
+                    logprobs_trajs = np.empty((0))
+                    for obs_t, samples_t in zip(obs_ts, samples_ts):
+                        logprobs = (
+                            self.model.get_logprobs(obs_t, samples_t)[0].cpu().numpy()
+                        )
+                        logprobs_trajs = np.concatenate(
+                            (
+                                logprobs_trajs,
+                                logprobs.reshape(-1),
+                            )
+                        )
+
+                    # normalize reward with running variance if specified
+                    if self.reward_scale_running:
+                        reward_trajs_transpose = self.running_reward_scaler(
+                            reward=reward_trajs.T, first=firsts_trajs[:-1].T
+                        )
+                        reward_trajs = reward_trajs_transpose.T
+
+                    # bootstrap value with GAE if not done - apply reward scaling with constant if specified
+                    obs_venv_ts = {
+                        key: torch.from_numpy(obs_venv[key]).float().to(self.device)
+                        for key in self.obs_dims.keys()
+                    }
+                    with torch.no_grad():
+                        next_value = (
+                            self.model.critic(obs_venv_ts, no_augment=True)
+                            .reshape(1, -1)
+                            .cpu()
+                            .numpy()
+                        )
+                    advantages_trajs = np.zeros_like(reward_trajs)
+                    lastgaelam = 0
+                    for t in reversed(range(self.n_steps)):
+                        if t == self.n_steps - 1:
+                            nextnonterminal = 1.0 - done_venv
+                            nextvalues = next_value
+                        else:
+                            nextnonterminal = 1.0 - dones_trajs[t + 1]
+                            nextvalues = values_trajs[t + 1]
+                        # delta = r + gamma*V(st+1) - V(st)
+                        delta = (
+                            reward_trajs[t] * self.reward_scale_const
+                            + self.gamma * nextvalues * nextnonterminal
+                            - values_trajs[t]
+                        )
+                        # A = delta_t + gamma*lamdba*delta_{t+1} + ...
+                        advantages_trajs[t] = lastgaelam = (
+                            delta
+                            + self.gamma
+                            * self.gae_lambda
+                            * nextnonterminal
+                            * lastgaelam
+                        )
+                    returns_trajs = advantages_trajs + values_trajs
+
+                # k for environment step
+                obs_k = {
+                    k: einops.rearrange(
+                        obs_trajs[k],
+                        "s e ... -> (s e) ...",
+                    )
+                    for k in obs_trajs.keys()
+                }
+                samples_k = einops.rearrange(
+                    torch.tensor(samples_trajs).float().to(self.device),
+                    "s e h d -> (s e) h d",
+                )
+                returns_k = (
+                    torch.tensor(returns_trajs).float().to(self.device).reshape(-1)
+                )
+                values_k = (
+                    torch.tensor(values_trajs).float().to(self.device).reshape(-1)
+                )
+                advantages_k = (
+                    torch.tensor(advantages_trajs).float().to(self.device).reshape(-1)
+                )
+                logprobs_k = torch.tensor(logprobs_trajs).float().to(self.device)
+
+                # Update policy and critic
+                total_steps = self.n_steps * self.n_envs
+                inds_k = np.arange(total_steps)
+                clipfracs = []
+                for update_epoch in range(self.update_epochs):
+
+                    # for each epoch, go through all data in batches
+                    flag_break = False
+                    np.random.shuffle(inds_k)
+                    num_batch = max(1, total_steps // self.batch_size)  # skip last ones
+                    for batch in range(num_batch):
+                        start = batch * self.batch_size
+                        end = start + self.batch_size
+                        inds_b = inds_k[start:end]  # b for batch
+                        obs_b = {k: obs_k[k][inds_b] for k in obs_k.keys()}
+                        samples_b = samples_k[inds_b]
+                        returns_b = returns_k[inds_b]
+                        values_b = values_k[inds_b]
+                        advantages_b = advantages_k[inds_b]
+                        logprobs_b = logprobs_k[inds_b]
+
+                        # get loss
+                        (
+                            pg_loss,
+                            entropy_loss,
+                            v_loss,
+                            clipfrac,
+                            approx_kl,
+                            ratio,
+                            bc_loss,
+                            std,
+                        ) = self.model.loss(
+                            obs_b,
+                            samples_b,
+                            returns_b,
+                            values_b,
+                            advantages_b,
+                            logprobs_b,
+                            use_bc_loss=self.use_bc_loss,
+                        )
+                        loss = (
+                            pg_loss
+                            + entropy_loss * self.ent_coef
+                            + v_loss * self.vf_coef
+                            + bc_loss * self.bc_loss_coeff
+                        )
+                        clipfracs += [clipfrac]
+
+                        # update policy and critic
+                        loss.backward()
+                        if (batch + 1) % self.grad_accumulate == 0:
+                            if self.itr >= self.n_critic_warmup_itr:
+                                if self.max_grad_norm is not None:
+                                    torch.nn.utils.clip_grad_norm_(
+                                        self.model.actor_ft.parameters(),
+                                        self.max_grad_norm,
+                                    )
+                                self.actor_optimizer.step()
+                            self.critic_optimizer.step()
+                            self.actor_optimizer.zero_grad()
+                            self.critic_optimizer.zero_grad()
+                            log.info(f"run grad update at batch {batch}")
+                            log.info(
+                                f"approx_kl: {approx_kl}, update_epoch: {update_epoch}, num_batch: {num_batch}"
+                            )
+
+                            # Stop gradient update if KL difference reaches target
+                            if (
+                                self.target_kl is not None
+                                and approx_kl > self.target_kl
+                                and self.itr >= self.n_critic_warmup_itr
+                            ):
+                                flag_break = True
+                                break
+                    if flag_break:
+                        break
+
+                # Explained variation of future rewards using value function
+                y_pred, y_true = values_k.cpu().numpy(), returns_k.cpu().numpy()
+                var_y = np.var(y_true)
+                explained_var = (
+                    np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y
+                )
+
+            # Update lr
+            if self.itr >= self.n_critic_warmup_itr:
+                self.actor_lr_scheduler.step()
+            self.critic_lr_scheduler.step()
+
+            # Save model
+            if self.itr % self.save_model_freq == 0 or self.itr == self.n_train_itr - 1:
+                self.save_model()
+
+            # Log loss and save metrics
+            run_results.append(
+                {
+                    "itr": self.itr,
+                }
+            )
+            if self.itr % self.log_freq == 0:
+                if eval_mode:
+                    log.info(
+                        f"eval: success rate {success_rate:8.4f} | avg episode reward {avg_episode_reward:8.4f} | avg best reward {avg_best_reward:8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "success rate - eval": success_rate,
+                                "avg episode reward - eval": avg_episode_reward,
+                                "avg best reward - eval": avg_best_reward,
+                                "num episode - eval": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=False,
+                        )
+                    run_results[-1]["eval_success_rate"] = success_rate
+                    run_results[-1]["eval_episode_reward"] = avg_episode_reward
+                    run_results[-1]["eval_best_reward"] = avg_best_reward
+                else:
+                    log.info(
+                        f"{self.itr}: loss {loss:8.4f} | pg loss {pg_loss:8.4f} | value loss {v_loss:8.4f} | bc loss {bc_loss:8.4f} | reward {avg_episode_reward:8.4f} | t:{timer():8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "loss": loss,
+                                "pg loss": pg_loss,
+                                "value loss": v_loss,
+                                "bc loss": bc_loss,
+                                "std": std,
+                                "approx kl": approx_kl,
+                                "ratio": ratio,
+                                "clipfrac": np.mean(clipfracs),
+                                "explained variance": explained_var,
+                                "avg episode reward - train": avg_episode_reward,
+                                "num episode - train": num_episode_finished,
+                                "actor lr": self.actor_optimizer.param_groups[0]["lr"],
+                                "critic lr": self.critic_optimizer.param_groups[0][
+                                    "lr"
+                                ],
+                            },
+                            step=self.itr,
+                            commit=True,
+                        )
+                    run_results[-1]["loss"] = loss
+                    run_results[-1]["pg_loss"] = pg_loss
+                    run_results[-1]["value_loss"] = v_loss
+                    run_results[-1]["bc_loss"] = bc_loss
+                    run_results[-1]["std"] = std
+                    run_results[-1]["approx_kl"] = approx_kl
+                    run_results[-1]["ratio"] = ratio
+                    run_results[-1]["clip_frac"] = np.mean(clipfracs)
+                    run_results[-1]["explained_variance"] = explained_var
+                    run_results[-1]["train_episode_reward"] = avg_episode_reward
+                run_results[-1]["time"] = timer()
+                with open(self.result_path, "wb") as f:
+                    pickle.dump(run_results, f)
+            self.itr += 1

agent/finetune/train_qsm_diffusion_agent.py

@@ -0,0 +1,316 @@
+"""
+QSM (Q-Score Matching) for diffusion policy.
+
+"""
+
+import os
+import pickle
+import einops
+import numpy as np
+import torch
+import logging
+import wandb
+from copy import deepcopy
+
+log = logging.getLogger(__name__)
+from util.timer import Timer
+from collections import deque
+from agent.finetune.train_agent import TrainAgent
+from util.scheduler import CosineAnnealingWarmupRestarts
+
+
+class TrainQSMDiffusionAgent(TrainAgent):
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+        # note the discount factor gamma here is applied to reward every act_steps, instead of every env step
+        self.gamma = cfg.train.gamma
+
+        # Wwarm up period for critic before actor updates
+        self.n_critic_warmup_itr = cfg.train.n_critic_warmup_itr
+
+        # Optimizer
+        self.actor_optimizer = torch.optim.AdamW(
+            self.model.actor.parameters(),
+            lr=cfg.train.actor_lr,
+            weight_decay=cfg.train.actor_weight_decay,
+        )
+        self.actor_lr_scheduler = CosineAnnealingWarmupRestarts(
+            self.actor_optimizer,
+            first_cycle_steps=cfg.train.actor_lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=cfg.train.actor_lr,
+            min_lr=cfg.train.actor_lr_scheduler.min_lr,
+            warmup_steps=cfg.train.actor_lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+        self.critic_optimizer = torch.optim.AdamW(
+            self.model.critic_q.parameters(),
+            lr=cfg.train.critic_lr,
+            weight_decay=cfg.train.critic_weight_decay,
+        )
+        self.critic_lr_scheduler = CosineAnnealingWarmupRestarts(
+            self.critic_optimizer,
+            first_cycle_steps=cfg.train.critic_lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=cfg.train.critic_lr,
+            min_lr=cfg.train.critic_lr_scheduler.min_lr,
+            warmup_steps=cfg.train.critic_lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+
+        # Buffer size
+        self.buffer_size = cfg.train.buffer_size
+
+        # Scaling reward
+        self.scale_reward_factor = cfg.train.scale_reward_factor
+
+        # Updates
+        self.replay_ratio = cfg.train.replay_ratio
+        self.critic_tau = cfg.train.critic_tau
+        self.q_grad_coeff = cfg.train.q_grad_coeff
+
+    def run(self):
+
+        # make a FIFO replay buffer for obs, action, and reward
+        obs_buffer = deque(maxlen=self.buffer_size)
+        action_buffer = deque(maxlen=self.buffer_size)
+        next_obs_buffer = deque(maxlen=self.buffer_size)
+        reward_buffer = deque(maxlen=self.buffer_size)
+        done_buffer = deque(maxlen=self.buffer_size)
+        first_buffer = deque(maxlen=self.buffer_size)
+
+        # Start training loop
+        timer = Timer()
+        run_results = []
+        done_venv = np.zeros((1, self.n_envs))
+        while self.itr < self.n_train_itr:
+
+            # Prepare video paths for each envs --- only applies for the first set of episodes if allowing reset within iteration and each iteration has multiple episodes from one env
+            options_venv = [{} for _ in range(self.n_envs)]
+            if self.itr % self.render_freq == 0 and self.render_video:
+                for env_ind in range(self.n_render):
+                    options_venv[env_ind]["video_path"] = os.path.join(
+                        self.render_dir, f"itr-{self.itr}_trial-{env_ind}.mp4"
+                    )
+
+            # Define train or eval - all envs restart
+            eval_mode = self.itr % self.val_freq == 0 and not self.force_train
+            self.model.eval() if eval_mode else self.model.train()
+            firsts_trajs = np.zeros((self.n_steps + 1, self.n_envs))
+
+            # Reset env at the beginning of an iteration
+            if self.reset_at_iteration or eval_mode or last_itr_eval:
+                prev_obs_venv = self.reset_env_all(options_venv=options_venv)
+                firsts_trajs[0] = 1
+            else:
+                firsts_trajs[0] = (
+                    done_venv  # if done at the end of last iteration, then the envs are just reset
+                )
+            last_itr_eval = eval_mode
+            reward_trajs = np.empty((0, self.n_envs))
+
+            # Collect a set of trajectories from env
+            for step in range(self.n_steps):
+                if step % 10 == 0:
+                    print(f"Processed step {step} of {self.n_steps}")
+
+                # Select action
+                with torch.no_grad():
+                    samples = (
+                        self.model(
+                            cond=torch.from_numpy(prev_obs_venv)
+                            .float()
+                            .to(self.device),
+                            deterministic=eval_mode,
+                        )
+                        .cpu()
+                        .numpy()
+                    )  # n_env x horizon x act
+                action_venv = samples[:, : self.act_steps]
+
+                # Apply multi-step action
+                obs_venv, reward_venv, done_venv, info_venv = self.venv.step(
+                    action_venv
+                )
+                reward_trajs = np.vstack((reward_trajs, reward_venv[None]))
+
+                # add to buffer
+                obs_buffer.append(prev_obs_venv)
+                action_buffer.append(action_venv)
+                next_obs_buffer.append(obs_venv)
+                reward_buffer.append(reward_venv * self.scale_reward_factor)
+                done_buffer.append(done_venv)
+                first_buffer.append(firsts_trajs[step])
+
+                firsts_trajs[step + 1] = done_venv
+                prev_obs_venv = obs_venv
+
+            # Summarize episode reward --- this needs to be handled differently depending on whether the environment is reset after each iteration. Only count episodes that finish within the iteration.
+            episodes_start_end = []
+            for env_ind in range(self.n_envs):
+                env_steps = np.where(firsts_trajs[:, env_ind] == 1)[0]
+                for i in range(len(env_steps) - 1):
+                    start = env_steps[i]
+                    end = env_steps[i + 1]
+                    if end - start > 1:
+                        episodes_start_end.append((env_ind, start, end - 1))
+            if len(episodes_start_end) > 0:
+                reward_trajs_split = [
+                    reward_trajs[start : end + 1, env_ind]
+                    for env_ind, start, end in episodes_start_end
+                ]
+                num_episode_finished = len(reward_trajs_split)
+                episode_reward = np.array(
+                    [np.sum(reward_traj) for reward_traj in reward_trajs_split]
+                )
+                episode_best_reward = np.array(
+                    [
+                        np.max(reward_traj) / self.act_steps
+                        for reward_traj in reward_trajs_split
+                    ]
+                )
+                avg_episode_reward = np.mean(episode_reward)
+                avg_best_reward = np.mean(episode_best_reward)
+                success_rate = np.mean(
+                    episode_best_reward >= self.best_reward_threshold_for_success
+                )
+            else:
+                episode_reward = np.array([])
+                num_episode_finished = 0
+                avg_episode_reward = 0
+                avg_best_reward = 0
+                success_rate = 0
+                log.info("[WARNING] No episode completed within the iteration!")
+
+            # Update
+            if not eval_mode:
+
+                obs_trajs = np.array(deepcopy(obs_buffer))
+                action_trajs = np.array(deepcopy(action_buffer))
+                next_obs_trajs = np.array(deepcopy(next_obs_buffer))
+                reward_trajs = np.array(deepcopy(reward_buffer))
+                done_trajs = np.array(deepcopy(done_buffer))
+                first_trajs = np.array(deepcopy(first_buffer))
+
+                # flatten
+                obs_trajs = einops.rearrange(
+                    obs_trajs,
+                    "s e h d -> (s e) h d",
+                )
+                next_obs_trajs = einops.rearrange(
+                    next_obs_trajs,
+                    "s e h d -> (s e) h d",
+                )
+                action_trajs = einops.rearrange(
+                    action_trajs,
+                    "s e h d -> (s e) h d",
+                )
+                reward_trajs = reward_trajs.reshape(-1)
+                done_trajs = done_trajs.reshape(-1)
+                first_trajs = first_trajs.reshape(-1)
+
+                num_batch = int(
+                    self.n_steps * self.n_envs / self.batch_size * self.replay_ratio
+                )
+
+                for _ in range(num_batch):
+
+                    # Sample batch
+                    inds = np.random.choice(len(obs_trajs), self.batch_size)
+                    obs_b = torch.from_numpy(obs_trajs[inds]).float().to(self.device)
+                    next_obs_b = (
+                        torch.from_numpy(next_obs_trajs[inds]).float().to(self.device)
+                    )
+                    actions_b = (
+                        torch.from_numpy(action_trajs[inds]).float().to(self.device)
+                    )
+                    reward_b = (
+                        torch.from_numpy(reward_trajs[inds]).float().to(self.device)
+                    )
+                    done_b = torch.from_numpy(done_trajs[inds]).float().to(self.device)
+
+                    # update critic q function
+                    critic_loss = self.model.loss_critic(
+                        obs_b, next_obs_b, actions_b, reward_b, done_b, self.gamma
+                    )
+                    self.critic_optimizer.zero_grad()
+                    critic_loss.backward()
+                    self.critic_optimizer.step()
+
+                    # update target q function
+                    self.model.update_target_critic(self.critic_tau)
+
+                    loss_critic = critic_loss.detach()
+
+                    # Update policy with collected trajectories
+                    loss = self.model.loss_actor(
+                        obs_b,
+                        actions_b,
+                        self.q_grad_coeff,
+                    )
+                    self.actor_optimizer.zero_grad()
+                    loss.backward()
+                    if self.itr >= self.n_critic_warmup_itr:
+                        if self.max_grad_norm is not None:
+                            torch.nn.utils.clip_grad_norm_(
+                                self.model.actor.parameters(), self.max_grad_norm
+                            )
+                        self.actor_optimizer.step()
+
+            # Update lr
+            self.actor_lr_scheduler.step()
+            self.critic_lr_scheduler.step()
+
+            # Save model
+            if self.itr % self.save_model_freq == 0 or self.itr == self.n_train_itr - 1:
+                self.save_model()
+
+            # Log loss and save metrics
+            run_results.append(
+                {
+                    "itr": self.itr,
+                }
+            )
+            if self.itr % self.log_freq == 0:
+                if eval_mode:
+                    log.info(
+                        f"eval: success rate {success_rate:8.4f} | avg episode reward {avg_episode_reward:8.4f} | avg best reward {avg_best_reward:8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "success rate - eval": success_rate,
+                                "avg episode reward - eval": avg_episode_reward,
+                                "avg best reward - eval": avg_best_reward,
+                                "num episode - eval": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=False,
+                        )
+                    run_results[-1]["eval_success_rate"] = success_rate
+                    run_results[-1]["eval_episode_reward"] = avg_episode_reward
+                    run_results[-1]["eval_best_reward"] = avg_best_reward
+                else:
+                    log.info(
+                        f"{self.itr}: loss {loss:8.4f} | reward {avg_episode_reward:8.4f} |t:{timer():8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "loss": loss,
+                                "loss - critic": loss_critic,
+                                "avg episode reward - train": avg_episode_reward,
+                                "num episode - train": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=True,
+                        )
+                    run_results[-1]["loss"] = loss
+                    run_results[-1]["loss_critic"] = loss_critic
+                    run_results[-1]["train_episode_reward"] = avg_episode_reward
+                run_results[-1]["time"] = timer()
+                with open(self.result_path, "wb") as f:
+                    pickle.dump(run_results, f)
+            self.itr += 1

agent/finetune/train_rwr_diffusion_agent.py

@@ -0,0 +1,301 @@
+"""
+Reward-weighted regression (RWR) for diffusion policy.
+
+"""
+
+import os
+import pickle
+import numpy as np
+import torch
+import logging
+import wandb
+
+log = logging.getLogger(__name__)
+from util.timer import Timer
+from agent.finetune.train_agent import TrainAgent
+from util.scheduler import CosineAnnealingWarmupRestarts
+
+
+class TrainRWRDiffusionAgent(TrainAgent):
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+        # note the discount factor gamma here is applied to reward every act_steps, instead of every env step
+        self.gamma = cfg.train.gamma
+
+        # Build optimizer
+        self.optimizer = torch.optim.AdamW(
+            self.model.parameters(),
+            lr=cfg.train.lr,
+            weight_decay=cfg.train.weight_decay,
+        )
+        self.lr_scheduler = CosineAnnealingWarmupRestarts(
+            self.optimizer,
+            first_cycle_steps=cfg.train.lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=cfg.train.lr,
+            min_lr=cfg.train.lr_scheduler.min_lr,
+            warmup_steps=cfg.train.lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+
+        # Reward exponential
+        self.beta = cfg.train.beta
+
+        # Max weight for AWR
+        self.max_reward_weight = cfg.train.max_reward_weight
+
+        # Updates
+        self.update_epochs = cfg.train.update_epochs
+
+    def run(self):
+
+        # Start training loop
+        timer = Timer()
+        run_results = []
+        done_venv = np.zeros((1, self.n_envs))
+        while self.itr < self.n_train_itr:
+
+            # Prepare video paths for each envs --- only applies for the first set of episodes if allowing reset within iteration and each iteration has multiple episodes from one env
+            options_venv = [{} for _ in range(self.n_envs)]
+            if self.itr % self.render_freq == 0 and self.render_video:
+                for env_ind in range(self.n_render):
+                    options_venv[env_ind]["video_path"] = os.path.join(
+                        self.render_dir, f"itr-{self.itr}_trial-{env_ind}.mp4"
+                    )
+
+            # Define train or eval - all envs restart
+            eval_mode = self.itr % self.val_freq == 0 and not self.force_train
+            self.model.eval() if eval_mode else self.model.train()
+            firsts_trajs = np.zeros((self.n_steps + 1, self.n_envs))
+
+            # Reset env at the beginning of an iteration
+            if self.reset_at_iteration or eval_mode or last_itr_eval:
+                prev_obs_venv = self.reset_env_all(options_venv=options_venv)
+                firsts_trajs[0] = 1
+            else:
+                firsts_trajs[0] = (
+                    done_venv  # if done at the end of last iteration, then the envs are just reset
+                )
+            last_itr_eval = eval_mode
+            reward_trajs = np.empty((0, self.n_envs))
+
+            # Holders
+            obs_trajs = np.empty((0, self.n_envs, self.n_cond_step, self.obs_dim))
+            samples_trajs = np.empty(
+                (
+                    0,
+                    self.n_envs,
+                    self.horizon_steps,
+                    self.action_dim,
+                )
+            )
+
+            # Collect a set of trajectories from env
+            for step in range(self.n_steps):
+                if step % 10 == 0:
+                    print(f"Processed step {step} of {self.n_steps}")
+
+                # Select action
+                with torch.no_grad():
+                    samples = (
+                        self.model(
+                            cond=torch.from_numpy(prev_obs_venv)
+                            .float()
+                            .to(self.device),
+                            deterministic=eval_mode,
+                        )
+                        .cpu()
+                        .numpy()
+                    )  # n_env x horizon x act
+                action_venv = samples[:, : self.act_steps]
+                obs_trajs = np.vstack((obs_trajs, prev_obs_venv[None]))
+                samples_trajs = np.vstack((samples_trajs, samples[None]))
+
+                # Apply multi-step action
+                obs_venv, reward_venv, done_venv, info_venv = self.venv.step(
+                    action_venv
+                )
+                reward_trajs = np.vstack((reward_trajs, reward_venv[None]))
+                firsts_trajs[step + 1] = done_venv
+                prev_obs_venv = obs_venv
+
+            # Summarize episode reward --- this needs to be handled differently depending on whether the environment is reset after each iteration. Only count episodes that finish within the iteration.
+            episodes_start_end = []
+            for env_ind in range(self.n_envs):
+                env_steps = np.where(firsts_trajs[:, env_ind] == 1)[0]
+                for i in range(len(env_steps) - 1):
+                    start = env_steps[i]
+                    end = env_steps[i + 1]
+                    if end - start > 1:
+                        episodes_start_end.append((env_ind, start, end - 1))
+            if len(episodes_start_end) > 0:
+                # Compute transitions for completed trajectories
+                obs_trajs_split = [
+                    obs_trajs[start : end + 1, env_ind]
+                    for env_ind, start, end in episodes_start_end
+                ]
+                samples_trajs_split = [
+                    samples_trajs[start : end + 1, env_ind]
+                    for env_ind, start, end in episodes_start_end
+                ]
+                reward_trajs_split = [
+                    reward_trajs[start : end + 1, env_ind]
+                    for env_ind, start, end in episodes_start_end
+                ]
+                num_episode_finished = len(reward_trajs_split)
+
+                # Compute episode returns
+                discounted_reward_trajs_split = [
+                    [
+                        self.gamma**t * r
+                        for t, r in zip(
+                            list(range(end - start + 1)),
+                            reward_trajs[start : end + 1, env_ind],
+                        )
+                    ]
+                    for env_ind, start, end in episodes_start_end
+                ]
+                returns_trajs_split = [
+                    np.cumsum(y[::-1])[::-1] for y in discounted_reward_trajs_split
+                ]
+                returns_trajs_split = np.concatenate(returns_trajs_split)
+                episode_reward = np.array(
+                    [np.sum(reward_traj) for reward_traj in reward_trajs_split]
+                )
+                episode_best_reward = np.array(
+                    [
+                        np.max(reward_traj) / self.act_steps
+                        for reward_traj in reward_trajs_split
+                    ]
+                )
+                avg_episode_reward = np.mean(episode_reward)
+                avg_best_reward = np.mean(episode_best_reward)
+                success_rate = np.mean(
+                    episode_best_reward >= self.best_reward_threshold_for_success
+                )
+            else:
+                episode_reward = np.array([])
+                num_episode_finished = 0
+                avg_episode_reward = 0
+                avg_best_reward = 0
+                success_rate = 0
+                log.info("[WARNING] No episode completed within the iteration!")
+
+            # Update
+            if not eval_mode:
+
+                # Tensorize data and put them to device
+                # k for environment step
+                obs_k = (
+                    torch.tensor(np.concatenate(obs_trajs_split))
+                    .float()
+                    .to(self.device)
+                )
+
+                samples_k = (
+                    torch.tensor(np.concatenate(samples_trajs_split))
+                    .float()
+                    .to(self.device)
+                )
+
+                # Normalize reward
+                returns_trajs_split = (
+                    returns_trajs_split - np.mean(returns_trajs_split)
+                ) / (returns_trajs_split.std() + 1e-3)
+
+                rewards_k = (
+                    torch.tensor(returns_trajs_split)
+                    .float()
+                    .to(self.device)
+                    .reshape(-1)
+                )
+
+                rewards_k_scaled = torch.exp(self.beta * rewards_k)
+                rewards_k_scaled.clamp_(max=self.max_reward_weight)
+
+                # rewards_k_scaled = rewards_k_scaled / rewards_k_scaled.mean()
+
+                # Update policy and critic
+                total_steps = len(rewards_k_scaled)
+                inds_k = np.arange(total_steps)
+                for _ in range(self.update_epochs):
+
+                    # for each epoch, go through all data in batches
+                    np.random.shuffle(inds_k)
+                    num_batch = max(1, total_steps // self.batch_size)  # skip last ones
+                    for batch in range(num_batch):
+                        start = batch * self.batch_size
+                        end = start + self.batch_size
+                        inds_b = inds_k[start:end]  # b for batch
+                        obs_b = obs_k[inds_b]
+                        samples_b = samples_k[inds_b]
+                        rewards_b = rewards_k_scaled[inds_b]
+
+                        # Update policy with collected trajectories
+                        loss = self.model.loss(
+                            samples_b,
+                            obs_b,
+                            rewards_b,
+                        )
+                        self.optimizer.zero_grad()
+                        loss.backward()
+                        if self.max_grad_norm is not None:
+                            torch.nn.utils.clip_grad_norm_(
+                                self.model.parameters(), self.max_grad_norm
+                            )
+                        self.optimizer.step()
+
+            # Update lr
+            self.lr_scheduler.step()
+
+            # Save model
+            if self.itr % self.save_model_freq == 0 or self.itr == self.n_train_itr - 1:
+                self.save_model()
+
+            # Log loss and save metrics
+            run_results.append(
+                {
+                    "itr": self.itr,
+                }
+            )
+            if self.itr % self.log_freq == 0:
+                if eval_mode:
+                    log.info(
+                        f"eval: success rate {success_rate:8.4f} | avg episode reward {avg_episode_reward:8.4f} | avg best reward {avg_best_reward:8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "success rate - eval": success_rate,
+                                "avg episode reward - eval": avg_episode_reward,
+                                "avg best reward - eval": avg_best_reward,
+                                "num episode - eval": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=False,
+                        )
+                    run_results[-1]["eval_success_rate"] = success_rate
+                    run_results[-1]["eval_episode_reward"] = avg_episode_reward
+                    run_results[-1]["eval_best_reward"] = avg_best_reward
+                else:
+                    log.info(
+                        f"{self.itr}: loss {loss:8.4f} | reward {avg_episode_reward:8.4f} |t:{timer():8.4f}"
+                    )
+                    if self.use_wandb:
+                        wandb.log(
+                            {
+                                "loss": loss,
+                                "avg episode reward - train": avg_episode_reward,
+                                "num episode - train": num_episode_finished,
+                            },
+                            step=self.itr,
+                            commit=True,
+                        )
+                    run_results[-1]["loss"] = loss
+                    run_results[-1]["train_episode_reward"] = avg_episode_reward
+                run_results[-1]["time"] = timer()
+                with open(self.result_path, "wb") as f:
+                    pickle.dump(run_results, f)
+            self.itr += 1

agent/pretrain/train_agent.py

@@ -0,0 +1,168 @@
+"""
+Parent pre-training agent class.
+
+"""
+
+import os
+import random
+import numpy as np
+from omegaconf import OmegaConf
+import torch
+import hydra
+import logging
+import wandb
+from copy import deepcopy
+
+log = logging.getLogger(__name__)
+from util.scheduler import CosineAnnealingWarmupRestarts
+
+DEVICE = "cuda:0"
+
+
+def to_device(x, device=DEVICE):
+    if torch.is_tensor(x):
+        return x.to(device)
+    elif type(x) is dict:
+        return {k: to_device(v, device) for k, v in x.items()}
+    else:
+        print(f"Unrecognized type in `to_device`: {type(x)}")
+
+
+def batch_to_device(batch, device="cuda:0"):
+    vals = [to_device(getattr(batch, field), device) for field in batch._fields]
+    return type(batch)(*vals)
+
+
+class EMA:
+    """
+    Empirical moving average
+
+    """
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.beta = cfg.decay
+
+    def update_model_average(self, ma_model, current_model):
+        for current_params, ma_params in zip(
+            current_model.parameters(), ma_model.parameters()
+        ):
+            old_weight, up_weight = ma_params.data, current_params.data
+            ma_params.data = self.update_average(old_weight, up_weight)
+
+    def update_average(self, old, new):
+        if old is None:
+            return new
+        return old * self.beta + (1 - self.beta) * new
+
+
+class PreTrainAgent:
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.seed = cfg.get("seed", 42)
+        random.seed(self.seed)
+        np.random.seed(self.seed)
+        torch.manual_seed(self.seed)
+
+        # Wandb
+        self.use_wandb = cfg.wandb is not None
+        if cfg.wandb is not None:
+            wandb.init(
+                entity=cfg.wandb.entity,
+                project=cfg.wandb.project,
+                name=cfg.wandb.run,
+                config=OmegaConf.to_container(cfg, resolve=True),
+            )
+
+        # Build model
+        self.model = hydra.utils.instantiate(cfg.model)
+        self.ema = EMA(cfg.ema)
+        self.ema_model = deepcopy(self.model)
+
+        # Training params
+        self.n_epochs = cfg.train.n_epochs
+        self.batch_size = cfg.train.batch_size
+        self.update_ema_freq = cfg.train.update_ema_freq
+        self.epoch_start_ema = cfg.train.epoch_start_ema
+        self.val_freq = cfg.train.get("val_freq", 100)
+
+        # Logging, checkpoints
+        self.logdir = cfg.logdir
+        self.checkpoint_dir = os.path.join(self.logdir, "checkpoint")
+        os.makedirs(self.checkpoint_dir, exist_ok=True)
+        self.log_freq = cfg.train.get("log_freq", 1)
+        self.save_model_freq = cfg.train.save_model_freq
+
+        # Build dataset
+        self.dataset_train = hydra.utils.instantiate(cfg.train_dataset)
+        self.dataloader_train = torch.utils.data.DataLoader(
+            self.dataset_train,
+            batch_size=self.batch_size,
+            num_workers=4 if self.dataset_train.device == "cpu" else 0,
+            shuffle=True,
+            pin_memory=True if self.dataset_train.device == "cpu" else False,
+        )
+        self.dataloader_val = None
+        if "train_split" in cfg.train and cfg.train.train_split < 1:
+            val_indices = self.dataset_train.set_train_val_split(cfg.train.train_split)
+            self.dataset_val = deepcopy(self.dataset_train)
+            self.dataset_val.set_indices(val_indices)
+            self.dataloader_val = torch.utils.data.DataLoader(
+                self.dataset_val,
+                batch_size=self.batch_size,
+                num_workers=4 if self.dataset_val.device == "cpu" else 0,
+                shuffle=True,
+                pin_memory=True if self.dataset_val.device == "cpu" else False,
+            )
+        self.optimizer = torch.optim.AdamW(
+            self.model.parameters(),
+            lr=cfg.train.learning_rate,
+            weight_decay=cfg.train.weight_decay,
+        )
+        self.lr_scheduler = CosineAnnealingWarmupRestarts(
+            self.optimizer,
+            first_cycle_steps=cfg.train.lr_scheduler.first_cycle_steps,
+            cycle_mult=1.0,
+            max_lr=cfg.train.learning_rate,
+            min_lr=cfg.train.lr_scheduler.min_lr,
+            warmup_steps=cfg.train.lr_scheduler.warmup_steps,
+            gamma=1.0,
+        )
+        self.reset_parameters()
+
+    def run(self):
+        raise NotImplementedError
+
+    def reset_parameters(self):
+        self.ema_model.load_state_dict(self.model.state_dict())
+
+    def step_ema(self):
+        if self.epoch < self.epoch_start_ema:
+            self.reset_parameters()
+            return
+        self.ema.update_model_average(self.ema_model, self.model)
+
+    def save_model(self):
+        """
+        saves model and ema to disk;
+        """
+        data = {
+            "epoch": self.epoch,
+            "model": self.model.state_dict(),
+            "ema": self.ema_model.state_dict(),
+        }
+        savepath = os.path.join(self.checkpoint_dir, f"state_{self.epoch}.pt")
+        torch.save(data, savepath)
+        log.info(f"Saved model to {savepath}")
+
+    def load(self, epoch):
+        """
+        loads model and ema from disk
+        """
+        loadpath = os.path.join(self.checkpoint_dir, f"state_{epoch}.pt")
+        data = torch.load(loadpath, weights_only=True)
+
+        self.epoch = data["epoch"]
+        self.model.load_state_dict(data["model"])
+        self.ema_model.load_state_dict(data["ema"])

agent/pretrain/train_diffusion_agent.py

@@ -0,0 +1,83 @@
+"""
+Pre-training diffusion policy
+
+"""
+
+import logging
+import wandb
+import numpy as np
+
+log = logging.getLogger(__name__)
+from util.timer import Timer
+from agent.pretrain.train_agent import PreTrainAgent, batch_to_device
+
+
+class TrainDiffusionAgent(PreTrainAgent):
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+    def run(self):
+
+        timer = Timer()
+        self.epoch = 1
+        for _ in range(self.n_epochs):
+
+            # train
+            loss_train_epoch = []
+            for batch_train in self.dataloader_train:
+                if self.dataset_train.device == "cpu":
+                    batch_train = batch_to_device(batch_train)
+
+                self.model.train()
+                loss_train = self.model.loss(*batch_train)
+                loss_train.backward()
+                loss_train_epoch.append(loss_train.item())
+
+                self.optimizer.step()
+                self.optimizer.zero_grad()
+            loss_train = np.mean(loss_train_epoch)
+
+            # validate
+            loss_val_epoch = []
+            if self.dataloader_val is not None and self.epoch % self.val_freq == 0:
+                self.model.eval()
+                for batch_val in self.dataloader_val:
+                    if self.dataset_val.device == "cpu":
+                        batch_val = batch_to_device(batch_val)
+                    loss_val, infos_val = self.model.loss(*batch_val)
+                    loss_val_epoch.append(loss_val.item())
+                self.model.train()
+            loss_val = np.mean(loss_val_epoch) if len(loss_val_epoch) > 0 else None
+
+            # update lr
+            self.lr_scheduler.step()
+
+            # update ema
+            if self.epoch % self.update_ema_freq == 0:
+                self.step_ema()
+
+            # save model
+            if self.epoch % self.save_model_freq == 0 or self.epoch == self.n_epochs:
+                self.save_model()
+
+            # log loss
+            if self.epoch % self.log_freq == 0:
+                log.info(
+                    f"{self.epoch}: train loss {loss_train:8.4f} | t:{timer():8.4f}"
+                )
+                if self.use_wandb:
+                    if loss_val is not None:
+                        wandb.log(
+                            {"loss - val": loss_val}, step=self.epoch, commit=False
+                        )
+                    wandb.log(
+                        {
+                            "loss - train": loss_train,
+                        },
+                        step=self.epoch,
+                        commit=True,
+                    )
+
+            # count
+            self.epoch += 1

agent/pretrain/train_gaussian_agent.py

@@ -0,0 +1,99 @@
+"""
+Pre-training Gaussian/GMM policy
+
+"""
+
+import logging
+import wandb
+import numpy as np
+
+log = logging.getLogger(__name__)
+from util.timer import Timer
+from agent.pretrain.train_agent import PreTrainAgent, batch_to_device
+
+
+class TrainGaussianAgent(PreTrainAgent):
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+        # Entropy bonus - not used right now since using fixed_std
+        self.ent_coef = cfg.train.get("ent_coef", 0)
+
+    def run(self):
+
+        timer = Timer()
+        self.epoch = 1
+        for _ in range(self.n_epochs):
+
+            # train
+            loss_train_epoch = []
+            ent_train_epoch = []
+            for batch_train in self.dataloader_train:
+                if self.dataset_train.device == "cpu":
+                    batch_train = batch_to_device(batch_train)
+
+                self.model.train()
+                loss_train, infos_train = self.model.loss(
+                    *batch_train,
+                    ent_coef=self.ent_coef,
+                )
+                loss_train.backward()
+                loss_train_epoch.append(loss_train.item())
+                ent_train_epoch.append(infos_train["entropy"].item())
+
+                self.optimizer.step()
+                self.optimizer.zero_grad()
+            loss_train = np.mean(loss_train_epoch)
+            ent_train = np.mean(ent_train_epoch)
+
+            # validate
+            loss_val_epoch = []
+            if self.dataloader_val is not None and self.epoch % self.val_freq == 0:
+                self.model.eval()
+                for batch_val in self.dataloader_val:
+                    if self.dataset_val.device == "cpu":
+                        batch_val = batch_to_device(batch_val)
+                    loss_val, infos_val = self.model.loss(
+                        *batch_val,
+                        ent_coef=self.ent_coef,
+                    )
+                    loss_val_epoch.append(loss_val.item())
+                self.model.train()
+            loss_val = np.mean(loss_val_epoch) if len(loss_val_epoch) > 0 else None
+
+            # update lr
+            self.lr_scheduler.step()
+
+            # update ema
+            if self.epoch % self.update_ema_freq == 0:
+                self.step_ema()
+
+            # save model
+            if self.epoch % self.save_model_freq == 0 or self.epoch == self.n_epochs:
+                self.save_model()
+
+            # log loss
+            if self.epoch % self.log_freq == 0:
+                infos_str = " | ".join(
+                    [f"{key}: {val:8.4f}" for key, val in infos_train.items()]
+                )
+                log.info(
+                    f"{self.epoch}: train loss {loss_train:8.4f} | {infos_str} | t:{timer():8.4f}"
+                )
+                if self.use_wandb:
+                    if loss_val is not None:
+                        wandb.log(
+                            {"loss - val": loss_val}, step=self.epoch, commit=False
+                        )
+                    wandb.log(
+                        {
+                            "loss - train": loss_train,
+                            "entropy - train": ent_train,
+                        },
+                        step=self.epoch,
+                        commit=True,
+                    )
+
+            # count
+            self.epoch += 1