release

script/dataset/README.md

@@ -0,0 +1,3 @@
+## Data processing scripts
+
+These are some scripts used for processing the raw datasets from the benchmarks. We already pre-processed them and provide the final datasets. These scripts are for information only.

script/dataset/filter_d3il_avoid_data.py

@@ -0,0 +1,357 @@
+"""
+Filter avoid data based on modes.
+
+Trajectories are normalized with filtered data, not the original data.
+"""
+
+import os
+import numpy as np
+from tqdm import tqdm
+import pickle
+import random
+import matplotlib.pyplot as plt
+from copy import deepcopy
+
+from agent.dataset.d3il_dataset.avoiding_dataset import Avoiding_Dataset
+
+
+def make_dataset(
+    load_path,
+    save_dir,
+    save_name_prefix,
+    val_split,
+    desired_modes,
+    desired_mode_ratios,
+    required_modes,
+    avoid_modes,
+):
+    print("Desired modes:", desired_modes)
+    print("Required modes:", required_modes)
+    print("Avoid modes:", avoid_modes)
+    print("Desired mode ratios:", desired_mode_ratios)
+    demo_dataset = Avoiding_Dataset(
+        load_path,
+        action_dim=2,
+        obs_dim=4,
+        max_len_data=200,
+    )
+    # from avoiding env
+    level_distance = 0.18
+    obstacle_offset = 0.075
+    l1_ypos = -0.1
+    l2_ypos = -0.1 + level_distance
+    l3_ypos = -0.1 + 2 * level_distance
+    # goal_ypos = -0.1 + 2.5 * level_distance
+    l1_xpos = 0.5
+    l2_top_xpos = 0.5 - obstacle_offset
+    l2_bottom_xpos = 0.5 + obstacle_offset
+    l3_top_xpos = 0.5 - 2 * obstacle_offset
+    l3_mid_xpos = 0.5
+    l3_bottom_xpos = 0.5 + 2 * obstacle_offset
+
+    def check_mode(x):
+        r_x_pos = x[0]
+        r_y_pos = x[1]
+        mode_encoding = np.zeros((9))
+        if r_y_pos - 0.01 <= l1_ypos <= r_y_pos + 0.01:
+            if r_x_pos < l1_xpos:
+                mode_encoding[0] = 1
+            elif r_x_pos > l1_xpos:
+                mode_encoding[1] = 1
+
+        if r_y_pos - 0.01 <= l2_ypos <= r_y_pos + 0.01:
+            if r_x_pos < l2_top_xpos:
+                mode_encoding[2] = 1
+            elif l2_top_xpos < r_x_pos < l2_bottom_xpos:
+                mode_encoding[3] = 1
+            elif r_x_pos > l2_bottom_xpos:
+                mode_encoding[4] = 1
+
+        # if r_y_pos - 0.015 <= self.l3_ypos and (not self.l3_passed):
+        if r_y_pos >= l3_ypos:
+            if r_x_pos < l3_top_xpos:
+                mode_encoding[5] = 1
+            if l3_top_xpos < r_x_pos < l3_mid_xpos:
+                mode_encoding[6] = 1
+            elif l3_mid_xpos < r_x_pos < l3_bottom_xpos:
+                mode_encoding[7] = 1
+            elif r_x_pos > l3_top_xpos:
+                mode_encoding[8] = 1
+        return mode_encoding
+
+    # extract length of each trajectory in the file
+    full_traj_lengths = []
+    full_actions = demo_dataset.actions
+    full_obs = demo_dataset.observations
+    masks = demo_dataset.masks
+    action_dim = full_actions.shape[2]
+    obs_dim = full_obs.shape[2]
+    for ep in range(masks.shape[0]):
+        full_traj_lengths.append(int(masks[ep].sum().item()))
+    full_traj_lengths = np.array(full_traj_lengths)
+
+    # take the max and min of obs and action
+    obs_min = np.zeros((obs_dim))
+    obs_max = np.zeros((obs_dim))
+    action_min = np.zeros((action_dim))
+    action_max = np.zeros((action_dim))
+    chosen_indices = []
+    for i in tqdm(range(len(masks))):
+        T = full_traj_lengths[i]
+        obs_traj = full_obs[i, :T].numpy()
+        action_traj = full_actions[i, :T].numpy()
+
+        # check if trajectory pass through desired hole
+        flag_desired = False
+        flag_required = [False for _ in required_modes] if required_modes else [True]
+        flag_avoid = False
+        for ob in obs_traj:
+            modes = check_mode(ob)
+            if any(modes[desired] for desired in desired_modes):
+                desired_mode_idx = np.argmax(
+                    [modes[desired] for desired in desired_modes]
+                )
+                flag_desired = True
+            if any(modes[avoid] for avoid in avoid_modes):
+                flag_avoid = True
+                break
+            for j, required in enumerate(required_modes):
+                if modes[required]:
+                    flag_required[j] = True
+        if flag_avoid or not flag_desired or not all(flag_required):
+            continue
+        if desired_mode_ratios:
+            if random.random() > desired_mode_ratios[desired_mode_idx]:
+                continue
+        chosen_indices.append(i)
+
+        obs_min = np.min(np.vstack((obs_min, np.min(obs_traj, axis=0))), axis=0)
+        obs_max = np.max(np.vstack((obs_max, np.max(obs_traj, axis=0))), axis=0)
+        action_min = np.min(
+            np.vstack((action_min, np.min(action_traj, axis=0))), axis=0
+        )
+        action_max = np.max(
+            np.vstack((action_max, np.max(action_traj, axis=0))), axis=0
+        )
+    if len(chosen_indices) == 0:
+        raise ValueError("No data found for the desired/required modes")
+    chosen_indices = np.array(chosen_indices)
+    traj_lengths = full_traj_lengths[chosen_indices]
+    actions = demo_dataset.actions[chosen_indices]
+    obs = demo_dataset.observations[chosen_indices]
+    max_traj_length = np.max(traj_lengths)
+
+    # split indices in train and val
+    num_traj = len(traj_lengths)
+    num_train = int(num_traj * (1 - val_split))
+    train_indices = random.sample(range(num_traj), k=num_train)
+
+    logger.info("\n========== Basic Info ===========")
+    logger.info("total transitions: {}".format(np.sum(traj_lengths)))
+    logger.info("total trajectories: {}".format(len(traj_lengths)))
+    logger.info(
+        f"traj length mean/std: {np.mean(traj_lengths)}, {np.std(traj_lengths)}"
+    )
+    logger.info(f"traj length min/max: {np.min(traj_lengths)}, {np.max(traj_lengths)}")
+    logger.info(f"obs min: {obs_min}")
+    logger.info(f"obs max: {obs_max}")
+    logger.info(f"action min: {action_min}")
+    logger.info(f"action max: {action_max}")
+
+    # do over all indices
+    out_train = {}
+    keys = [
+        "observations",
+        "actions",
+        "rewards",
+    ]
+    total_timesteps = actions.shape[1]
+    out_train["observations"] = np.empty((0, total_timesteps, obs_dim))
+    out_train["actions"] = np.empty((0, total_timesteps, action_dim))
+    out_train["rewards"] = np.empty((0, total_timesteps))
+    out_train["traj_length"] = []
+    out_val = deepcopy(out_train)
+    for i in tqdm(range(len(traj_lengths))):
+        if i in train_indices:
+            out = out_train
+        else:
+            out = out_val
+        T = traj_lengths[i]
+        obs_traj = obs[i].numpy()
+        action_traj = actions[i].numpy()
+
+        # scale to [-1, 1] for both ob and action
+        obs_traj = 2 * (obs_traj - obs_min) / (obs_max - obs_min + 1e-6) - 1
+        action_traj = (
+            2 * (action_traj - action_min) / (action_max - action_min + 1e-6) - 1
+        )
+
+        # get episode length
+        traj_length = T
+        out["traj_length"].append(traj_length)
+
+        # extract
+        rewards = np.zeros(total_timesteps)  # no reward from d3il dataset
+        data_traj = {
+            "observations": obs_traj,
+            "actions": action_traj,
+            "rewards": rewards,
+        }
+        for key in keys:
+            traj = data_traj[key]
+            out[key] = np.vstack((out[key], traj[None]))
+
+    # plot all trajectories and save in a figure
+    def plot(out, name):
+        def get_obj_xy_list():
+            mid_pos = 0.5
+            offset = 0.075
+            first_level_y = -0.1
+            level_distance = 0.18
+            return [
+                [mid_pos, first_level_y],
+                [mid_pos - offset, first_level_y + level_distance],
+                [mid_pos + offset, first_level_y + level_distance],
+                [mid_pos - 2 * offset, first_level_y + 2 * level_distance],
+                [mid_pos, first_level_y + 2 * level_distance],
+                [mid_pos + 2 * offset, first_level_y + 2 * level_distance],
+            ]
+
+        pillar_xys = get_obj_xy_list()
+        fig = plt.figure()
+        all_trajs = out["observations"]  # num x timestep x obs
+        for traj, traj_length in zip(all_trajs, out["traj_length"]):
+            # unnormalize
+            traj = (traj + 1) / 2  # [-1, 1] -> [0, 1]
+            traj = traj * (obs_max - obs_min) + obs_min
+            plt.plot(
+                traj[:traj_length, 2], traj[:traj_length, 3], color=(0.3, 0.3, 0.3)
+            )
+        plt.axhline(y=0.4, color=np.array([31, 119, 180]) / 255, linestyle="-")
+        for xy in pillar_xys:
+            circle = plt.Circle(xy, 0.01, color=(0.0, 0.0, 0.0), fill=True)
+            plt.gca().add_patch(circle)
+        plt.xlabel("X pos")
+        plt.ylabel("Y pos")
+        plt.xlim([0.2, 0.8])
+        plt.ylim([-0.3, 0.5])
+        ax = plt.gca()
+        ax.set_aspect("equal", adjustable="box")
+        ax.set_facecolor("white")
+        plt.savefig(os.path.join(save_dir, name))
+        plt.close(fig)
+
+    plot(out_train, name="train-trajs.png")
+    plot(out_val, name="val-trajs.png")
+
+    # Save to np file
+    save_train_path = os.path.join(save_dir, save_name_prefix + "train.pkl")
+    save_val_path = os.path.join(save_dir, save_name_prefix + "val.pkl")
+    with open(save_train_path, "wb") as f:
+        pickle.dump(out_train, f)
+    with open(save_val_path, "wb") as f:
+        pickle.dump(out_val, f)
+    normalization_save_path = os.path.join(
+        save_dir, save_name_prefix + "normalization.npz"
+    )
+    np.savez(
+        normalization_save_path,
+        obs_min=obs_min,
+        obs_max=obs_max,
+        action_min=action_min,
+        action_max=action_max,
+    )
+
+    # debug
+    logger.info("\n========== Final ===========")
+    logger.info(
+        f"Train - Number of episodes and transitions: {len(out_train['traj_length'])}, {np.sum(out_train['traj_length'])}"
+    )
+    logger.info(
+        f"Val - Number of episodes and transitions: {len(out_val['traj_length'])}, {np.sum(out_val['traj_length'])}"
+    )
+    logger.info(
+        f"Train - Mean/Std trajectory length: {np.mean(out_train['traj_length'])}, {np.std(out_train['traj_length'])}"
+    )
+    logger.info(
+        f"Train - Max/Min trajectory length: {np.max(out_train['traj_length'])}, {np.min(out_train['traj_length'])}"
+    )
+    if val_split > 0:
+        logger.info(
+            f"Val - Mean/Std trajectory length: {np.mean(out_val['traj_length'])}, {np.std(out_val['traj_length'])}"
+        )
+        logger.info(
+            f"Val - Max/Min trajectory length: {np.max(out_val['traj_length'])}, {np.min(out_val['traj_length'])}"
+        )
+    for obs_dim_ind in range(obs_dim):
+        obs = out_train["observations"][:, :, obs_dim_ind]
+        logger.info(
+            f"Train - Obs dim {obs_dim_ind+1} mean {np.mean(obs)} std {np.std(obs)} min {np.min(obs)} max {np.max(obs)}"
+        )
+    for action_dim_ind in range(action_dim):
+        action = out_train["actions"][:, :, action_dim_ind]
+        logger.info(
+            f"Train - Action dim {action_dim_ind+1} mean {np.mean(action)} std {np.std(action)} min {np.min(action)} max {np.max(action)}"
+        )
+    if val_split > 0:
+        for obs_dim_ind in range(obs_dim):
+            obs = out_val["observations"][:, :, obs_dim_ind]
+            logger.info(
+                f"Val - Obs dim {obs_dim_ind+1} mean {np.mean(obs)} std {np.std(obs)} min {np.min(obs)} max {np.max(obs)}"
+            )
+        for action_dim_ind in range(action_dim):
+            action = out_val["actions"][:, :, action_dim_ind]
+            logger.info(
+                f"Val - Action dim {action_dim_ind+1} mean {np.mean(action)} std {np.std(action)} min {np.min(action)} max {np.max(action)}"
+            )
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--load_path", type=str, default=".")
+    parser.add_argument("--save_dir", type=str, default=".")
+    parser.add_argument("--save_name_prefix", type=str, default="")
+    parser.add_argument("--val_split", type=float, default="0.2")
+    parser.add_argument("--desired_modes", nargs="+", type=int)
+    parser.add_argument("--desired_mode_ratios", nargs="+", type=float, default=[])
+    parser.add_argument("--required_modes", nargs="+", type=int, default=[])
+    parser.add_argument("--avoid_modes", nargs="+", type=int, default=[])
+    args = parser.parse_args()
+    if len(args.desired_mode_ratios) > 0:
+        assert len(args.desired_modes) == len(
+            args.desired_mode_ratios
+        ), "Desired modes and desired mode ratios should have the same length"
+
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    import logging
+    import datetime
+
+    os.makedirs(args.save_dir, exist_ok=True)
+    log_path = os.path.join(
+        args.save_dir,
+        args.save_name_prefix
+        + f"{datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')}.log",
+    )
+    logger = logging.getLogger("get_D4RL_dataset")
+    logger.setLevel(logging.INFO)
+    file_handler = logging.FileHandler(log_path)
+    file_handler.setLevel(logging.INFO)  # Set the minimum level for this handler
+    formatter = logging.Formatter(
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    )
+    file_handler.setFormatter(formatter)
+    logger.addHandler(file_handler)
+
+    make_dataset(
+        args.load_path,
+        args.save_dir,
+        args.save_name_prefix,
+        args.val_split,
+        args.desired_modes,
+        args.desired_mode_ratios,
+        args.required_modes,
+        args.avoid_modes,
+    )

script/dataset/get_d4rl_dataset.py

@@ -0,0 +1,247 @@
+"""
+Download D4RL dataset and save it into our custom format so it can be loaded for diffusion training.
+
+"""
+
+import os
+import logging
+import gym
+import random
+from copy import deepcopy
+import numpy as np
+from tqdm import tqdm
+import pickle
+
+import d4rl.gym_mujoco  # Import required to register environments
+
+
+def make_dataset(env_name, save_dir, save_name_prefix, val_split, logger):
+    # Create the environment
+    env = gym.make(env_name)
+
+    # d4rl abides by the OpenAI gym interface
+    env.reset()
+    env.step(env.action_space.sample())
+
+    # Each task is associated with a dataset
+    # dataset contains observations, actions, rewards, terminals, and infos
+    dataset = env.get_dataset()
+    logger.info("\n========== Basic Info ===========")
+    logger.info(f"Keys in the dataset: {dataset.keys()}")
+    logger.info(f"Observation shape: {dataset['observations'].shape}")
+    logger.info(f"Action shape: {dataset['actions'].shape}")
+    terminal_indices = np.argwhere(dataset["terminals"])[:, 0]
+    timeout_indices = np.argwhere(dataset["timeouts"])[:, 0]
+    obs_dim = dataset["observations"].shape[1]
+    action_dim = dataset["actions"].shape[1]
+    done_indices = np.concatenate([terminal_indices, timeout_indices])
+    done_indices = np.sort(done_indices)
+    traj_lengths = []
+    prev_index = 0
+    for i in tqdm(range(len(done_indices))):
+        # get episode length
+        cur_index = done_indices[i]
+        traj_lengths.append(cur_index - prev_index + 1)
+        prev_index = cur_index + 1
+    obs_min = np.min(dataset["observations"], axis=0)
+    obs_max = np.max(dataset["observations"], axis=0)
+    action_min = np.min(dataset["actions"], axis=0)
+    action_max = np.max(dataset["actions"], axis=0)
+    max_episode_steps = max(traj_lengths)
+    logger.info("total transitions: {}".format(np.sum(traj_lengths)))
+    logger.info("total trajectories: {}".format(len(traj_lengths)))
+    logger.info(
+        f"traj length mean/std: {np.mean(traj_lengths)}, {np.std(traj_lengths)}"
+    )
+    logger.info(f"traj length min/max: {np.min(traj_lengths)}, {np.max(traj_lengths)}")
+    logger.info(f"obs min: {obs_min}")
+    logger.info(f"obs max: {obs_max}")
+    logger.info(f"action min: {action_min}")
+    logger.info(f"action max: {action_max}")
+
+    # Subsample episodes by taking the first ones
+    if args.max_episodes > 0:
+        traj_lengths = traj_lengths[: args.max_episodes]
+        done_indices = done_indices[: args.max_episodes]
+        max_episode_steps = max(traj_lengths)
+
+    # split indices in train and val
+    num_traj = len(traj_lengths)
+    num_train = int(num_traj * (1 - val_split))
+    train_indices = random.sample(range(num_traj), k=num_train)
+
+    # do over all indices
+    out_train = {}
+    keys = [
+        "observations",
+        "actions",
+        "rewards",
+    ]
+    out_train["observations"] = np.empty(
+        (0, max_episode_steps, dataset["observations"].shape[-1])
+    )
+    out_train["actions"] = np.empty(
+        (0, max_episode_steps, dataset["actions"].shape[-1])
+    )
+    out_train["rewards"] = np.empty((0, max_episode_steps))
+    out_train["traj_length"] = []
+    out_val = deepcopy(out_train)
+    prev_index = 0
+    train_episode_reward_all = []
+    val_episode_reward_all = []
+    for i in tqdm(range(len(done_indices))):
+        if i in train_indices:
+            out = out_train
+            episode_reward_all = train_episode_reward_all
+        else:
+            out = out_val
+            episode_reward_all = val_episode_reward_all
+
+        # get episode length
+        cur_index = done_indices[i]
+        traj_length = cur_index - prev_index + 1
+
+        # Skip if the episode has no reward
+        if np.sum(dataset["rewards"][prev_index : cur_index + 1]) > 0:
+            out["traj_length"].append(traj_length)
+
+            # apply padding to make all episodes have the same max steps
+            for key in keys:
+                traj = dataset[key][prev_index : cur_index + 1]
+
+                # also scale
+                if key == "observations":
+                    traj = 2 * (traj - obs_min) / (obs_max - obs_min + 1e-6) - 1
+                elif key == "actions":
+                    traj = (
+                        2 * (traj - action_min) / (action_max - action_min + 1e-6) - 1
+                    )
+
+                if traj.ndim == 1:
+                    traj = np.pad(
+                        traj,
+                        (0, max_episode_steps - len(traj)),
+                        mode="constant",
+                        constant_values=0,
+                    )
+                else:
+                    traj = np.pad(
+                        traj,
+                        ((0, max_episode_steps - traj.shape[0]), (0, 0)),
+                        mode="constant",
+                        constant_values=0,
+                    )
+                out[key] = np.vstack((out[key], traj[None]))
+
+            # check reward
+            episode_reward_all.append(np.sum(out["rewards"][-1]))
+        else:
+            print(f"skipping {i} / {len(done_indices)}")
+
+        # update prev index
+        prev_index = cur_index + 1
+
+    # Save to np file
+    save_train_path = os.path.join(save_dir, save_name_prefix + "train.pkl")
+    save_val_path = os.path.join(save_dir, save_name_prefix + "val.pkl")
+    with open(save_train_path, "wb") as f:
+        pickle.dump(out_train, f)
+    with open(save_val_path, "wb") as f:
+        pickle.dump(out_val, f)
+    normalization_save_path = os.path.join(
+        save_dir, save_name_prefix + "normalization.npz"
+    )
+    np.savez(
+        normalization_save_path,
+        obs_min=obs_min,
+        obs_max=obs_max,
+        action_min=action_min,
+        action_max=action_max,
+    )
+
+    # debug
+    logger.info("\n========== Final ===========")
+    logger.info(
+        f"Train - Number of episodes and transitions: {len(out_train['traj_length'])}, {np.sum(out_train['traj_length'])}"
+    )
+    logger.info(
+        f"Val - Number of episodes and transitions: {len(out_val['traj_length'])}, {np.sum(out_val['traj_length'])}"
+    )
+    logger.info(
+        f"Train - Mean/Std trajectory length: {np.mean(out_train['traj_length'])}, {np.std(out_train['traj_length'])}"
+    )
+    logger.info(
+        f"Train - Max/Min trajectory length: {np.max(out_train['traj_length'])}, {np.min(out_train['traj_length'])}"
+    )
+    if val_split > 0:
+        logger.info(
+            f"Val - Mean/Std trajectory length: {np.mean(out_val['traj_length'])}, {np.std(out_val['traj_length'])}"
+        )
+        logger.info(
+            f"Val - Max/Min trajectory length: {np.max(out_val['traj_length'])}, {np.min(out_val['traj_length'])}"
+        )
+    logger.info(
+        f"Train - Mean/Std episode reward: {np.mean(train_episode_reward_all)},  {np.std(train_episode_reward_all)}"
+    )
+    if val_split > 0:
+        logger.info(
+            f"Val - Mean/Std episode reward: {np.mean(val_episode_reward_all)},  {np.std(val_episode_reward_all)}"
+        )
+    for obs_dim_ind in range(obs_dim):
+        obs = out_train["observations"][:, :, obs_dim_ind]
+        logger.info(
+            f"Train - Obs dim {obs_dim_ind+1} mean {np.mean(obs)} std {np.std(obs)} min {np.min(obs)} max {np.max(obs)}"
+        )
+    for action_dim_ind in range(action_dim):
+        action = out_train["actions"][:, :, action_dim_ind]
+        logger.info(
+            f"Train - Action dim {action_dim_ind+1} mean {np.mean(action)} std {np.std(action)} min {np.min(action)} max {np.max(action)}"
+        )
+    if val_split > 0:
+        for obs_dim_ind in range(obs_dim):
+            obs = out_val["observations"][:, :, obs_dim_ind]
+            logger.info(
+                f"Val - Obs dim {obs_dim_ind+1} mean {np.mean(obs)} std {np.std(obs)} min {np.min(obs)} max {np.max(obs)}"
+            )
+        for action_dim_ind in range(action_dim):
+            action = out_val["actions"][:, :, action_dim_ind]
+            logger.info(
+                f"Val - Action dim {action_dim_ind+1} mean {np.mean(action)} std {np.std(action)} min {np.min(action)} max {np.max(action)}"
+            )
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--env_name", type=str, default="hopper-medium-v2")
+    parser.add_argument("--save_dir", type=str, default=".")
+    parser.add_argument("--save_name_prefix", type=str, default="")
+    parser.add_argument("--val_split", type=float, default="0.2")
+    parser.add_argument("--max_episodes", type=int, default="-1")
+    args = parser.parse_args()
+
+    import datetime
+
+    # import logging.config
+    if args.max_episodes > 0:
+        args.save_name_prefix += f"max_episodes_{args.max_episodes}_"
+    os.makedirs(args.save_dir, exist_ok=True)
+    log_path = os.path.join(
+        args.save_dir,
+        args.save_name_prefix
+        + f"_{datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')}.log",
+    )
+    logger = logging.getLogger("get_D4RL_dataset")
+    logger.setLevel(logging.INFO)
+    file_handler = logging.FileHandler(log_path)
+    file_handler.setLevel(logging.INFO)  # Set the minimum level for this handler
+    formatter = logging.Formatter(
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    )
+    file_handler.setFormatter(formatter)
+    logger.addHandler(file_handler)
+
+    make_dataset(
+        args.env_name, args.save_dir, args.save_name_prefix, args.val_split, logger
+    )

script/dataset/process_d3il_dataset.py

@@ -0,0 +1,294 @@
+"""
+Process d3il dataset and save it into our custom format so it can be loaded for diffusion training.
+"""
+
+import os
+import numpy as np
+from tqdm import tqdm
+import pickle
+import random
+import matplotlib.pyplot as plt
+from copy import deepcopy
+
+from agent.dataset.d3il_dataset.aligning_dataset import Aligning_Dataset
+from agent.dataset.d3il_dataset.avoiding_dataset import Avoiding_Dataset
+from agent.dataset.d3il_dataset.pushing_dataset import Pushing_Dataset
+from agent.dataset.d3il_dataset.sorting_dataset import Sorting_Dataset
+from agent.dataset.d3il_dataset.stacking_dataset import Stacking_Dataset
+
+
+def make_dataset(load_path, save_dir, save_name_prefix, env_type, val_split):
+    if env_type == "align":
+        demo_dataset = Aligning_Dataset(
+            load_path,
+            action_dim=3,
+            obs_dim=20,
+            max_len_data=512,
+        )
+    elif env_type == "avoid":
+        demo_dataset = Avoiding_Dataset(
+            load_path,
+            action_dim=2,
+            obs_dim=4,
+            max_len_data=200,
+        )
+    elif env_type == "push":
+        demo_dataset = Pushing_Dataset(
+            load_path,
+            action_dim=2,
+            obs_dim=10,
+            max_len_data=512,
+        )
+    elif env_type == "sort":
+        # Can config number of boxes to be 2, 4, or 6.
+        # TODO: add other numbers of boxes
+        demo_dataset = Sorting_Dataset(
+            load_path,
+            action_dim=2,
+            obs_dim=10,
+            max_len_data=600,
+            num_boxes=2,
+        )
+    elif env_type == "stack":
+        demo_dataset = Stacking_Dataset(
+            load_path,
+            action_dim=8,
+            obs_dim=20,
+            max_len_data=1000,
+        )
+    else:
+        raise ValueError("Invalid dataset type.")
+    # extract length of each trajectory in the file
+    traj_lengths = []
+    actions = demo_dataset.actions
+    obs = demo_dataset.observations
+    masks = demo_dataset.masks
+    action_dim = actions.shape[2]
+    obs_dim = obs.shape[2]
+    for ep in range(masks.shape[0]):
+        traj_lengths.append(int(masks[ep].sum().item()))
+    traj_lengths = np.array(traj_lengths)
+    max_traj_length = np.max(traj_lengths)
+
+    # split indices in train and val
+    num_traj = len(traj_lengths)
+    num_train = int(num_traj * (1 - val_split))
+    train_indices = random.sample(range(num_traj), k=num_train)
+
+    # take the max and min of obs and action
+    obs_min = np.zeros((obs_dim))
+    obs_max = np.zeros((obs_dim))
+    action_min = np.zeros((action_dim))
+    action_max = np.zeros((action_dim))
+    for i in tqdm(range(len(traj_lengths))):
+        T = traj_lengths[i]
+        obs_traj = obs[i, :T].numpy()
+        action_traj = actions[i, :T].numpy()
+        obs_min = np.min(np.vstack((obs_min, np.min(obs_traj, axis=0))), axis=0)
+        obs_max = np.max(np.vstack((obs_max, np.max(obs_traj, axis=0))), axis=0)
+        action_min = np.min(
+            np.vstack((action_min, np.min(action_traj, axis=0))), axis=0
+        )
+        action_max = np.max(
+            np.vstack((action_max, np.max(action_traj, axis=0))), axis=0
+        )
+    logger.info("\n========== Basic Info ===========")
+    logger.info("total transitions: {}".format(np.sum(traj_lengths)))
+    logger.info("total trajectories: {}".format(len(traj_lengths)))
+    logger.info(
+        f"traj length mean/std: {np.mean(traj_lengths)}, {np.std(traj_lengths)}"
+    )
+    logger.info(f"traj length min/max: {np.min(traj_lengths)}, {np.max(traj_lengths)}")
+    logger.info(f"obs min: {obs_min}")
+    logger.info(f"obs max: {obs_max}")
+    logger.info(f"action min: {action_min}")
+    logger.info(f"action max: {action_max}")
+
+    # do over all indices
+    out_train = {}
+    keys = [
+        "observations",
+        "actions",
+        "rewards",
+    ]
+    total_timesteps = actions.shape[1]
+    out_train["observations"] = np.empty((0, total_timesteps, obs_dim))
+    out_train["actions"] = np.empty((0, total_timesteps, action_dim))
+    out_train["rewards"] = np.empty((0, total_timesteps))
+    out_train["traj_length"] = []
+    out_val = deepcopy(out_train)
+    for i in tqdm(range(len(traj_lengths))):
+        if i in train_indices:
+            out = out_train
+        else:
+            out = out_val
+
+        T = traj_lengths[i]
+        obs_traj = obs[i].numpy()
+        action_traj = actions[i].numpy()
+
+        # scale to [-1, 1] for both ob and action
+        obs_traj = 2 * (obs_traj - obs_min) / (obs_max - obs_min + 1e-6) - 1
+        action_traj = (
+            2 * (action_traj - action_min) / (action_max - action_min + 1e-6) - 1
+        )
+
+        # get episode length
+        traj_length = T
+        out["traj_length"].append(traj_length)
+
+        # extract
+        rewards = np.zeros(total_timesteps)  # no reward from d3il dataset
+        data_traj = {
+            "observations": obs_traj,
+            "actions": action_traj,
+            "rewards": rewards,
+        }
+        for key in keys:
+            traj = data_traj[key]
+            out[key] = np.vstack((out[key], traj[None]))
+
+    # plot all trajectories and save in a figure
+    def plot(out, name):
+        def get_obj_xy_list():
+            mid_pos = 0.5
+            offset = 0.075
+            first_level_y = -0.1
+            level_distance = 0.18
+            return [
+                [mid_pos, first_level_y],
+                [mid_pos - offset, first_level_y + level_distance],
+                [mid_pos + offset, first_level_y + level_distance],
+                [mid_pos - 2 * offset, first_level_y + 2 * level_distance],
+                [mid_pos, first_level_y + 2 * level_distance],
+                [mid_pos + 2 * offset, first_level_y + 2 * level_distance],
+            ]
+
+        pillar_xys = get_obj_xy_list()
+        fig = plt.figure()
+        all_trajs = out["observations"]  # num x timestep x obs
+        for traj, traj_length in zip(all_trajs, out["traj_length"]):
+            # unnormalize
+            traj = (traj + 1) / 2  # [-1, 1] -> [0, 1]
+            traj = traj * (obs_max - obs_min) + obs_min
+            plt.plot(
+                traj[:traj_length, 2], traj[:traj_length, 3], color=(0.3, 0.3, 0.3)
+            )
+        plt.axhline(y=0.4, color=np.array([31, 119, 180]) / 255, linestyle="-")
+        for xy in pillar_xys:
+            circle = plt.Circle(xy, 0.01, color=(0.0, 0.0, 0.0), fill=True)
+            plt.gca().add_patch(circle)
+        plt.xlabel("X pos")
+        plt.ylabel("Y pos")
+        plt.xlim([0.2, 0.8])
+        plt.ylim([-0.3, 0.5])
+        ax = plt.gca()
+        ax.set_aspect("equal", adjustable="box")
+        ax.set_facecolor("white")
+        plt.savefig(os.path.join(save_dir, name))
+        plt.close(fig)
+
+    plot(out_train, name="train-trajs.png")
+    plot(out_val, name="val-trajs.png")
+
+    # Save to np file
+    save_train_path = os.path.join(save_dir, save_name_prefix + "train.pkl")
+    save_val_path = os.path.join(save_dir, save_name_prefix + "val.pkl")
+    with open(save_train_path, "wb") as f:
+        pickle.dump(out_train, f)
+    with open(save_val_path, "wb") as f:
+        pickle.dump(out_val, f)
+    normalization_save_path = os.path.join(
+        save_dir, save_name_prefix + "normalization.npz"
+    )
+    np.savez(
+        normalization_save_path,
+        obs_min=obs_min,
+        obs_max=obs_max,
+        action_min=action_min,
+        action_max=action_max,
+    )
+
+    # debug
+    logger.info("\n========== Final ===========")
+    logger.info(
+        f"Train - Number of episodes and transitions: {len(out_train['traj_length'])}, {np.sum(out_train['traj_length'])}"
+    )
+    logger.info(
+        f"Val - Number of episodes and transitions: {len(out_val['traj_length'])}, {np.sum(out_val['traj_length'])}"
+    )
+    logger.info(
+        f"Train - Mean/Std trajectory length: {np.mean(out_train['traj_length'])}, {np.std(out_train['traj_length'])}"
+    )
+    logger.info(
+        f"Train - Max/Min trajectory length: {np.max(out_train['traj_length'])}, {np.min(out_train['traj_length'])}"
+    )
+    if val_split > 0:
+        logger.info(
+            f"Val - Mean/Std trajectory length: {np.mean(out_val['traj_length'])}, {np.std(out_val['traj_length'])}"
+        )
+        logger.info(
+            f"Val - Max/Min trajectory length: {np.max(out_val['traj_length'])}, {np.min(out_val['traj_length'])}"
+        )
+    for obs_dim_ind in range(obs_dim):
+        obs = out_train["observations"][:, :, obs_dim_ind]
+        logger.info(
+            f"Train - Obs dim {obs_dim_ind+1} mean {np.mean(obs)} std {np.std(obs)} min {np.min(obs)} max {np.max(obs)}"
+        )
+    for action_dim_ind in range(action_dim):
+        action = out_train["actions"][:, :, action_dim_ind]
+        logger.info(
+            f"Train - Action dim {action_dim_ind+1} mean {np.mean(action)} std {np.std(action)} min {np.min(action)} max {np.max(action)}"
+        )
+    if val_split > 0:
+        for obs_dim_ind in range(obs_dim):
+            obs = out_val["observations"][:, :, obs_dim_ind]
+            logger.info(
+                f"Val - Obs dim {obs_dim_ind+1} mean {np.mean(obs)} std {np.std(obs)} min {np.min(obs)} max {np.max(obs)}"
+            )
+        for action_dim_ind in range(action_dim):
+            action = out_val["actions"][:, :, action_dim_ind]
+            logger.info(
+                f"Val - Action dim {action_dim_ind+1} mean {np.mean(action)} std {np.std(action)} min {np.min(action)} max {np.max(action)}"
+            )
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--load_path", type=str, default=".")
+    parser.add_argument("--save_dir", type=str, default=".")
+    parser.add_argument("--save_name_prefix", type=str, default="")
+    parser.add_argument("--env_type", type=str, default="align")
+    parser.add_argument("--val_split", type=float, default="0.2")
+    args = parser.parse_args()
+
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    import logging
+    import datetime
+
+    os.makedirs(args.save_dir, exist_ok=True)
+    log_path = os.path.join(
+        args.save_dir,
+        args.save_name_prefix
+        + f"_{datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')}.log",
+    )
+    logger = logging.getLogger("get_D4RL_dataset")
+    logger.setLevel(logging.INFO)
+    file_handler = logging.FileHandler(log_path)
+    file_handler.setLevel(logging.INFO)  # Set the minimum level for this handler
+    formatter = logging.Formatter(
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    )
+    file_handler.setFormatter(formatter)
+    logger.addHandler(file_handler)
+
+    make_dataset(
+        args.load_path,
+        args.save_dir,
+        args.save_name_prefix,
+        args.env_type,
+        args.val_split,
+    )

script/dataset/process_robomimic_dataset.py

@@ -0,0 +1,421 @@
+"""
+Process robomimic dataset and save it into our custom format so it can be loaded for diffusion training.
+
+Using some code from robomimic/robomimic/scripts/get_dataset_info.py
+
+can-mh:
+    total transitions: 62756
+    total trajectories: 300
+    traj length mean: 209.18666666666667
+    traj length std: 114.42181532479817
+    traj length min: 98
+    traj length max: 1050
+    action min: -1.0
+    action max: 1.0
+
+    {
+        "env_name": "PickPlaceCan",
+        "env_version": "1.4.1",
+        "type": 1,
+        "env_kwargs": {
+            "has_renderer": false,
+            "has_offscreen_renderer": false,
+            "ignore_done": true,
+            "use_object_obs": true,
+            "use_camera_obs": false,
+            "control_freq": 20,
+            "controller_configs": {
+                "type": "OSC_POSE",
+                "input_max": 1,
+                "input_min": -1,
+                "output_max": [
+                    0.05,
+                    0.05,
+                    0.05,
+                    0.5,
+                    0.5,
+                    0.5
+                ],
+                "output_min": [
+                    -0.05,
+                    -0.05,
+                    -0.05,
+                    -0.5,
+                    -0.5,
+                    -0.5
+                ],
+                "kp": 150,
+                "damping": 1,
+                "impedance_mode": "fixed",
+                "kp_limits": [
+                    0,
+                    300
+                ],
+                "damping_limits": [
+                    0,
+                    10
+                ],
+                "position_limits": null,
+                "orientation_limits": null,
+                "uncouple_pos_ori": true,
+                "control_delta": true,
+                "interpolation": null,
+                "ramp_ratio": 0.2
+            },
+            "robots": [
+                "Panda"
+            ],
+            "camera_depths": false,
+            "camera_heights": 84,
+            "camera_widths": 84,
+            "reward_shaping": false
+        }
+    }
+
+robomimic dataset normalizes action to [-1, 1], observation roughly? to [-1, 1]. Seems sometimes the upper value is a bit larger than 1 (but within 1.1).
+
+"""
+
+import numpy as np
+from tqdm import tqdm
+import pickle
+
+try:
+    import h5py  # not included in pyproject.toml
+except:
+    print("Installing h5py")
+    os.system("pip install h5py")
+import os
+import random
+from copy import deepcopy
+import logging
+
+
+def make_dataset(
+    load_path,
+    save_dir,
+    save_name_prefix,
+    val_split,
+    normalize,
+):
+    # Load hdf5 file from load_path
+    with h5py.File(load_path, "r") as f:
+        # put demonstration list in increasing episode order
+        demos = sorted(list(f["data"].keys()))
+        inds = np.argsort([int(elem[5:]) for elem in demos])
+        demos = [demos[i] for i in inds]
+
+        if args.max_episodes > 0:
+            demos = demos[: args.max_episodes]
+
+        # From generate_paper_configs.py: default observation is eef pose, gripper finger position, and object information, all of which are low-dim.
+        low_dim_obs_names = [
+            "robot0_eef_pos",
+            "robot0_eef_quat",
+            "robot0_gripper_qpos",
+        ]
+        if "transport" in load_path:
+            low_dim_obs_names += [
+                "robot1_eef_pos",
+                "robot1_eef_quat",
+                "robot1_gripper_qpos",
+            ]
+        if args.cameras is None:  # state-only
+            low_dim_obs_names.append("object")
+        obs_dim = 0
+        for low_dim_obs_name in low_dim_obs_names:
+            dim = f["data/demo_0/obs/{}".format(low_dim_obs_name)].shape[1]
+            obs_dim += dim
+            logging.info(f"Using {low_dim_obs_name} with dim {dim} for observation")
+        action_dim = f["data/demo_0/actions"].shape[1]
+        logging.info(f"Total low-dim observation dim: {obs_dim}")
+        logging.info(f"Action dim: {action_dim}")
+
+        # get basic stats
+        traj_lengths = []
+        obs_min = np.zeros((obs_dim))
+        obs_max = np.zeros((obs_dim))
+        action_min = np.zeros((action_dim))
+        action_max = np.zeros((action_dim))
+        for ep in demos:
+            traj_lengths.append(f[f"data/{ep}/actions"].shape[0])
+            obs = np.hstack(
+                [
+                    f[f"data/{ep}/obs/{low_dim_obs_name}"][()]
+                    for low_dim_obs_name in low_dim_obs_names
+                ]
+            )
+            actions = f[f"data/{ep}/actions"]
+            obs_min = np.minimum(obs_min, np.min(obs, axis=0))
+            obs_max = np.maximum(obs_max, np.max(obs, axis=0))
+            action_min = np.minimum(action_min, np.min(actions, axis=0))
+            action_max = np.maximum(action_max, np.max(actions, axis=0))
+        traj_lengths = np.array(traj_lengths)
+        max_traj_length = np.max(traj_lengths)
+
+        # report statistics on the data
+        logging.info("===== Basic stats =====")
+        logging.info("total transitions: {}".format(np.sum(traj_lengths)))
+        logging.info("total trajectories: {}".format(traj_lengths.shape[0]))
+        logging.info(
+            f"traj length mean/std: {np.mean(traj_lengths)}, {np.std(traj_lengths)}"
+        )
+        logging.info(
+            f"traj length min/max: {np.min(traj_lengths)}, {np.max(traj_lengths)}"
+        )
+        logging.info(f"obs min: {obs_min}")
+        logging.info(f"obs max: {obs_max}")
+        logging.info(f"action min: {action_min}")
+        logging.info(f"action max: {action_max}")
+
+        # deal with images
+        if args.cameras is not None:
+            img_shapes = []
+            img_names = []  # not necessary but keep old implementation
+            for camera in args.cameras:
+                if f"{camera}_image" in f["data/demo_0/obs"]:
+                    img_shape = f["data/demo_0/obs/{}_image".format(camera)].shape[1:]
+                    img_shapes.append(img_shape)
+                    img_names.append(f"{camera}_image")
+            # ensure all images have the same height and width
+            assert all(
+                [
+                    img_shape[0] == img_shapes[0][0]
+                    and img_shape[1] == img_shapes[0][1]
+                    for img_shape in img_shapes
+                ]
+            )
+            combined_img_shape = (
+                img_shapes[0][0],
+                img_shapes[0][1],
+                sum([img_shape[2] for img_shape in img_shapes]),
+            )
+            logging.info(f"Image shapes: {img_shapes}")
+
+        # split indices in train and val
+        num_traj = len(traj_lengths)
+        num_train = int(num_traj * (1 - val_split))
+        train_indices = random.sample(range(num_traj), k=num_train)
+
+        # do over all indices
+        out_train = {}
+        keys = [
+            "observations",
+            "actions",
+            "rewards",
+        ]
+        if args.cameras is not None:
+            keys.append("images")
+        out_train["observations"] = np.empty((0, max_traj_length, obs_dim))
+        out_train["actions"] = np.empty((0, max_traj_length, action_dim))
+        out_train["rewards"] = np.empty((0, max_traj_length))
+        out_train["traj_length"] = []
+        if args.cameras is not None:
+            out_train["images"] = np.empty(
+                (
+                    0,
+                    max_traj_length,
+                    *combined_img_shape,
+                ),
+                dtype=np.uint8,
+            )
+        out_val = deepcopy(out_train)
+        train_episode_reward_all = []
+        val_episode_reward_all = []
+        for i in tqdm(range(len(demos))):
+            ep = demos[i]
+            if i in train_indices:
+                out = out_train
+            else:
+                out = out_val
+
+            # get episode length
+            traj_length = f[f"data/{ep}"].attrs["num_samples"]
+            out["traj_length"].append(traj_length)
+            # print("Episode:", i, "Trajectory length:", traj_length)
+
+            # extract
+            raw_actions = f[f"data/{ep}/actions"][()]
+            rewards = f[f"data/{ep}/rewards"][()]
+            raw_obs = np.hstack(
+                [
+                    f[f"data/{ep}/obs/{low_dim_obs_name}"][()]
+                    for low_dim_obs_name in low_dim_obs_names
+                ]
+            )  # not normalized
+
+            # scale to [-1, 1] for both ob and action
+            if normalize:
+                obs = 2 * (raw_obs - obs_min) / (obs_max - obs_min + 1e-6) - 1
+                actions = (
+                    2 * (raw_actions - action_min) / (action_max - action_min + 1e-6)
+                    - 1
+                )
+            else:
+                obs = raw_obs
+                actions = raw_actions
+
+            data_traj = {
+                "observations": obs,
+                "actions": actions,
+                "rewards": rewards,
+            }
+            if args.cameras is not None:  # no normalization
+                data_traj["images"] = np.concatenate(
+                    (
+                        [
+                            f["data/{}/obs/{}".format(ep, img_name)][()]
+                            for img_name in img_names
+                        ]
+                    ),
+                    axis=-1,
+                )
+
+            # apply padding to make all episodes have the same max steps
+            # later when we load this dataset, we will use the traj_length to slice the data
+            for key in keys:
+                traj = data_traj[key]
+                if traj.ndim == 1:
+                    pad_width = (0, max_traj_length - len(traj))
+                elif traj.ndim == 2:
+                    pad_width = ((0, max_traj_length - traj.shape[0]), (0, 0))
+                elif traj.ndim == 4:
+                    pad_width = (
+                        (0, max_traj_length - traj.shape[0]),
+                        (0, 0),
+                        (0, 0),
+                        (0, 0),
+                    )
+                else:
+                    raise ValueError("Unsupported dimension")
+                traj = np.pad(
+                    traj,
+                    pad_width,
+                    mode="constant",
+                    constant_values=0,
+                )
+                out[key] = np.vstack((out[key], traj[None]))
+
+            # check reward
+            if i in train_indices:
+                train_episode_reward_all.append(np.sum(data_traj["rewards"]))
+            else:
+                val_episode_reward_all.append(np.sum(data_traj["rewards"]))
+
+    # Save to np file
+    save_train_path = os.path.join(save_dir, save_name_prefix + "train.pkl")
+    save_val_path = os.path.join(save_dir, save_name_prefix + "val.pkl")
+    with open(save_train_path, "wb") as f:
+        pickle.dump(out_train, f)
+    with open(save_val_path, "wb") as f:
+        pickle.dump(out_val, f)
+    if normalize:
+        normalization_save_path = os.path.join(
+            save_dir, save_name_prefix + "normalization.npz"
+        )
+        np.savez(
+            normalization_save_path,
+            obs_min=obs_min,
+            obs_max=obs_max,
+            action_min=action_min,
+            action_max=action_max,
+        )
+
+    # debug
+    logging.info("\n========== Final ===========")
+    logging.info(
+        f"Train - Number of episodes and transitions: {len(out_train['traj_length'])}, {np.sum(out_train['traj_length'])}"
+    )
+    logging.info(
+        f"Val - Number of episodes and transitions: {len(out_val['traj_length'])}, {np.sum(out_val['traj_length'])}"
+    )
+    logging.info(
+        f"Train - Mean/Std trajectory length: {np.mean(out_train['traj_length'])}, {np.std(out_train['traj_length'])}"
+    )
+    logging.info(
+        f"Train - Max/Min trajectory length: {np.max(out_train['traj_length'])}, {np.min(out_train['traj_length'])}"
+    )
+    logging.info(
+        f"Train - Mean/Std episode reward: {np.mean(train_episode_reward_all)},  {np.std(train_episode_reward_all)}"
+    )
+    if val_split > 0:
+        logging.info(
+            f"Val - Mean/Std trajectory length: {np.mean(out_val['traj_length'])}, {np.std(out_val['traj_length'])}"
+        )
+        logging.info(
+            f"Val - Max/Min trajectory length: {np.max(out_val['traj_length'])}, {np.min(out_val['traj_length'])}"
+        )
+        logging.info(
+            f"Val - Mean/Std episode reward: {np.mean(val_episode_reward_all)},  {np.std(val_episode_reward_all)}"
+        )
+    for obs_dim_ind in range(obs_dim):
+        obs = out_train["observations"][:, :, obs_dim_ind]
+        logging.info(
+            f"Train - Obs dim {obs_dim_ind+1} mean {np.mean(obs)} std {np.std(obs)} min {np.min(obs)} max {np.max(obs)}"
+        )
+    for action_dim_ind in range(action_dim):
+        action = out_train["actions"][:, :, action_dim_ind]
+        logging.info(
+            f"Train - Action dim {action_dim_ind+1} mean {np.mean(action)} std {np.std(action)} min {np.min(action)} max {np.max(action)}"
+        )
+    if val_split > 0:
+        for obs_dim_ind in range(obs_dim):
+            obs = out_val["observations"][:, :, obs_dim_ind]
+            logging.info(
+                f"Val - Obs dim {obs_dim_ind+1} mean {np.mean(obs)} std {np.std(obs)} min {np.min(obs)} max {np.max(obs)}"
+            )
+        for action_dim_ind in range(action_dim):
+            action = out_val["actions"][:, :, action_dim_ind]
+            logging.info(
+                f"Val - Action dim {action_dim_ind+1} mean {np.mean(action)} std {np.std(action)} min {np.min(action)} max {np.max(action)}"
+            )
+    # logging.info("Train - Observation shape:", out_train["observations"].shape)
+    # logging.info("Train - Action shape:", out_train["actions"].shape)
+    # logging.info("Train - Reward shape:", out_train["rewards"].shape)
+    # logging.info("Val - Observation shape:", out_val["observations"].shape)
+    # logging.info("Val - Action shape:", out_val["actions"].shape)
+    # logging.info("Val - Reward shape:", out_val["rewards"].shape)
+    # if use_img:
+    # logging.info("Image shapes:", img_shapes)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--load_path", type=str, default=".")
+    parser.add_argument("--save_dir", type=str, default=".")
+    parser.add_argument("--save_name_prefix", type=str, default="")
+    parser.add_argument("--val_split", type=float, default="0.2")
+    parser.add_argument("--max_episodes", type=int, default="-1")
+    parser.add_argument("--normalize", action="store_true")
+    parser.add_argument("--cameras", nargs="*", default=None)
+    args = parser.parse_args()
+
+    import datetime
+
+    if args.max_episodes > 0:
+        args.save_name_prefix += f"max_episodes_{args.max_episodes}_"
+
+    os.makedirs(args.save_dir, exist_ok=True)
+    log_path = os.path.join(
+        args.save_dir,
+        args.save_name_prefix
+        + f"_{datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')}.log",
+    )
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(message)s",
+        handlers=[
+            logging.FileHandler(log_path, mode="w"),
+            logging.StreamHandler(),
+        ],
+    )
+
+    make_dataset(
+        args.load_path,
+        args.save_dir,
+        args.save_name_prefix,
+        args.val_split,
+        args.normalize,
+    )