update

equidiff/equi_diffpo/common/robomimic_config_util.py

@@ -0,0 +1,47 @@
+from omegaconf import OmegaConf
+from robomimic.config import config_factory
+import robomimic.scripts.generate_paper_configs as gpc
+from robomimic.scripts.generate_paper_configs import (
+    modify_config_for_default_image_exp,
+    modify_config_for_default_low_dim_exp,
+    modify_config_for_dataset,
+)
+
+def get_robomimic_config(
+        algo_name='bc_rnn', 
+        hdf5_type='low_dim', 
+        task_name='square', 
+        dataset_type='ph'
+    ):
+    base_dataset_dir = '/tmp/null'
+    filter_key = None
+
+    # decide whether to use low-dim or image training defaults
+    modifier_for_obs = modify_config_for_default_image_exp
+    if hdf5_type in ["low_dim", "low_dim_sparse", "low_dim_dense"]:
+        modifier_for_obs = modify_config_for_default_low_dim_exp
+
+    algo_config_name = "bc" if algo_name == "bc_rnn" else algo_name
+    config = config_factory(algo_name=algo_config_name)
+    # turn into default config for observation modalities (e.g.: low-dim or rgb)
+    config = modifier_for_obs(config)
+    # add in config based on the dataset
+    config = modify_config_for_dataset(
+        config=config, 
+        task_name=task_name, 
+        dataset_type=dataset_type, 
+        hdf5_type=hdf5_type, 
+        base_dataset_dir=base_dataset_dir,
+        filter_key=filter_key,
+    )
+    # add in algo hypers based on dataset
+    algo_config_modifier = getattr(gpc, f'modify_{algo_name}_config_for_dataset')
+    config = algo_config_modifier(
+        config=config, 
+        task_name=task_name, 
+        dataset_type=dataset_type, 
+        hdf5_type=hdf5_type,
+    )
+    return config
+    
+

equidiff/equi_diffpo/common/robomimic_util.py

@@ -0,0 +1,240 @@
+import numpy as np
+import copy
+
+import h5py
+import robomimic.utils.obs_utils as ObsUtils
+import robomimic.utils.file_utils as FileUtils
+import robomimic.utils.env_utils as EnvUtils
+import robomimic.utils.tensor_utils as TensorUtils
+from scipy.spatial.transform import Rotation
+
+from robomimic.config import config_factory
+
+
+class RobomimicAbsoluteActionConverter:
+    def __init__(self, dataset_path, algo_name='bc'):
+        # default BC config
+        config = config_factory(algo_name=algo_name)
+
+        # read config to set up metadata for observation modalities (e.g. detecting rgb observations)
+        # must ran before create dataset
+        ObsUtils.initialize_obs_utils_with_config(config)
+
+        env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path)
+        abs_env_meta = copy.deepcopy(env_meta)
+        abs_env_meta['env_kwargs']['controller_configs']['control_delta'] = False
+
+        env = EnvUtils.create_env_from_metadata(
+            env_meta=env_meta,
+            render=False, 
+            render_offscreen=False,
+            use_image_obs=False, 
+        )
+        assert len(env.env.robots) in (1, 2)
+
+        abs_env = EnvUtils.create_env_from_metadata(
+            env_meta=abs_env_meta,
+            render=False, 
+            render_offscreen=False,
+            use_image_obs=False, 
+        )
+        assert not abs_env.env.robots[0].controller.use_delta
+
+        self.env = env
+        self.abs_env = abs_env
+        self.file = h5py.File(dataset_path, 'r')
+    
+    def __len__(self):
+        return len(self.file['data'])
+
+    def convert_actions(self, 
+            states: np.ndarray, 
+            actions: np.ndarray) -> np.ndarray:
+        """
+        Given state and delta action sequence
+        generate equivalent goal position and orientation for each step
+        keep the original gripper action intact.
+        """
+        # in case of multi robot
+        # reshape (N,14) to (N,2,7)
+        # or (N,7) to (N,1,7)
+        stacked_actions = actions.reshape(*actions.shape[:-1],-1,7)
+
+        env = self.env
+        # generate abs actions
+        action_goal_pos = np.zeros(
+            stacked_actions.shape[:-1]+(3,), 
+            dtype=stacked_actions.dtype)
+        action_goal_ori = np.zeros(
+            stacked_actions.shape[:-1]+(3,), 
+            dtype=stacked_actions.dtype)
+        action_gripper = stacked_actions[...,[-1]]
+        for i in range(len(states)):
+            _ = env.reset_to({'states': states[i]})
+
+            # taken from robot_env.py L#454
+            for idx, robot in enumerate(env.env.robots):
+                # run controller goal generator
+                robot.control(stacked_actions[i,idx], policy_step=True)
+            
+                # read pos and ori from robots
+                controller = robot.controller
+                action_goal_pos[i,idx] = controller.goal_pos
+                action_goal_ori[i,idx] = Rotation.from_matrix(
+                    controller.goal_ori).as_rotvec()
+
+        stacked_abs_actions = np.concatenate([
+            action_goal_pos,
+            action_goal_ori,
+            action_gripper
+        ], axis=-1)
+        abs_actions = stacked_abs_actions.reshape(actions.shape)
+        return abs_actions
+
+    def convert_idx(self, idx):
+        file = self.file
+        demo = file[f'data/demo_{idx}']
+        # input
+        states = demo['states'][:]
+        actions = demo['actions'][:]
+
+        # generate abs actions
+        abs_actions = self.convert_actions(states, actions)
+        return abs_actions
+
+    def convert_and_eval_idx(self, idx):
+        env = self.env
+        abs_env = self.abs_env
+        file = self.file
+        # first step have high error for some reason, not representative
+        eval_skip_steps = 1
+
+        demo = file[f'data/demo_{idx}']
+        # input
+        states = demo['states'][:]
+        actions = demo['actions'][:]
+
+        # generate abs actions
+        abs_actions = self.convert_actions(states, actions)
+
+        # verify
+        robot0_eef_pos = demo['obs']['robot0_eef_pos'][:]
+        robot0_eef_quat = demo['obs']['robot0_eef_quat'][:]
+
+        delta_error_info = self.evaluate_rollout_error(
+            env, states, actions, robot0_eef_pos, robot0_eef_quat, 
+            metric_skip_steps=eval_skip_steps)
+        abs_error_info = self.evaluate_rollout_error(
+            abs_env, states, abs_actions, robot0_eef_pos, robot0_eef_quat,
+            metric_skip_steps=eval_skip_steps)
+
+        info = {
+            'delta_max_error': delta_error_info,
+            'abs_max_error': abs_error_info
+        }
+        return abs_actions, info
+
+    @staticmethod
+    def evaluate_rollout_error(env, 
+            states, actions, 
+            robot0_eef_pos, 
+            robot0_eef_quat, 
+            metric_skip_steps=1):
+        # first step have high error for some reason, not representative
+
+        # evaluate abs actions
+        rollout_next_states = list()
+        rollout_next_eef_pos = list()
+        rollout_next_eef_quat = list()
+        obs = env.reset_to({'states': states[0]})
+        for i in range(len(states)):
+            obs = env.reset_to({'states': states[i]})
+            obs, reward, done, info = env.step(actions[i])
+            obs = env.get_observation()
+            rollout_next_states.append(env.get_state()['states'])
+            rollout_next_eef_pos.append(obs['robot0_eef_pos'])
+            rollout_next_eef_quat.append(obs['robot0_eef_quat'])
+        rollout_next_states = np.array(rollout_next_states)
+        rollout_next_eef_pos = np.array(rollout_next_eef_pos)
+        rollout_next_eef_quat = np.array(rollout_next_eef_quat)
+
+        next_state_diff = states[1:] - rollout_next_states[:-1]
+        max_next_state_diff = np.max(np.abs(next_state_diff[metric_skip_steps:]))
+
+        next_eef_pos_diff = robot0_eef_pos[1:] - rollout_next_eef_pos[:-1]
+        next_eef_pos_dist = np.linalg.norm(next_eef_pos_diff, axis=-1)
+        max_next_eef_pos_dist = next_eef_pos_dist[metric_skip_steps:].max()
+
+        next_eef_rot_diff = Rotation.from_quat(robot0_eef_quat[1:]) \
+            * Rotation.from_quat(rollout_next_eef_quat[:-1]).inv()
+        next_eef_rot_dist = next_eef_rot_diff.magnitude()
+        max_next_eef_rot_dist = next_eef_rot_dist[metric_skip_steps:].max()
+
+        info = {
+            'state': max_next_state_diff,
+            'pos': max_next_eef_pos_dist,
+            'rot': max_next_eef_rot_dist
+        }
+        return info
+
+class RobomimicObsConverter:
+    def __init__(self, dataset_path, algo_name='bc'):
+        # default BC config
+        # config = config_factory(algo_name=algo_name)
+
+        # read config to set up metadata for observation modalities (e.g. detecting rgb observations)
+        # must ran before create dataset
+        # ObsUtils.initialize_obs_utils_with_config(config)
+
+        env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path)
+        # env_meta['env_kwargs']['camera_names'] = ['birdview', 'agentview', 'sideview', 'robot0_eye_in_hand']
+
+        env = EnvUtils.create_env_for_data_processing(
+            env_meta=env_meta,
+            # camera_names=['frontview', 'birdview', 'agentview', 'sideview', 'agentview_full', 'robot0_robotview', 'robot0_eye_in_hand'], 
+            camera_names=['birdview', 'agentview', 'sideview', 'robot0_eye_in_hand'], 
+            camera_height=84, 
+            camera_width=84, 
+            reward_shaping=False,
+        )
+        # env = EnvUtils.create_env_from_metadata(
+        #     env_meta=env_meta,
+        #     render=True, 
+        #     render_offscreen=True,
+        #     use_image_obs=True,
+        # )
+
+        self.env = env
+        self.file = h5py.File(dataset_path, 'r')
+    
+    def __len__(self):
+        return len(self.file['data'])
+
+    def convert_obs(self, initial_state, states):
+        obss = []
+        self.env.reset()
+        obs = self.env.reset_to(initial_state)
+        obss.append(obs)
+        for i in range(1, len(states)):
+            obs = self.env.reset_to({'states': states[i]})
+            obss.append(obs)
+        return TensorUtils.list_of_flat_dict_to_dict_of_list(obss)
+
+    def convert_idx(self, idx):
+        file = self.file
+        demo = file[f'data/demo_{idx}']
+        # input
+
+        states = demo['states'][:]
+        initial_state = dict(states=states[0])
+        initial_state["model"] = demo.attrs["model_file"]
+
+        # generate abs actions
+        obss = self.convert_obs(initial_state, states)
+        del obss['birdview_image']
+        del obss['birdview_depth']
+        del obss['agentview_depth']
+        del obss['sideview_image']
+        del obss['sideview_depth']
+        del obss['robot0_eye_in_hand_depth']
+        return obss

equidiff/equi_diffpo/dataset/robomimic_replay_image_dataset.py

@@ -0,0 +1,377 @@
+from typing import Dict, List
+import torch
+import numpy as np
+import h5py
+from tqdm import tqdm
+import zarr
+import os
+import shutil
+import copy
+import json
+import hashlib
+from filelock import FileLock
+from threadpoolctl import threadpool_limits
+import concurrent.futures
+import multiprocessing
+from omegaconf import OmegaConf
+from equi_diffpo.common.pytorch_util import dict_apply
+from equi_diffpo.dataset.base_dataset import BaseImageDataset, LinearNormalizer
+from equi_diffpo.model.common.normalizer import LinearNormalizer, SingleFieldLinearNormalizer
+from equi_diffpo.model.common.rotation_transformer import RotationTransformer
+from equi_diffpo.codecs.imagecodecs_numcodecs import register_codecs, Jpeg2k
+from equi_diffpo.common.replay_buffer import ReplayBuffer
+from equi_diffpo.common.sampler import SequenceSampler, get_val_mask
+from equi_diffpo.common.normalize_util import (
+    robomimic_abs_action_only_normalizer_from_stat,
+    robomimic_abs_action_only_dual_arm_normalizer_from_stat,
+    get_range_normalizer_from_stat,
+    get_image_range_normalizer,
+    get_identity_normalizer_from_stat,
+    array_to_stats
+)
+register_codecs()
+
+class RobomimicReplayImageDataset(BaseImageDataset):
+    def __init__(self,
+            shape_meta: dict,
+            dataset_path: str,
+            horizon=1,
+            pad_before=0,
+            pad_after=0,
+            n_obs_steps=None,
+            abs_action=False,
+            rotation_rep='rotation_6d', # ignored when abs_action=False
+            use_legacy_normalizer=False,
+            use_cache=False,
+            seed=42,
+            val_ratio=0.0,
+            n_demo=100
+        ):
+        self.n_demo = n_demo
+        rotation_transformer = RotationTransformer(
+            from_rep='axis_angle', to_rep=rotation_rep)
+
+        replay_buffer = None
+        if use_cache:
+            cache_zarr_path = dataset_path + f'.{n_demo}.' + '.zarr.zip'
+            cache_lock_path = cache_zarr_path + '.lock'
+            print('Acquiring lock on cache.')
+            with FileLock(cache_lock_path):
+                if not os.path.exists(cache_zarr_path):
+                    # cache does not exists
+                    try:
+                        print('Cache does not exist. Creating!')
+                        # store = zarr.DirectoryStore(cache_zarr_path)
+                        replay_buffer = _convert_robomimic_to_replay(
+                            store=zarr.MemoryStore(), 
+                            shape_meta=shape_meta, 
+                            dataset_path=dataset_path, 
+                            abs_action=abs_action, 
+                            rotation_transformer=rotation_transformer,
+                            n_demo=n_demo)
+                        print('Saving cache to disk.')
+                        with zarr.ZipStore(cache_zarr_path) as zip_store:
+                            replay_buffer.save_to_store(
+                                store=zip_store
+                            )
+                    except Exception as e:
+                        shutil.rmtree(cache_zarr_path)
+                        raise e
+                else:
+                    print('Loading cached ReplayBuffer from Disk.')
+                    with zarr.ZipStore(cache_zarr_path, mode='r') as zip_store:
+                        replay_buffer = ReplayBuffer.copy_from_store(
+                            src_store=zip_store, store=zarr.MemoryStore())
+                    print('Loaded!')
+        else:
+            replay_buffer = _convert_robomimic_to_replay(
+                store=zarr.MemoryStore(), 
+                shape_meta=shape_meta, 
+                dataset_path=dataset_path, 
+                abs_action=abs_action, 
+                rotation_transformer=rotation_transformer,
+                n_demo=n_demo)
+
+        rgb_keys = list()
+        lowdim_keys = list()
+        obs_shape_meta = shape_meta['obs']
+        for key, attr in obs_shape_meta.items():
+            type = attr.get('type', 'low_dim')
+            if type == 'rgb':
+                rgb_keys.append(key)
+            elif type == 'low_dim':
+                lowdim_keys.append(key)
+        
+        # for key in rgb_keys:
+        #     replay_buffer[key].compressor.numthreads=1
+
+        key_first_k = dict()
+        if n_obs_steps is not None:
+            # only take first k obs from images
+            for key in rgb_keys + lowdim_keys:
+                key_first_k[key] = n_obs_steps
+
+        val_mask = get_val_mask(
+            n_episodes=replay_buffer.n_episodes, 
+            val_ratio=val_ratio,
+            seed=seed)
+        train_mask = ~val_mask
+        sampler = SequenceSampler(
+            replay_buffer=replay_buffer, 
+            sequence_length=horizon,
+            pad_before=pad_before, 
+            pad_after=pad_after,
+            episode_mask=train_mask,
+            key_first_k=key_first_k)
+        
+        self.replay_buffer = replay_buffer
+        self.sampler = sampler
+        self.shape_meta = shape_meta
+        self.rgb_keys = rgb_keys
+        self.lowdim_keys = lowdim_keys
+        self.abs_action = abs_action
+        self.n_obs_steps = n_obs_steps
+        self.train_mask = train_mask
+        self.horizon = horizon
+        self.pad_before = pad_before
+        self.pad_after = pad_after
+        self.use_legacy_normalizer = use_legacy_normalizer
+
+    def get_validation_dataset(self):
+        val_set = copy.copy(self)
+        val_set.sampler = SequenceSampler(
+            replay_buffer=self.replay_buffer, 
+            sequence_length=self.horizon,
+            pad_before=self.pad_before, 
+            pad_after=self.pad_after,
+            episode_mask=~self.train_mask
+            )
+        val_set.train_mask = ~self.train_mask
+        return val_set
+
+    def get_normalizer(self, **kwargs) -> LinearNormalizer:
+        normalizer = LinearNormalizer()
+
+        # action
+        stat = array_to_stats(self.replay_buffer['action'])
+        if self.abs_action:
+            if stat['mean'].shape[-1] > 10:
+                # dual arm
+                this_normalizer = robomimic_abs_action_only_dual_arm_normalizer_from_stat(stat)
+            else:
+                this_normalizer = robomimic_abs_action_only_normalizer_from_stat(stat)
+            
+            if self.use_legacy_normalizer:
+                this_normalizer = normalizer_from_stat(stat)
+        else:
+            # already normalized
+            this_normalizer = get_identity_normalizer_from_stat(stat)
+        normalizer['action'] = this_normalizer
+
+        # obs
+        for key in self.lowdim_keys:
+            stat = array_to_stats(self.replay_buffer[key])
+
+            if key.endswith('pos'):
+                this_normalizer = get_range_normalizer_from_stat(stat)
+            elif key.endswith('quat'):
+                # quaternion is in [-1,1] already
+                this_normalizer = get_identity_normalizer_from_stat(stat)
+            elif key.endswith('qpos'):
+                this_normalizer = get_range_normalizer_from_stat(stat)
+            else:
+                raise RuntimeError('unsupported')
+            normalizer[key] = this_normalizer
+
+        # image
+        for key in self.rgb_keys:
+            normalizer[key] = get_image_range_normalizer()
+        return normalizer
+
+    def get_all_actions(self) -> torch.Tensor:
+        return torch.from_numpy(self.replay_buffer['action'])
+
+    def __len__(self):
+        return len(self.sampler)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        threadpool_limits(1)
+        data = self.sampler.sample_sequence(idx)
+
+        # to save RAM, only return first n_obs_steps of OBS
+        # since the rest will be discarded anyway.
+        # when self.n_obs_steps is None
+        # this slice does nothing (takes all)
+        T_slice = slice(self.n_obs_steps)
+
+        obs_dict = dict()
+        for key in self.rgb_keys:
+            # move channel last to channel first
+            # T,H,W,C
+            # convert uint8 image to float32
+            obs_dict[key] = np.moveaxis(data[key][T_slice],-1,1
+                ).astype(np.float32) / 255.
+            # T,C,H,W
+            del data[key]
+        for key in self.lowdim_keys:
+            obs_dict[key] = data[key][T_slice].astype(np.float32)
+            del data[key]
+
+        torch_data = {
+            'obs': dict_apply(obs_dict, torch.from_numpy),
+            'action': torch.from_numpy(data['action'].astype(np.float32))
+        }
+        return torch_data
+
+
+def _convert_actions(raw_actions, abs_action, rotation_transformer):
+    actions = raw_actions
+    if abs_action:
+        is_dual_arm = False
+        if raw_actions.shape[-1] == 14:
+            # dual arm
+            raw_actions = raw_actions.reshape(-1,2,7)
+            is_dual_arm = True
+
+        pos = raw_actions[...,:3]
+        rot = raw_actions[...,3:6]
+        gripper = raw_actions[...,6:]
+        rot = rotation_transformer.forward(rot)
+        raw_actions = np.concatenate([
+            pos, rot, gripper
+        ], axis=-1).astype(np.float32)
+    
+        if is_dual_arm:
+            raw_actions = raw_actions.reshape(-1,20)
+        actions = raw_actions
+    return actions
+
+
+def _convert_robomimic_to_replay(store, shape_meta, dataset_path, abs_action, rotation_transformer, 
+        n_workers=None, max_inflight_tasks=None, n_demo=100):
+    if n_workers is None:
+        n_workers = multiprocessing.cpu_count()
+    if max_inflight_tasks is None:
+        max_inflight_tasks = n_workers * 5
+
+    # parse shape_meta
+    rgb_keys = list()
+    lowdim_keys = list()
+    # construct compressors and chunks
+    obs_shape_meta = shape_meta['obs']
+    for key, attr in obs_shape_meta.items():
+        shape = attr['shape']
+        type = attr.get('type', 'low_dim')
+        if type == 'rgb':
+            rgb_keys.append(key)
+        elif type == 'low_dim':
+            lowdim_keys.append(key)
+    
+    root = zarr.group(store)
+    data_group = root.require_group('data', overwrite=True)
+    meta_group = root.require_group('meta', overwrite=True)
+
+    with h5py.File(dataset_path) as file:
+        # count total steps
+        demos = file['data']
+        episode_ends = list()
+        prev_end = 0
+        for i in range(n_demo):
+            demo = demos[f'demo_{i}']
+            episode_length = demo['actions'].shape[0]
+            episode_end = prev_end + episode_length
+            prev_end = episode_end
+            episode_ends.append(episode_end)
+        n_steps = episode_ends[-1]
+        episode_starts = [0] + episode_ends[:-1]
+        _ = meta_group.array('episode_ends', episode_ends, 
+            dtype=np.int64, compressor=None, overwrite=True)
+
+        # save lowdim data
+        for key in tqdm(lowdim_keys + ['action'], desc="Loading lowdim data"):
+            data_key = 'obs/' + key
+            if key == 'action':
+                data_key = 'actions'
+            this_data = list()
+            for i in range(n_demo):
+                demo = demos[f'demo_{i}']
+                this_data.append(demo[data_key][:].astype(np.float32))
+            this_data = np.concatenate(this_data, axis=0)
+            if key == 'action':
+                this_data = _convert_actions(
+                    raw_actions=this_data,
+                    abs_action=abs_action,
+                    rotation_transformer=rotation_transformer
+                )
+                assert this_data.shape == (n_steps,) + tuple(shape_meta['action']['shape'])
+            else:
+                assert this_data.shape == (n_steps,) + tuple(shape_meta['obs'][key]['shape'])
+            _ = data_group.array(
+                name=key,
+                data=this_data,
+                shape=this_data.shape,
+                chunks=this_data.shape,
+                compressor=None,
+                dtype=this_data.dtype
+            )
+        
+        def img_copy(zarr_arr, zarr_idx, hdf5_arr, hdf5_idx):
+            try:
+                zarr_arr[zarr_idx] = hdf5_arr[hdf5_idx]
+                # make sure we can successfully decode
+                _ = zarr_arr[zarr_idx]
+                return True
+            except Exception as e:
+                return False
+        
+        with tqdm(total=n_steps*len(rgb_keys), desc="Loading image data", mininterval=1.0) as pbar:
+            # one chunk per thread, therefore no synchronization needed
+            with concurrent.futures.ThreadPoolExecutor(max_workers=n_workers) as executor:
+                futures = set()
+                for key in rgb_keys:
+                    data_key = 'obs/' + key
+                    shape = tuple(shape_meta['obs'][key]['shape'])
+                    c,h,w = shape
+                    this_compressor = Jpeg2k(level=50)
+                    img_arr = data_group.require_dataset(
+                        name=key,
+                        shape=(n_steps,h,w,c),
+                        chunks=(1,h,w,c),
+                        compressor=this_compressor,
+                        dtype=np.uint8
+                    )
+                    for episode_idx in range(n_demo):
+                        demo = demos[f'demo_{episode_idx}']
+                        hdf5_arr = demo['obs'][key]
+                        for hdf5_idx in range(hdf5_arr.shape[0]):
+                            if len(futures) >= max_inflight_tasks:
+                                # limit number of inflight tasks
+                                completed, futures = concurrent.futures.wait(futures, 
+                                    return_when=concurrent.futures.FIRST_COMPLETED)
+                                for f in completed:
+                                    if not f.result():
+                                        raise RuntimeError('Failed to encode image!')
+                                pbar.update(len(completed))
+
+                            zarr_idx = episode_starts[episode_idx] + hdf5_idx
+                            futures.add(
+                                executor.submit(img_copy, 
+                                    img_arr, zarr_idx, hdf5_arr, hdf5_idx))
+                completed, futures = concurrent.futures.wait(futures)
+                for f in completed:
+                    if not f.result():
+                        raise RuntimeError('Failed to encode image!')
+                pbar.update(len(completed))
+
+    replay_buffer = ReplayBuffer(root)
+    return replay_buffer
+
+def normalizer_from_stat(stat):
+    max_abs = np.maximum(stat['max'].max(), np.abs(stat['min']).max())
+    scale = np.full_like(stat['max'], fill_value=1/max_abs)
+    offset = np.zeros_like(stat['max'])
+    return SingleFieldLinearNormalizer.create_manual(
+        scale=scale,
+        offset=offset,
+        input_stats_dict=stat
+    )

equidiff/equi_diffpo/dataset/robomimic_replay_image_sym_dataset.py

@@ -0,0 +1,90 @@
+from equi_diffpo.dataset.base_dataset import LinearNormalizer
+from equi_diffpo.model.common.normalizer import LinearNormalizer
+from equi_diffpo.dataset.robomimic_replay_image_dataset import RobomimicReplayImageDataset, normalizer_from_stat
+from equi_diffpo.common.normalize_util import (
+    robomimic_abs_action_only_symmetric_normalizer_from_stat,
+    get_range_normalizer_from_stat,
+    get_range_symmetric_normalizer_from_stat,
+    get_image_range_normalizer,
+    get_identity_normalizer_from_stat,
+    array_to_stats
+)
+import numpy as np
+
+class RobomimicReplayImageSymDataset(RobomimicReplayImageDataset):
+    def __init__(self,
+            shape_meta: dict,
+            dataset_path: str,
+            horizon=1,
+            pad_before=0,
+            pad_after=0,
+            n_obs_steps=None,
+            abs_action=False,
+            rotation_rep='rotation_6d', # ignored when abs_action=False
+            use_legacy_normalizer=False,
+            use_cache=False,
+            seed=42,
+            val_ratio=0.0,
+            n_demo=100
+        ):
+        super().__init__(
+            shape_meta,
+            dataset_path,
+            horizon,
+            pad_before,
+            pad_after,
+            n_obs_steps,
+            abs_action,
+            rotation_rep,
+            use_legacy_normalizer,
+            use_cache,
+            seed,
+            val_ratio,
+            n_demo
+        )
+        
+    def get_normalizer(self, **kwargs) -> LinearNormalizer:
+        normalizer = LinearNormalizer()
+
+        # action
+        stat = array_to_stats(self.replay_buffer['action'])
+        if self.abs_action:
+            if stat['mean'].shape[-1] > 10:
+                # dual arm
+                raise NotImplementedError
+            else:
+                this_normalizer = robomimic_abs_action_only_symmetric_normalizer_from_stat(stat)
+            
+            if self.use_legacy_normalizer:
+                this_normalizer = normalizer_from_stat(stat)
+        else:
+            # already normalized
+            this_normalizer = get_identity_normalizer_from_stat(stat)
+        normalizer['action'] = this_normalizer
+
+        # obs
+        for key in self.lowdim_keys:
+            stat = array_to_stats(self.replay_buffer[key])
+
+            if key.endswith('qpos'):
+                this_normalizer = get_range_normalizer_from_stat(stat)
+            elif key.endswith('pos'):
+                this_normalizer = get_range_symmetric_normalizer_from_stat(stat)
+            elif key.endswith('quat'):
+                # quaternion is in [-1,1] already
+                this_normalizer = get_identity_normalizer_from_stat(stat)
+            elif key.find('bbox') > -1:
+                this_normalizer = get_identity_normalizer_from_stat(stat)
+            else:
+                raise RuntimeError('unsupported')
+            normalizer[key] = this_normalizer
+
+        # image
+        for key in self.rgb_keys:
+            normalizer[key] = get_image_range_normalizer()
+        
+        normalizer['pos_vecs'] = get_identity_normalizer_from_stat({'min': -1 * np.ones([10, 2], np.float32), 'max': np.ones([10, 2], np.float32)})
+        normalizer['crops'] = get_image_range_normalizer()
+
+        return normalizer
+

equidiff/equi_diffpo/dataset/robomimic_replay_lowdim_dataset.py

@@ -0,0 +1,169 @@
+from typing import Dict, List
+import torch
+import numpy as np
+import h5py
+from tqdm import tqdm
+import copy
+from equi_diffpo.common.pytorch_util import dict_apply
+from equi_diffpo.dataset.base_dataset import BaseLowdimDataset, LinearNormalizer
+from equi_diffpo.model.common.normalizer import LinearNormalizer, SingleFieldLinearNormalizer
+from equi_diffpo.model.common.rotation_transformer import RotationTransformer
+from equi_diffpo.common.replay_buffer import ReplayBuffer
+from equi_diffpo.common.sampler import (
+    SequenceSampler, get_val_mask, downsample_mask)
+from equi_diffpo.common.normalize_util import (
+    robomimic_abs_action_only_normalizer_from_stat,
+    robomimic_abs_action_only_dual_arm_normalizer_from_stat,
+    get_identity_normalizer_from_stat,
+    array_to_stats
+)
+
+class RobomimicReplayLowdimDataset(BaseLowdimDataset):
+    def __init__(self,
+            dataset_path: str,
+            horizon=1,
+            pad_before=0,
+            pad_after=0,
+            obs_keys: List[str]=[
+                'object', 
+                'robot0_eef_pos', 
+                'robot0_eef_quat', 
+                'robot0_gripper_qpos'],
+            abs_action=False,
+            rotation_rep='rotation_6d',
+            use_legacy_normalizer=False,
+            seed=42,
+            val_ratio=0.0,
+            max_train_episodes=None,
+            n_demo=100
+        ):
+        obs_keys = list(obs_keys)
+        rotation_transformer = RotationTransformer(
+            from_rep='axis_angle', to_rep=rotation_rep)
+
+        replay_buffer = ReplayBuffer.create_empty_numpy()
+        with h5py.File(dataset_path) as file:
+            demos = file['data']
+            for i in tqdm(range(n_demo), desc="Loading hdf5 to ReplayBuffer"):
+                demo = demos[f'demo_{i}']
+                episode = _data_to_obs(
+                    raw_obs=demo['obs'],
+                    raw_actions=demo['actions'][:].astype(np.float32),
+                    obs_keys=obs_keys,
+                    abs_action=abs_action,
+                    rotation_transformer=rotation_transformer)
+                replay_buffer.add_episode(episode)
+
+        val_mask = get_val_mask(
+            n_episodes=replay_buffer.n_episodes, 
+            val_ratio=val_ratio,
+            seed=seed)
+        train_mask = ~val_mask
+        train_mask = downsample_mask(
+            mask=train_mask, 
+            max_n=max_train_episodes, 
+            seed=seed)
+
+        sampler = SequenceSampler(
+            replay_buffer=replay_buffer, 
+            sequence_length=horizon,
+            pad_before=pad_before, 
+            pad_after=pad_after,
+            episode_mask=train_mask)
+        
+        self.replay_buffer = replay_buffer
+        self.sampler = sampler
+        self.abs_action = abs_action
+        self.train_mask = train_mask
+        self.horizon = horizon
+        self.pad_before = pad_before
+        self.pad_after = pad_after
+        self.use_legacy_normalizer = use_legacy_normalizer
+    
+    def get_validation_dataset(self):
+        val_set = copy.copy(self)
+        val_set.sampler = SequenceSampler(
+            replay_buffer=self.replay_buffer, 
+            sequence_length=self.horizon,
+            pad_before=self.pad_before, 
+            pad_after=self.pad_after,
+            episode_mask=~self.train_mask
+            )
+        val_set.train_mask = ~self.train_mask
+        return val_set
+
+    def get_normalizer(self, **kwargs) -> LinearNormalizer:
+        normalizer = LinearNormalizer()
+
+        # action
+        stat = array_to_stats(self.replay_buffer['action'])
+        if self.abs_action:
+            if stat['mean'].shape[-1] > 10:
+                # dual arm
+                this_normalizer = robomimic_abs_action_only_dual_arm_normalizer_from_stat(stat)
+            else:
+                this_normalizer = robomimic_abs_action_only_normalizer_from_stat(stat)
+            
+            if self.use_legacy_normalizer:
+                this_normalizer = normalizer_from_stat(stat)
+        else:
+            # already normalized
+            this_normalizer = get_identity_normalizer_from_stat(stat)
+        normalizer['action'] = this_normalizer
+        
+        # aggregate obs stats
+        obs_stat = array_to_stats(self.replay_buffer['obs'])
+
+
+        normalizer['obs'] = normalizer_from_stat(obs_stat)
+        return normalizer
+
+    def get_all_actions(self) -> torch.Tensor:
+        return torch.from_numpy(self.replay_buffer['action'])
+    
+    def __len__(self):
+        return len(self.sampler)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        data = self.sampler.sample_sequence(idx)
+        torch_data = dict_apply(data, torch.from_numpy)
+        return torch_data
+
+def normalizer_from_stat(stat):
+    max_abs = np.maximum(stat['max'].max(), np.abs(stat['min']).max())
+    scale = np.full_like(stat['max'], fill_value=1/max_abs)
+    offset = np.zeros_like(stat['max'])
+    return SingleFieldLinearNormalizer.create_manual(
+        scale=scale,
+        offset=offset,
+        input_stats_dict=stat
+    )
+    
+def _data_to_obs(raw_obs, raw_actions, obs_keys, abs_action, rotation_transformer):
+    obs = np.concatenate([
+        raw_obs[key] for key in obs_keys
+    ], axis=-1).astype(np.float32)
+
+    if abs_action:
+        is_dual_arm = False
+        if raw_actions.shape[-1] == 14:
+            # dual arm
+            raw_actions = raw_actions.reshape(-1,2,7)
+            is_dual_arm = True
+
+        pos = raw_actions[...,:3]
+        rot = raw_actions[...,3:6]
+        gripper = raw_actions[...,6:]
+        rot = rotation_transformer.forward(rot)
+        raw_actions = np.concatenate([
+            pos, rot, gripper
+        ], axis=-1).astype(np.float32)
+    
+        if is_dual_arm:
+            raw_actions = raw_actions.reshape(-1,20)
+    
+    data = {
+        'obs': obs,
+        'action': raw_actions
+    }
+    return data

equidiff/equi_diffpo/dataset/robomimic_replay_lowdim_sym_dataset.py

@@ -0,0 +1,73 @@
+from typing import Dict, List
+import torch
+import numpy as np
+from equi_diffpo.common.pytorch_util import dict_apply
+from equi_diffpo.dataset.base_dataset import LinearNormalizer
+from equi_diffpo.dataset.robomimic_replay_lowdim_dataset import RobomimicReplayLowdimDataset, normalizer_from_stat
+from equi_diffpo.common.normalize_util import robomimic_abs_action_only_symmetric_normalizer_from_stat
+from equi_diffpo.common.normalize_util import (
+    robomimic_abs_action_only_symmetric_normalizer_from_stat,
+    get_identity_normalizer_from_stat,
+    array_to_stats
+)
+
+
+class RobomimicReplayLowdimSymDataset(RobomimicReplayLowdimDataset):
+    def __init__(self,
+            dataset_path: str,
+            horizon=1,
+            pad_before=0,
+            pad_after=0,
+            obs_keys: List[str]=[
+                'object', 
+                'robot0_eef_pos', 
+                'robot0_eef_quat', 
+                'robot0_gripper_qpos'],
+            abs_action=False,
+            rotation_rep='rotation_6d',
+            use_legacy_normalizer=False,
+            seed=42,
+            val_ratio=0.0,
+            max_train_episodes=None,
+            n_demo=100
+        ):
+        super().__init__(
+            dataset_path,
+            horizon,
+            pad_before,
+            pad_after,
+            obs_keys,
+            abs_action,
+            rotation_rep,
+            use_legacy_normalizer,
+            seed,
+            val_ratio,
+            max_train_episodes,
+            n_demo,
+        )
+
+    def get_normalizer(self, **kwargs) -> LinearNormalizer:
+        normalizer = LinearNormalizer()
+
+        # action
+        stat = array_to_stats(self.replay_buffer['action'])
+        if self.abs_action:
+            if stat['mean'].shape[-1] > 10:
+                # dual arm
+                raise NotImplementedError
+            else:
+                this_normalizer = robomimic_abs_action_only_symmetric_normalizer_from_stat(stat)
+            
+            if self.use_legacy_normalizer:
+                this_normalizer = normalizer_from_stat(stat)
+        else:
+            # already normalized
+            this_normalizer = get_identity_normalizer_from_stat(stat)
+        normalizer['action'] = this_normalizer
+        
+        # aggregate obs stats
+        obs_stat = array_to_stats(self.replay_buffer['obs'])
+
+
+        normalizer['obs'] = normalizer_from_stat(obs_stat)
+        return normalizer

equidiff/equi_diffpo/dataset/robomimic_replay_point_cloud_dataset.py

@@ -0,0 +1,407 @@
+from typing import Dict, List
+import torch
+import numpy as np
+import h5py
+from tqdm import tqdm
+import zarr
+import os
+import shutil
+import copy
+import json
+import hashlib
+from filelock import FileLock
+from threadpoolctl import threadpool_limits
+import concurrent.futures
+import multiprocessing
+from omegaconf import OmegaConf
+from equi_diffpo.common.pytorch_util import dict_apply
+from equi_diffpo.dataset.base_dataset import BaseImageDataset, LinearNormalizer
+from equi_diffpo.model.common.normalizer import LinearNormalizer, SingleFieldLinearNormalizer
+from equi_diffpo.model.common.rotation_transformer import RotationTransformer
+from equi_diffpo.codecs.imagecodecs_numcodecs import register_codecs, Jpeg2k
+from equi_diffpo.common.replay_buffer import ReplayBuffer
+from equi_diffpo.common.sampler import SequenceSampler, get_val_mask
+from equi_diffpo.common.normalize_util import (
+    robomimic_abs_action_only_normalizer_from_stat,
+    get_range_normalizer_from_stat,
+    get_voxel_identity_normalizer,
+    get_image_range_normalizer,
+    get_identity_normalizer_from_stat,
+    array_to_stats
+)
+register_codecs()
+
+class RobomimicReplayPointCloudDataset(BaseImageDataset):
+    def __init__(self,
+            shape_meta: dict,
+            dataset_path: str,
+            horizon=1,
+            pad_before=0,
+            pad_after=0,
+            n_obs_steps=None,
+            abs_action=False,
+            rotation_rep='rotation_6d', # ignored when abs_action=False
+            use_legacy_normalizer=False,
+            use_cache=False,
+            seed=42,
+            val_ratio=0.0,
+            n_demo=100,
+        ):
+        self.n_demo = n_demo
+        rotation_transformer = RotationTransformer(
+            from_rep='axis_angle', to_rep=rotation_rep)
+
+        replay_buffer = None
+        if use_cache:
+            cache_zarr_path = dataset_path + f'.{n_demo}.' + '.zarr.zip'
+            cache_lock_path = cache_zarr_path + '.lock'
+            print('Acquiring lock on cache.')
+            with FileLock(cache_lock_path):
+                if not os.path.exists(cache_zarr_path):
+                    # cache does not exists
+                    try:
+                        print('Cache does not exist. Creating!')
+                        # store = zarr.DirectoryStore(cache_zarr_path)
+                        replay_buffer = _convert_point_cloud_to_replay(
+                            store=zarr.MemoryStore(), 
+                            shape_meta=shape_meta, 
+                            dataset_path=dataset_path, 
+                            abs_action=abs_action, 
+                            rotation_transformer=rotation_transformer,
+                            n_demo=n_demo)
+                        print('Saving cache to disk.')
+                        with zarr.ZipStore(cache_zarr_path) as zip_store:
+                            replay_buffer.save_to_store(
+                                store=zip_store
+                            )
+                    except Exception as e:
+                        shutil.rmtree(cache_zarr_path)
+                        raise e
+                else:
+                    print('Loading cached ReplayBuffer from Disk.')
+                    with zarr.ZipStore(cache_zarr_path, mode='r') as zip_store:
+                        replay_buffer = ReplayBuffer.copy_from_store(
+                            src_store=zip_store, store=zarr.MemoryStore())
+                    print('Loaded!')
+        else:
+            replay_buffer = _convert_point_cloud_to_replay(
+                store=zarr.MemoryStore(), 
+                shape_meta=shape_meta, 
+                dataset_path=dataset_path, 
+                abs_action=abs_action, 
+                rotation_transformer=rotation_transformer,
+                n_demo=n_demo)
+
+        rgb_keys = list()
+        pc_keys = list()
+        lowdim_keys = list()
+        obs_shape_meta = shape_meta['obs']
+        for key, attr in obs_shape_meta.items():
+            type = attr.get('type', 'low_dim')
+            if type == 'rgb':
+                rgb_keys.append(key)
+            if type == 'point_cloud':
+                pc_keys.append(key)
+            elif type == 'low_dim':
+                lowdim_keys.append(key)
+        
+        # for key in rgb_keys:
+        #     replay_buffer[key].compressor.numthreads=1
+
+        key_first_k = dict()
+        if n_obs_steps is not None:
+            # only take first k obs from images
+            for key in rgb_keys + pc_keys + lowdim_keys:
+                key_first_k[key] = n_obs_steps
+
+        val_mask = get_val_mask(
+            n_episodes=replay_buffer.n_episodes, 
+            val_ratio=val_ratio,
+            seed=seed)
+        train_mask = ~val_mask
+        sampler = SequenceSampler(
+            replay_buffer=replay_buffer, 
+            sequence_length=horizon,
+            pad_before=pad_before, 
+            pad_after=pad_after,
+            episode_mask=train_mask,
+            key_first_k=key_first_k)
+        
+        self.replay_buffer = replay_buffer
+        self.sampler = sampler
+        self.shape_meta = shape_meta
+        self.rgb_keys = rgb_keys
+        self.pc_keys = pc_keys
+        self.lowdim_keys = lowdim_keys
+        self.abs_action = abs_action
+        self.n_obs_steps = n_obs_steps
+        self.train_mask = train_mask
+        self.horizon = horizon
+        self.pad_before = pad_before
+        self.pad_after = pad_after
+        self.use_legacy_normalizer = use_legacy_normalizer
+
+    def get_validation_dataset(self):
+        val_set = copy.copy(self)
+        val_set.sampler = SequenceSampler(
+            replay_buffer=self.replay_buffer, 
+            sequence_length=self.horizon,
+            pad_before=self.pad_before, 
+            pad_after=self.pad_after,
+            episode_mask=~self.train_mask
+            )
+        val_set.train_mask = ~self.train_mask
+        return val_set
+
+    def get_normalizer(self, mode='limits', **kwargs) -> LinearNormalizer:
+        data = {
+            'action': self.replay_buffer['action'],
+            'robot0_eef_pos': self.replay_buffer['robot0_eef_pos'][...,:],
+            'robot0_eef_quat': self.replay_buffer['robot0_eef_quat'][...,:],
+            'robot0_gripper_qpos': self.replay_buffer['robot0_gripper_qpos'][...,:],
+            'point_cloud': self.replay_buffer['point_cloud'],
+        }
+        normalizer = LinearNormalizer()
+        normalizer.fit(data=data, last_n_dims=1, mode=mode, **kwargs)
+        # normalizer['point_cloud'] = SingleFieldLinearNormalizer.create_identity()
+        return normalizer
+
+    def get_all_actions(self) -> torch.Tensor:
+        return torch.from_numpy(self.replay_buffer['action'])
+
+    def __len__(self):
+        return len(self.sampler)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        threadpool_limits(1)
+        data = self.sampler.sample_sequence(idx)
+
+        # to save RAM, only return first n_obs_steps of OBS
+        # since the rest will be discarded anyway.
+        # when self.n_obs_steps is None
+        # this slice does nothing (takes all)
+        T_slice = slice(self.n_obs_steps)
+
+        obs_dict = dict()
+        for key in self.rgb_keys:
+            # move channel last to channel first
+            # T,H,W,C
+            # convert uint8 image to float32
+            obs_dict[key] = np.moveaxis(data[key][T_slice],-1,1
+                ).astype(np.float32) / 255.
+            # T,C,H,W
+            del data[key]
+        for key in self.pc_keys:
+            obs_dict[key] = data[key][T_slice].astype(np.float32)
+            del data[key]
+        for key in self.lowdim_keys:
+            obs_dict[key] = data[key][T_slice].astype(np.float32)
+            del data[key]
+
+        torch_data = {
+            'obs': dict_apply(obs_dict, torch.from_numpy),
+            'action': torch.from_numpy(data['action'].astype(np.float32))
+        }
+        return torch_data
+
+
+def _convert_actions(raw_actions, abs_action, rotation_transformer):
+    actions = raw_actions
+    if abs_action:
+        is_dual_arm = False
+        if raw_actions.shape[-1] == 14:
+            # dual arm
+            raw_actions = raw_actions.reshape(-1,2,7)
+            is_dual_arm = True
+
+        pos = raw_actions[...,:3]
+        rot = raw_actions[...,3:6]
+        gripper = raw_actions[...,6:]
+        rot = rotation_transformer.forward(rot)
+        raw_actions = np.concatenate([
+            pos, rot, gripper
+        ], axis=-1).astype(np.float32)
+    
+        if is_dual_arm:
+            raw_actions = raw_actions.reshape(-1,20)
+        actions = raw_actions
+    return actions
+
+
+def _convert_point_cloud_to_replay(store, shape_meta, dataset_path, abs_action, rotation_transformer, 
+        n_workers=None, max_inflight_tasks=None, n_demo=100):
+    if n_workers is None:
+        n_workers = multiprocessing.cpu_count()
+    if max_inflight_tasks is None:
+        max_inflight_tasks = n_workers * 5
+
+    # parse shape_meta
+    pc_keys = list()
+    rgb_keys = list()
+    lowdim_keys = list()
+    # construct compressors and chunks
+    obs_shape_meta = shape_meta['obs']
+    for key, attr in obs_shape_meta.items():
+        shape = attr['shape']
+        type = attr.get('type', 'low_dim')
+        if type == 'rgb':
+            rgb_keys.append(key)
+        elif type == 'point_cloud':
+            pc_keys.append(key)
+        elif type == 'low_dim':
+            lowdim_keys.append(key)
+    
+    root = zarr.group(store)
+    data_group = root.require_group('data', overwrite=True)
+    meta_group = root.require_group('meta', overwrite=True)
+
+    with h5py.File(dataset_path) as file:
+        # count total steps
+        demos = file['data']
+        episode_ends = list()
+        prev_end = 0
+        n_demo = min(n_demo, len(demos))
+        for i in range(n_demo):
+            demo = demos[f'demo_{i}']
+            episode_length = demo['actions'].shape[0]
+            episode_end = prev_end + episode_length
+            prev_end = episode_end
+            episode_ends.append(episode_end)
+        n_steps = episode_ends[-1]
+        episode_starts = [0] + episode_ends[:-1]
+        _ = meta_group.array('episode_ends', episode_ends, 
+            dtype=np.int64, compressor=None, overwrite=True)
+
+        # save lowdim data
+        for key in tqdm(lowdim_keys + ['action'], desc="Loading lowdim data"):
+            data_key = 'obs/' + key
+            if key == 'action':
+                data_key = 'actions'
+            this_data = list()
+            for i in range(n_demo):
+                demo = demos[f'demo_{i}']
+                this_data.append(demo[data_key][:].astype(np.float32))
+            this_data = np.concatenate(this_data, axis=0)
+            if key == 'action':
+                this_data = _convert_actions(
+                    raw_actions=this_data,
+                    abs_action=abs_action,
+                    rotation_transformer=rotation_transformer
+                )
+                assert this_data.shape == (n_steps,) + tuple(shape_meta['action']['shape'])
+            else:
+                assert this_data.shape == (n_steps,) + tuple(shape_meta['obs'][key]['shape'])
+            _ = data_group.array(
+                name=key,
+                data=this_data,
+                shape=this_data.shape,
+                chunks=this_data.shape,
+                compressor=None,
+                dtype=this_data.dtype
+            )
+        
+        def pc_copy(zarr_arr, zarr_idx, hdf5_arr, hdf5_idx):
+            try:
+                zarr_arr[zarr_idx] = hdf5_arr[hdf5_idx]
+                _ = zarr_arr[zarr_idx]
+                return True
+            except Exception as e:
+                return False
+        
+        def img_copy(zarr_arr, zarr_idx, hdf5_arr, hdf5_idx):
+            try:
+                zarr_arr[zarr_idx] = hdf5_arr[hdf5_idx]
+                # make sure we can successfully decode
+                _ = zarr_arr[zarr_idx]
+                return True
+            except Exception as e:
+                return False
+            
+        with tqdm(total=n_steps*len(rgb_keys), desc="Loading image data", mininterval=1.0) as pbar:
+            # one chunk per thread, therefore no synchronization needed
+            with concurrent.futures.ThreadPoolExecutor(max_workers=n_workers) as executor:
+                futures = set()
+                for key in rgb_keys:
+                    data_key = 'obs/' + key
+                    shape = tuple(shape_meta['obs'][key]['shape'])
+                    c,h,w = shape
+                    this_compressor = Jpeg2k(level=50)
+                    img_arr = data_group.require_dataset(
+                        name=key,
+                        shape=(n_steps,h,w,c),
+                        chunks=(1,h,w,c),
+                        compressor=this_compressor,
+                        dtype=np.uint8
+                    )
+                    for episode_idx in range(n_demo):
+                        demo = demos[f'demo_{episode_idx}']
+                        hdf5_arr = demo['obs'][key]
+                        for hdf5_idx in range(hdf5_arr.shape[0]):
+                            if len(futures) >= max_inflight_tasks:
+                                # limit number of inflight tasks
+                                completed, futures = concurrent.futures.wait(futures, 
+                                    return_when=concurrent.futures.FIRST_COMPLETED)
+                                for f in completed:
+                                    if not f.result():
+                                        raise RuntimeError('Failed to encode image!')
+                                pbar.update(len(completed))
+
+                            zarr_idx = episode_starts[episode_idx] + hdf5_idx
+                            futures.add(
+                                executor.submit(img_copy, 
+                                    img_arr, zarr_idx, hdf5_arr, hdf5_idx))
+                completed, futures = concurrent.futures.wait(futures)
+                for f in completed:
+                    if not f.result():
+                        raise RuntimeError('Failed to encode image!')
+                pbar.update(len(completed))
+        
+        with tqdm(total=n_steps*len(pc_keys), desc="Loading point cloud data", mininterval=1.0) as pbar:
+            # one chunk per thread, therefore no synchronization needed
+            with concurrent.futures.ThreadPoolExecutor(max_workers=n_workers) as executor:
+                futures = set()
+                for key in pc_keys:
+                    data_key = key
+                    shape = tuple(shape_meta['obs'][key]['shape'])
+                    n, c = shape
+                    img_arr = data_group.require_dataset(
+                        name=key,
+                        shape=(n_steps, n, c),
+                        chunks=(1, n, c),
+                        dtype=np.float32
+                    )
+                    for episode_idx in range(n_demo):
+                        demo = demos[f'demo_{episode_idx}']
+                        hdf5_arr = demo['obs'][key]
+                        for hdf5_idx in range(hdf5_arr.shape[0]):
+                            if len(futures) >= max_inflight_tasks:
+                                # limit number of inflight tasks
+                                completed, futures = concurrent.futures.wait(futures, 
+                                    return_when=concurrent.futures.FIRST_COMPLETED)
+                                for f in completed:
+                                    if not f.result():
+                                        raise RuntimeError('Failed to encode image!')
+                                pbar.update(len(completed))
+
+                            zarr_idx = episode_starts[episode_idx] + hdf5_idx
+                            futures.add(
+                                executor.submit(pc_copy, 
+                                    img_arr, zarr_idx, hdf5_arr, hdf5_idx))
+                completed, futures = concurrent.futures.wait(futures)
+                for f in completed:
+                    if not f.result():
+                        raise RuntimeError('Failed to encode image!')
+                pbar.update(len(completed))
+
+    replay_buffer = ReplayBuffer(root)
+    return replay_buffer
+
+def normalizer_from_stat(stat):
+    max_abs = np.maximum(stat['max'].max(), np.abs(stat['min']).max())
+    scale = np.full_like(stat['max'], fill_value=1/max_abs)
+    offset = np.zeros_like(stat['max'])
+    return SingleFieldLinearNormalizer.create_manual(
+        scale=scale,
+        offset=offset,
+        input_stats_dict=stat
+    )

equidiff/equi_diffpo/dataset/robomimic_replay_voxel_sym_dataset.py

@@ -0,0 +1,452 @@
+from typing import Dict, List
+import torch
+import numpy as np
+import h5py
+from tqdm import tqdm
+import zarr
+import os
+import shutil
+import copy
+import json
+import hashlib
+from filelock import FileLock
+from threadpoolctl import threadpool_limits
+import concurrent.futures
+import multiprocessing
+from omegaconf import OmegaConf
+from equi_diffpo.common.pytorch_util import dict_apply
+from equi_diffpo.dataset.base_dataset import BaseImageDataset, LinearNormalizer
+from equi_diffpo.model.common.normalizer import LinearNormalizer, SingleFieldLinearNormalizer
+from equi_diffpo.model.common.rotation_transformer import RotationTransformer
+from equi_diffpo.codecs.imagecodecs_numcodecs import register_codecs, Jpeg2k
+from equi_diffpo.common.replay_buffer import ReplayBuffer
+from equi_diffpo.common.sampler import SequenceSampler, get_val_mask
+from equi_diffpo.common.normalize_util import (
+    robomimic_abs_action_only_normalizer_from_stat,
+    get_range_normalizer_from_stat,
+    get_voxel_identity_normalizer,
+    get_image_range_normalizer,
+    get_identity_normalizer_from_stat,
+    array_to_stats
+)
+register_codecs()
+
+class RobomimicReplayVoxelSymDataset(BaseImageDataset):
+    def __init__(self,
+            shape_meta: dict,
+            dataset_path: str,
+            horizon=1,
+            pad_before=0,
+            pad_after=0,
+            n_obs_steps=None,
+            abs_action=False,
+            rotation_rep='rotation_6d', # ignored when abs_action=False
+            use_legacy_normalizer=False,
+            use_cache=False,
+            seed=42,
+            val_ratio=0.0,
+            n_demo=100,
+            ws_size=0.6,
+            ws_x_center=0,
+            ws_y_center=0,
+        ):
+        self.n_demo = n_demo
+        self.ws_size = ws_size
+        self.ws_center = np.array([ws_x_center, ws_y_center])
+        rotation_transformer = RotationTransformer(
+            from_rep='axis_angle', to_rep=rotation_rep)
+
+        replay_buffer = None
+        if use_cache:
+            cache_zarr_path = dataset_path + f'.{n_demo}.' + '.zarr.zip'
+            cache_lock_path = cache_zarr_path + '.lock'
+            print('Acquiring lock on cache.')
+            with FileLock(cache_lock_path):
+                if not os.path.exists(cache_zarr_path):
+                    # cache does not exists
+                    try:
+                        print('Cache does not exist. Creating!')
+                        # store = zarr.DirectoryStore(cache_zarr_path)
+                        replay_buffer = _convert_voxel_to_replay(
+                            store=zarr.MemoryStore(), 
+                            shape_meta=shape_meta, 
+                            dataset_path=dataset_path, 
+                            abs_action=abs_action, 
+                            rotation_transformer=rotation_transformer,
+                            n_demo=n_demo)
+                        print('Saving cache to disk.')
+                        with zarr.ZipStore(cache_zarr_path) as zip_store:
+                            replay_buffer.save_to_store(
+                                store=zip_store
+                            )
+                    except Exception as e:
+                        shutil.rmtree(cache_zarr_path)
+                        raise e
+                else:
+                    print('Loading cached ReplayBuffer from Disk.')
+                    with zarr.ZipStore(cache_zarr_path, mode='r') as zip_store:
+                        replay_buffer = ReplayBuffer.copy_from_store(
+                            src_store=zip_store, store=zarr.MemoryStore())
+                    print('Loaded!')
+        else:
+            replay_buffer = _convert_voxel_to_replay(
+                store=zarr.MemoryStore(), 
+                shape_meta=shape_meta, 
+                dataset_path=dataset_path, 
+                abs_action=abs_action, 
+                rotation_transformer=rotation_transformer,
+                n_demo=n_demo)
+
+        rgb_keys = list()
+        voxel_keys = list()
+        lowdim_keys = list()
+        obs_shape_meta = shape_meta['obs']
+        for key, attr in obs_shape_meta.items():
+            type = attr.get('type', 'low_dim')
+            if type == 'rgb':
+                rgb_keys.append(key)
+            if type == 'voxel':
+                voxel_keys.append(key)
+            elif type == 'low_dim':
+                lowdim_keys.append(key)
+        
+        # for key in rgb_keys:
+        #     replay_buffer[key].compressor.numthreads=1
+
+        key_first_k = dict()
+        if n_obs_steps is not None:
+            # only take first k obs from images
+            for key in rgb_keys + voxel_keys + lowdim_keys:
+                key_first_k[key] = n_obs_steps
+
+        val_mask = get_val_mask(
+            n_episodes=replay_buffer.n_episodes, 
+            val_ratio=val_ratio,
+            seed=seed)
+        train_mask = ~val_mask
+        sampler = SequenceSampler(
+            replay_buffer=replay_buffer, 
+            sequence_length=horizon,
+            pad_before=pad_before, 
+            pad_after=pad_after,
+            episode_mask=train_mask,
+            key_first_k=key_first_k)
+        
+        self.replay_buffer = replay_buffer
+        self.sampler = sampler
+        self.shape_meta = shape_meta
+        self.rgb_keys = rgb_keys
+        self.voxel_keys = voxel_keys
+        self.lowdim_keys = lowdim_keys
+        self.abs_action = abs_action
+        self.n_obs_steps = n_obs_steps
+        self.train_mask = train_mask
+        self.horizon = horizon
+        self.pad_before = pad_before
+        self.pad_after = pad_after
+        self.use_legacy_normalizer = use_legacy_normalizer
+
+    def get_validation_dataset(self):
+        val_set = copy.copy(self)
+        val_set.sampler = SequenceSampler(
+            replay_buffer=self.replay_buffer, 
+            sequence_length=self.horizon,
+            pad_before=self.pad_before, 
+            pad_after=self.pad_after,
+            episode_mask=~self.train_mask
+            )
+        val_set.train_mask = ~self.train_mask
+        return val_set
+
+    def get_normalizer(self, **kwargs) -> LinearNormalizer:
+        normalizer = LinearNormalizer()
+        # action
+        stat = array_to_stats(self.replay_buffer['action'])
+        if self.abs_action:
+            if stat['mean'].shape[-1] > 10:
+                # dual arm
+                raise NotImplementedError
+            else:
+                magnitute = max(np.max([stat['max'][:2] - self.ws_center, self.ws_center - stat['min'][:2]]), self.ws_size/2)
+                stat['min'][:2] = self.ws_center - magnitute
+                stat['max'][:2] = self.ws_center + magnitute
+                stat['mean'][:2] = self.ws_center
+                this_normalizer = robomimic_abs_action_only_normalizer_from_stat(stat)
+            
+            if self.use_legacy_normalizer:
+                this_normalizer = normalizer_from_stat(stat)
+        else:
+            # already normalized
+            this_normalizer = get_identity_normalizer_from_stat(stat)
+        normalizer['action'] = this_normalizer
+
+        # obs
+        for key in self.lowdim_keys:
+            stat = array_to_stats(self.replay_buffer[key])
+
+            if key.endswith('qpos'):
+                this_normalizer = get_range_normalizer_from_stat(stat)
+            elif key.endswith('pos'):
+                magnitute = max(np.max([stat['max'][:2] - self.ws_center, self.ws_center - stat['min'][:2]]), self.ws_size/2)
+                stat['min'][:2] = self.ws_center - magnitute
+                stat['max'][:2] = self.ws_center + magnitute
+                stat['mean'][:2] = self.ws_center
+                this_normalizer = get_range_normalizer_from_stat(stat)
+            elif key.endswith('quat'):
+                # quaternion is in [-1,1] already
+                this_normalizer = get_identity_normalizer_from_stat(stat)
+            else:
+                raise RuntimeError('unsupported')
+            normalizer[key] = this_normalizer
+
+        # image
+        for key in self.rgb_keys:
+            normalizer[key] = get_image_range_normalizer()
+        for key in self.voxel_keys:
+            normalizer[key] = get_voxel_identity_normalizer()
+
+        return normalizer
+
+    def get_all_actions(self) -> torch.Tensor:
+        return torch.from_numpy(self.replay_buffer['action'])
+
+    def __len__(self):
+        return len(self.sampler)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        threadpool_limits(1)
+        data = self.sampler.sample_sequence(idx)
+
+        # to save RAM, only return first n_obs_steps of OBS
+        # since the rest will be discarded anyway.
+        # when self.n_obs_steps is None
+        # this slice does nothing (takes all)
+        T_slice = slice(self.n_obs_steps)
+
+        obs_dict = dict()
+        for key in self.rgb_keys:
+            # move channel last to channel first
+            # T,H,W,C
+            # convert uint8 image to float32
+            obs_dict[key] = np.moveaxis(data[key][T_slice],-1,1
+                ).astype(np.float32) / 255.
+            # T,C,H,W
+            del data[key]
+        for key in self.voxel_keys:
+            obs_dict[key] = data[key][T_slice].astype(np.float32)
+            obs_dict[key][:, 1:] /= 255.
+            # # convert uint8 image to float32
+            # voxels = np.moveaxis(data[key][T_slice].astype(np.float32), [0, 1, 2, 3, 4], [0, 1, 4, 3, 2])
+            # voxels = np.flip(voxels, (2, 3))
+            # voxels[:, 1:] /= 255.
+            # obs_dict[key] = voxels.copy()
+            del data[key]
+        for key in self.lowdim_keys:
+            obs_dict[key] = data[key][T_slice].astype(np.float32)
+            del data[key]
+
+        torch_data = {
+            'obs': dict_apply(obs_dict, torch.from_numpy),
+            'action': torch.from_numpy(data['action'].astype(np.float32))
+        }
+        return torch_data
+
+
+def _convert_actions(raw_actions, abs_action, rotation_transformer):
+    actions = raw_actions
+    if abs_action:
+        is_dual_arm = False
+        if raw_actions.shape[-1] == 14:
+            # dual arm
+            raw_actions = raw_actions.reshape(-1,2,7)
+            is_dual_arm = True
+
+        pos = raw_actions[...,:3]
+        rot = raw_actions[...,3:6]
+        gripper = raw_actions[...,6:]
+        rot = rotation_transformer.forward(rot)
+        raw_actions = np.concatenate([
+            pos, rot, gripper
+        ], axis=-1).astype(np.float32)
+    
+        if is_dual_arm:
+            raw_actions = raw_actions.reshape(-1,20)
+        actions = raw_actions
+    return actions
+
+
+def _convert_voxel_to_replay(store, shape_meta, dataset_path, abs_action, rotation_transformer, 
+        n_workers=None, max_inflight_tasks=None, n_demo=100):
+    if n_workers is None:
+        n_workers = 24
+    if max_inflight_tasks is None:
+        max_inflight_tasks = n_workers * 5
+
+    # parse shape_meta
+    voxel_keys = list()
+    rgb_keys = list()
+    lowdim_keys = list()
+    # construct compressors and chunks
+    obs_shape_meta = shape_meta['obs']
+    for key, attr in obs_shape_meta.items():
+        shape = attr['shape']
+        type = attr.get('type', 'low_dim')
+        if type == 'rgb':
+            rgb_keys.append(key)
+        elif type == 'voxel':
+            voxel_keys.append(key)
+        elif type == 'low_dim':
+            lowdim_keys.append(key)
+    
+    root = zarr.group(store)
+    data_group = root.require_group('data', overwrite=True)
+    meta_group = root.require_group('meta', overwrite=True)
+
+    with h5py.File(dataset_path) as file:
+        # count total steps
+        demos = file['data']
+        episode_ends = list()
+        prev_end = 0
+        n_demo = min(n_demo, len(demos))
+        for i in range(n_demo):
+            demo = demos[f'demo_{i}']
+            episode_length = demo['actions'].shape[0]
+            episode_end = prev_end + episode_length
+            prev_end = episode_end
+            episode_ends.append(episode_end)
+        n_steps = episode_ends[-1]
+        episode_starts = [0] + episode_ends[:-1]
+        _ = meta_group.array('episode_ends', episode_ends, 
+            dtype=np.int64, compressor=None, overwrite=True)
+
+        # save lowdim data
+        for key in tqdm(lowdim_keys + ['action'], desc="Loading lowdim data"):
+            data_key = 'obs/' + key
+            if key == 'action':
+                data_key = 'actions'
+            this_data = list()
+            for i in range(n_demo):
+                demo = demos[f'demo_{i}']
+                this_data.append(demo[data_key][:].astype(np.float32))
+            this_data = np.concatenate(this_data, axis=0)
+            if key == 'action':
+                this_data = _convert_actions(
+                    raw_actions=this_data,
+                    abs_action=abs_action,
+                    rotation_transformer=rotation_transformer
+                )
+                assert this_data.shape == (n_steps,) + tuple(shape_meta['action']['shape'])
+            else:
+                assert this_data.shape == (n_steps,) + tuple(shape_meta['obs'][key]['shape'])
+            _ = data_group.array(
+                name=key,
+                data=this_data,
+                shape=this_data.shape,
+                chunks=this_data.shape,
+                compressor=None,
+                dtype=this_data.dtype
+            )
+
+        def copy_to_zarr(zarr_arr, hdf5_arr, start_idx, end_idx):
+            try:
+                zarr_arr[start_idx:end_idx] = hdf5_arr
+                # make sure we can successfully decode
+                _ = zarr_arr[start_idx:end_idx]
+                return True
+            except Exception as e:
+                return False
+            
+        with tqdm(total=n_demo*len(rgb_keys), desc="Loading image data", mininterval=1.0) as pbar:
+            # one chunk per thread, therefore no synchronization needed
+            with concurrent.futures.ThreadPoolExecutor(max_workers=n_workers) as executor:
+                futures = set()
+                for key in rgb_keys:
+                    data_key = 'obs/' + key
+                    shape = tuple(shape_meta['obs'][key]['shape'])
+                    c,h,w = shape
+                    this_compressor = Jpeg2k(level=50)
+                    img_arr = data_group.require_dataset(
+                        name=key,
+                        shape=(n_steps,h,w,c),
+                        chunks=(1,h,w,c),
+                        compressor=this_compressor,
+                        dtype=np.uint8
+                    )
+                    for episode_idx in range(n_demo):
+                        demo = demos[f'demo_{episode_idx}']
+                        hdf5_arr = demo['obs'][key][:]
+                        start_idx = episode_starts[episode_idx]
+                        if episode_idx < n_demo - 1:
+                            end_idx = episode_starts[episode_idx+1]
+                        else:
+                            end_idx = n_steps
+                        if len(futures) >= max_inflight_tasks:
+                            # limit number of inflight tasks
+                            completed, futures = concurrent.futures.wait(futures, 
+                                return_when=concurrent.futures.FIRST_COMPLETED)
+                            for f in completed:
+                                if not f.result():
+                                    raise RuntimeError('Failed to encode image!')
+                            pbar.update(len(completed))
+
+                        futures.add(
+                            executor.submit(copy_to_zarr, 
+                                img_arr, hdf5_arr, start_idx, end_idx))
+                completed, futures = concurrent.futures.wait(futures)
+                for f in completed:
+                    if not f.result():
+                        raise RuntimeError('Failed to encode image!')
+                pbar.update(len(completed))
+        
+        with tqdm(total=n_demo*len(voxel_keys), desc="Loading voxel data", mininterval=1.0) as pbar:
+            # one chunk per thread, therefore no synchronization needed
+            with concurrent.futures.ThreadPoolExecutor(max_workers=n_workers) as executor:
+                futures = set()
+                for key in voxel_keys:
+                    data_key = key
+                    shape = tuple(shape_meta['obs'][key]['shape'])
+                    c,h,w,l = shape
+                    img_arr = data_group.require_dataset(
+                        name=key,
+                        shape=(n_steps,c,h,w,l),
+                        chunks=(1,c,h,w,l),
+                        dtype=np.uint8
+                    )
+                    for episode_idx in range(n_demo):
+                        demo = demos[f'demo_{episode_idx}']
+                        hdf5_arr = demo['obs'][key][:]
+                        start_idx = episode_starts[episode_idx]
+                        if episode_idx < n_demo - 1:
+                            end_idx = episode_starts[episode_idx+1]
+                        else:
+                            end_idx = n_steps
+                        if len(futures) >= max_inflight_tasks:
+                            # limit number of inflight tasks
+                            completed, futures = concurrent.futures.wait(futures, 
+                                return_when=concurrent.futures.FIRST_COMPLETED)
+                            for f in completed:
+                                if not f.result():
+                                    raise RuntimeError('Failed to encode image!')
+                            pbar.update(len(completed))
+
+                        futures.add(
+                            executor.submit(copy_to_zarr, 
+                                img_arr, hdf5_arr, start_idx, end_idx))
+                completed, futures = concurrent.futures.wait(futures)
+                for f in completed:
+                    if not f.result():
+                        raise RuntimeError('Failed to encode image!')
+                pbar.update(len(completed))
+
+    replay_buffer = ReplayBuffer(root)
+    return replay_buffer
+
+def normalizer_from_stat(stat):
+    max_abs = np.maximum(stat['max'].max(), np.abs(stat['min']).max())
+    scale = np.full_like(stat['max'], fill_value=1/max_abs)
+    offset = np.zeros_like(stat['max'])
+    return SingleFieldLinearNormalizer.create_manual(
+        scale=scale,
+        offset=offset,
+        input_stats_dict=stat
+    )

equidiff/equi_diffpo/env/robomimic/robomimic_image_wrapper.py

@@ -0,0 +1,170 @@
+from typing import List, Optional
+from matplotlib.pyplot import fill
+import numpy as np
+import gym
+from gym import spaces
+from omegaconf import OmegaConf
+from robomimic.envs.env_robosuite import EnvRobosuite
+
+class RobomimicImageWrapper(gym.Env):
+    def __init__(self, 
+        env: EnvRobosuite,
+        shape_meta: dict,
+        init_state: Optional[np.ndarray]=None,
+        render_obs_key='agentview_image',
+        ):
+
+        self.env = env
+        self.render_obs_key = render_obs_key
+        self.init_state = init_state
+        self.seed_state_map = dict()
+        self._seed = None
+        self.shape_meta = shape_meta
+        self.render_cache = None
+        self.has_reset_before = False
+        
+        # setup spaces
+        action_shape = shape_meta['action']['shape']
+        action_space = spaces.Box(
+            low=-1,
+            high=1,
+            shape=action_shape,
+            dtype=np.float32
+        )
+        self.action_space = action_space
+
+        observation_space = spaces.Dict()
+        for key, value in shape_meta['obs'].items():
+            shape = value['shape']
+            min_value, max_value = -1, 1
+            if key.endswith('image'):
+                min_value, max_value = 0, 1
+            elif key.endswith('depth'):
+                min_value, max_value = 0, 1
+            elif key.endswith('voxels'):
+                min_value, max_value = 0, 1
+            elif key.endswith('point_cloud'):
+                min_value, max_value = -10, 10
+            elif key.endswith('quat'):
+                min_value, max_value = -1, 1
+            elif key.endswith('qpos'):
+                min_value, max_value = -1, 1
+            elif key.endswith('pos'):
+                # better range?
+                min_value, max_value = -1, 1
+            else:
+                raise RuntimeError(f"Unsupported type {key}")
+            
+            this_space = spaces.Box(
+                low=min_value,
+                high=max_value,
+                shape=shape,
+                dtype=np.float32
+            )
+            observation_space[key] = this_space
+        self.observation_space = observation_space
+
+
+    def get_observation(self, raw_obs=None):
+        if raw_obs is None:
+            raw_obs = self.env.get_observation()
+        
+        self.render_cache = raw_obs[self.render_obs_key]
+
+        obs = dict()
+        for key in self.observation_space.keys():
+            obs[key] = raw_obs[key]
+        return obs
+
+    def seed(self, seed=None):
+        np.random.seed(seed=seed)
+        self._seed = seed
+    
+    def reset(self):
+        if self.init_state is not None:
+            if not self.has_reset_before:
+                # the env must be fully reset at least once to ensure correct rendering
+                self.env.reset()
+                self.has_reset_before = True
+
+            # always reset to the same state
+            # to be compatible with gym
+            raw_obs = self.env.reset_to({'states': self.init_state})
+        elif self._seed is not None:
+            # reset to a specific seed
+            seed = self._seed
+            if seed in self.seed_state_map:
+                # env.reset is expensive, use cache
+                raw_obs = self.env.reset_to({'states': self.seed_state_map[seed]})
+            else:
+                # robosuite's initializes all use numpy global random state
+                np.random.seed(seed=seed)
+                raw_obs = self.env.reset()
+                state = self.env.get_state()['states']
+                self.seed_state_map[seed] = state
+            self._seed = None
+        else:
+            # random reset
+            raw_obs = self.env.reset()
+
+        # return obs
+        obs = self.get_observation(raw_obs)
+        return obs
+    
+    def step(self, action):
+        raw_obs, reward, done, info = self.env.step(action)
+        obs = self.get_observation(raw_obs)
+        return obs, reward, done, info
+    
+    def render(self, mode='rgb_array'):
+        if self.render_cache is None:
+            raise RuntimeError('Must run reset or step before render.')
+        img = np.moveaxis(self.render_cache, 0, -1)
+        img = (img * 255).astype(np.uint8)
+        return img
+
+
+def test():
+    import os
+    from omegaconf import OmegaConf
+    cfg_path = os.path.expanduser('~/dev/diffusion_policy/diffusion_policy/config/task/lift_image.yaml')
+    cfg = OmegaConf.load(cfg_path)
+    shape_meta = cfg['shape_meta']
+
+
+    import robomimic.utils.file_utils as FileUtils
+    import robomimic.utils.env_utils as EnvUtils
+    from matplotlib import pyplot as plt
+
+    dataset_path = os.path.expanduser('~/dev/diffusion_policy/data/robomimic/datasets/square/ph/image.hdf5')
+    env_meta = FileUtils.get_env_metadata_from_dataset(
+        dataset_path)
+
+    env = EnvUtils.create_env_from_metadata(
+        env_meta=env_meta,
+        render=False, 
+        render_offscreen=False,
+        use_image_obs=True, 
+    )
+
+    wrapper = RobomimicImageWrapper(
+        env=env,
+        shape_meta=shape_meta
+    )
+    wrapper.seed(0)
+    obs = wrapper.reset()
+    img = wrapper.render()
+    plt.imshow(img)
+
+
+    # states = list()
+    # for _ in range(2):
+    #     wrapper.seed(0)
+    #     wrapper.reset()
+    #     states.append(wrapper.env.get_state()['states'])
+    # assert np.allclose(states[0], states[1])
+
+    # img = wrapper.render()
+    # plt.imshow(img)
+    # wrapper.seed()
+    # states.append(wrapper.env.get_state()['states'])

equidiff/equi_diffpo/env/robomimic/robomimic_lowdim_wrapper.py

@@ -0,0 +1,133 @@
+from typing import List, Dict, Optional
+import numpy as np
+import gym
+from gym.spaces import Box
+from robomimic.envs.env_robosuite import EnvRobosuite
+
+class RobomimicLowdimWrapper(gym.Env):
+    def __init__(self, 
+        env: EnvRobosuite,
+        obs_keys: List[str]=[
+            'object', 
+            'robot0_eef_pos', 
+            'robot0_eef_quat', 
+            'robot0_gripper_qpos'],
+        init_state: Optional[np.ndarray]=None,
+        render_hw=(256,256),
+        render_camera_name='agentview'
+        ):
+
+        self.env = env
+        self.obs_keys = obs_keys
+        self.init_state = init_state
+        self.render_hw = render_hw
+        self.render_camera_name = render_camera_name
+        self.seed_state_map = dict()
+        self._seed = None
+        
+        # setup spaces
+        low = np.full(env.action_dimension, fill_value=-1)
+        high = np.full(env.action_dimension, fill_value=1)
+        self.action_space = Box(
+            low=low,
+            high=high,
+            shape=low.shape,
+            dtype=low.dtype
+        )
+        obs_example = self.get_observation()
+        low = np.full_like(obs_example, fill_value=-1)
+        high = np.full_like(obs_example, fill_value=1)
+        self.observation_space = Box(
+            low=low,
+            high=high,
+            shape=low.shape,
+            dtype=low.dtype
+        )
+
+    def get_observation(self):
+        raw_obs = self.env.get_observation()
+        obs = np.concatenate([
+            raw_obs[key] for key in self.obs_keys
+        ], axis=0)
+        return obs
+
+    def seed(self, seed=None):
+        np.random.seed(seed=seed)
+        self._seed = seed
+    
+    def reset(self):
+        if self.init_state is not None:
+            # always reset to the same state
+            # to be compatible with gym
+            self.env.reset_to({'states': self.init_state})
+        elif self._seed is not None:
+            # reset to a specific seed
+            seed = self._seed
+            if seed in self.seed_state_map:
+                # env.reset is expensive, use cache
+                self.env.reset_to({'states': self.seed_state_map[seed]})
+            else:
+                # robosuite's initializes all use numpy global random state
+                np.random.seed(seed=seed)
+                self.env.reset()
+                state = self.env.get_state()['states']
+                self.seed_state_map[seed] = state
+            self._seed = None
+        else:
+            # random reset
+            self.env.reset()
+
+        # return obs
+        obs = self.get_observation()
+        return obs
+    
+    def step(self, action):
+        raw_obs, reward, done, info = self.env.step(action)
+        obs = np.concatenate([
+            raw_obs[key] for key in self.obs_keys
+        ], axis=0)
+        return obs, reward, done, info
+    
+    def render(self, mode='rgb_array'):
+        h, w = self.render_hw
+        return self.env.render(mode=mode, 
+            height=h, width=w, 
+            camera_name=self.render_camera_name)
+
+
+def test():
+    import robomimic.utils.file_utils as FileUtils
+    import robomimic.utils.env_utils as EnvUtils
+    from matplotlib import pyplot as plt
+
+    dataset_path = '/home/cchi/dev/diffusion_policy/data/robomimic/datasets/square/ph/low_dim.hdf5'
+    env_meta = FileUtils.get_env_metadata_from_dataset(
+        dataset_path)
+
+    env = EnvUtils.create_env_from_metadata(
+        env_meta=env_meta,
+        render=False, 
+        render_offscreen=False,
+        use_image_obs=False, 
+    )
+    wrapper = RobomimicLowdimWrapper(
+        env=env,
+        obs_keys=[
+            'object', 
+            'robot0_eef_pos', 
+            'robot0_eef_quat', 
+            'robot0_gripper_qpos'
+        ]
+    )
+
+    states = list()
+    for _ in range(2):
+        wrapper.seed(0)
+        wrapper.reset()
+        states.append(wrapper.env.get_state()['states'])
+    assert np.allclose(states[0], states[1])
+
+    img = wrapper.render()
+    plt.imshow(img)
+    # wrapper.seed()
+    # states.append(wrapper.env.get_state()['states'])

equidiff/equi_diffpo/env_runner/robomimic_image_runner.py

@@ -0,0 +1,378 @@
+import os
+import wandb
+import numpy as np
+import torch
+import collections
+import pathlib
+import tqdm
+import h5py
+import math
+import dill
+import wandb.sdk.data_types.video as wv
+from equi_diffpo.gym_util.async_vector_env import AsyncVectorEnv
+from equi_diffpo.gym_util.sync_vector_env import SyncVectorEnv
+from equi_diffpo.gym_util.multistep_wrapper import MultiStepWrapper
+from equi_diffpo.gym_util.video_recording_wrapper import VideoRecordingWrapper, VideoRecorder
+from equi_diffpo.model.common.rotation_transformer import RotationTransformer
+
+from equi_diffpo.policy.base_image_policy import BaseImagePolicy
+from equi_diffpo.common.pytorch_util import dict_apply
+from equi_diffpo.env_runner.base_image_runner import BaseImageRunner
+from equi_diffpo.env.robomimic.robomimic_image_wrapper import RobomimicImageWrapper
+import robomimic.utils.file_utils as FileUtils
+import robomimic.utils.env_utils as EnvUtils
+import robomimic.utils.obs_utils as ObsUtils
+
+
+def create_env(env_meta, shape_meta, enable_render=True):
+    modality_mapping = collections.defaultdict(list)
+    for key, attr in shape_meta['obs'].items():
+        modality_mapping[attr.get('type', 'low_dim')].append(key)
+    ObsUtils.initialize_obs_modality_mapping_from_dict(modality_mapping)
+
+    env = EnvUtils.create_env_from_metadata(
+        env_meta=env_meta,
+        render=False, 
+        render_offscreen=enable_render,
+        use_image_obs=enable_render, 
+    )
+    return env
+
+
+class RobomimicImageRunner(BaseImageRunner):
+    """
+    Robomimic envs already enforces number of steps.
+    """
+
+    def __init__(self, 
+            output_dir,
+            dataset_path,
+            shape_meta:dict,
+            n_train=10,
+            n_train_vis=3,
+            train_start_idx=0,
+            n_test=22,
+            n_test_vis=6,
+            test_start_seed=10000,
+            max_steps=400,
+            n_obs_steps=2,
+            n_action_steps=8,
+            render_obs_key='agentview_image',
+            fps=10,
+            crf=22,
+            past_action=False,
+            abs_action=False,
+            tqdm_interval_sec=5.0,
+            n_envs=None
+        ):
+        super().__init__(output_dir)
+
+        if n_envs is None:
+            n_envs = n_train + n_test
+
+        # assert n_obs_steps <= n_action_steps
+        dataset_path = os.path.expanduser(dataset_path)
+        robosuite_fps = 20
+        steps_per_render = max(robosuite_fps // fps, 1)
+
+        # read from dataset
+        env_meta = FileUtils.get_env_metadata_from_dataset(
+            dataset_path)
+        # disable object state observation
+        env_meta['env_kwargs']['use_object_obs'] = False
+
+        rotation_transformer = None
+        if abs_action:
+            env_meta['env_kwargs']['controller_configs']['control_delta'] = False
+            rotation_transformer = RotationTransformer('axis_angle', 'rotation_6d')
+
+        def env_fn():
+            robomimic_env = create_env(
+                env_meta=env_meta, 
+                shape_meta=shape_meta
+            )
+            # Robosuite's hard reset causes excessive memory consumption.
+            # Disabled to run more envs.
+            # https://github.com/ARISE-Initiative/robosuite/blob/92abf5595eddb3a845cd1093703e5a3ccd01e77e/robosuite/environments/base.py#L247-L248
+            robomimic_env.env.hard_reset = False
+            return MultiStepWrapper(
+                VideoRecordingWrapper(
+                    RobomimicImageWrapper(
+                        env=robomimic_env,
+                        shape_meta=shape_meta,
+                        init_state=None,
+                        render_obs_key=render_obs_key
+                    ),
+                    video_recoder=VideoRecorder.create_h264(
+                        fps=fps,
+                        codec='h264',
+                        input_pix_fmt='rgb24',
+                        crf=crf,
+                        thread_type='FRAME',
+                        thread_count=1
+                    ),
+                    file_path=None,
+                    steps_per_render=steps_per_render
+                ),
+                n_obs_steps=n_obs_steps,
+                n_action_steps=n_action_steps,
+                max_episode_steps=max_steps
+            )
+        
+        # For each process the OpenGL context can only be initialized once
+        # Since AsyncVectorEnv uses fork to create worker process,
+        # a separate env_fn that does not create OpenGL context (enable_render=False)
+        # is needed to initialize spaces.
+        def dummy_env_fn():
+            robomimic_env = create_env(
+                    env_meta=env_meta, 
+                    shape_meta=shape_meta,
+                    enable_render=False
+                )
+            return MultiStepWrapper(
+                VideoRecordingWrapper(
+                    RobomimicImageWrapper(
+                        env=robomimic_env,
+                        shape_meta=shape_meta,
+                        init_state=None,
+                        render_obs_key=render_obs_key
+                    ),
+                    video_recoder=VideoRecorder.create_h264(
+                        fps=fps,
+                        codec='h264',
+                        input_pix_fmt='rgb24',
+                        crf=crf,
+                        thread_type='FRAME',
+                        thread_count=1
+                    ),
+                    file_path=None,
+                    steps_per_render=steps_per_render
+                ),
+                n_obs_steps=n_obs_steps,
+                n_action_steps=n_action_steps,
+                max_episode_steps=max_steps
+            )
+
+        env_fns = [env_fn] * n_envs
+        env_seeds = list()
+        env_prefixs = list()
+        env_init_fn_dills = list()
+
+        # train
+        with h5py.File(dataset_path, 'r') as f:
+            for i in range(n_train):
+                train_idx = train_start_idx + i
+                enable_render = i < n_train_vis
+                init_state = f[f'data/demo_{train_idx}/states'][0]
+
+                def init_fn(env, init_state=init_state, 
+                    enable_render=enable_render):
+                    # setup rendering
+                    # video_wrapper
+                    assert isinstance(env.env, VideoRecordingWrapper)
+                    env.env.video_recoder.stop()
+                    env.env.file_path = None
+                    if enable_render:
+                        filename = pathlib.Path(output_dir).joinpath(
+                            'media', wv.util.generate_id() + ".mp4")
+                        filename.parent.mkdir(parents=False, exist_ok=True)
+                        filename = str(filename)
+                        env.env.file_path = filename
+
+                    # switch to init_state reset
+                    assert isinstance(env.env.env, RobomimicImageWrapper)
+                    env.env.env.init_state = init_state
+
+                env_seeds.append(train_idx)
+                env_prefixs.append('train/')
+                env_init_fn_dills.append(dill.dumps(init_fn))
+        
+        # test
+        for i in range(n_test):
+            seed = test_start_seed + i
+            enable_render = i < n_test_vis
+
+            def init_fn(env, seed=seed, 
+                enable_render=enable_render):
+                # setup rendering
+                # video_wrapper
+                assert isinstance(env.env, VideoRecordingWrapper)
+                env.env.video_recoder.stop()
+                env.env.file_path = None
+                if enable_render:
+                    filename = pathlib.Path(output_dir).joinpath(
+                        'media', wv.util.generate_id() + ".mp4")
+                    filename.parent.mkdir(parents=False, exist_ok=True)
+                    filename = str(filename)
+                    env.env.file_path = filename
+
+                # switch to seed reset
+                assert isinstance(env.env.env, RobomimicImageWrapper)
+                env.env.env.init_state = None
+                env.seed(seed)
+
+            env_seeds.append(seed)
+            env_prefixs.append('test/')
+            env_init_fn_dills.append(dill.dumps(init_fn))
+
+        env = AsyncVectorEnv(env_fns, dummy_env_fn=dummy_env_fn)
+
+        self.env_meta = env_meta
+        self.env = env
+        self.env_fns = env_fns
+        self.env_seeds = env_seeds
+        self.env_prefixs = env_prefixs
+        self.env_init_fn_dills = env_init_fn_dills
+        self.fps = fps
+        self.crf = crf
+        self.n_obs_steps = n_obs_steps
+        self.n_action_steps = n_action_steps
+        self.past_action = past_action
+        self.max_steps = max_steps
+        self.rotation_transformer = rotation_transformer
+        self.abs_action = abs_action
+        self.tqdm_interval_sec = tqdm_interval_sec
+        self.max_rewards = {}
+        for prefix in self.env_prefixs:
+            self.max_rewards[prefix] = 0
+
+    def run(self, policy: BaseImagePolicy):
+        device = policy.device
+        dtype = policy.dtype
+        env = self.env
+        
+        # plan for rollout
+        n_envs = len(self.env_fns)
+        n_inits = len(self.env_init_fn_dills)
+        n_chunks = math.ceil(n_inits / n_envs)
+
+        # allocate data
+        all_video_paths = [None] * n_inits
+        all_rewards = [None] * n_inits
+
+        for chunk_idx in range(n_chunks):
+            start = chunk_idx * n_envs
+            end = min(n_inits, start + n_envs)
+            this_global_slice = slice(start, end)
+            this_n_active_envs = end - start
+            this_local_slice = slice(0,this_n_active_envs)
+            
+            this_init_fns = self.env_init_fn_dills[this_global_slice]
+            n_diff = n_envs - len(this_init_fns)
+            if n_diff > 0:
+                this_init_fns.extend([self.env_init_fn_dills[0]]*n_diff)
+            assert len(this_init_fns) == n_envs
+
+            # init envs
+            env.call_each('run_dill_function', 
+                args_list=[(x,) for x in this_init_fns])
+
+            # start rollout
+            obs = env.reset()
+            past_action = None
+            policy.reset()
+
+            env_name = self.env_meta['env_name']
+            pbar = tqdm.tqdm(total=self.max_steps, desc=f"Eval {env_name}Image {chunk_idx+1}/{n_chunks}", 
+                leave=False, mininterval=self.tqdm_interval_sec)
+            
+            done = False
+            while not done:
+                # create obs dict
+                np_obs_dict = dict(obs)
+                if self.past_action and (past_action is not None):
+                    # TODO: not tested
+                    np_obs_dict['past_action'] = past_action[
+                        :,-(self.n_obs_steps-1):].astype(np.float32)
+                
+                # device transfer
+                obs_dict = dict_apply(np_obs_dict, 
+                    lambda x: torch.from_numpy(x).to(
+                        device=device))
+
+                # run policy
+                with torch.no_grad():
+                    action_dict = policy.predict_action(obs_dict)
+
+                # device_transfer
+                np_action_dict = dict_apply(action_dict,
+                    lambda x: x.detach().to('cpu').numpy())
+
+                action = np_action_dict['action']
+                if not np.all(np.isfinite(action)):
+                    print(action)
+                    raise RuntimeError("Nan or Inf action")
+                
+                # step env
+                env_action = action
+                if self.abs_action:
+                    env_action = self.undo_transform_action(action)
+
+                obs, reward, done, info = env.step(env_action)
+                done = np.all(done)
+                past_action = action
+
+                # update pbar
+                pbar.update(action.shape[1])
+            pbar.close()
+
+            # collect data for this round
+            all_video_paths[this_global_slice] = env.render()[this_local_slice]
+            all_rewards[this_global_slice] = env.call('get_attr', 'reward')[this_local_slice]
+        # clear out video buffer
+        _ = env.reset()
+        
+        # log
+        max_rewards = collections.defaultdict(list)
+        log_data = dict()
+        # results reported in the paper are generated using the commented out line below
+        # which will only report and average metrics from first n_envs initial condition and seeds
+        # fortunately this won't invalidate our conclusion since
+        # 1. This bug only affects the variance of metrics, not their mean
+        # 2. All baseline methods are evaluated using the same code
+        # to completely reproduce reported numbers, uncomment this line:
+        # for i in range(len(self.env_fns)):
+        # and comment out this line
+        for i in range(n_inits):
+            seed = self.env_seeds[i]
+            prefix = self.env_prefixs[i]
+            max_reward = np.max(all_rewards[i])
+            max_rewards[prefix].append(max_reward)
+            log_data[prefix+f'sim_max_reward_{seed}'] = max_reward
+
+            # visualize sim
+            video_path = all_video_paths[i]
+            if video_path is not None:
+                sim_video = wandb.Video(video_path)
+                log_data[prefix+f'sim_video_{seed}'] = sim_video
+        
+        # log aggregate metrics
+        for prefix, value in max_rewards.items():
+            name = prefix+'mean_score'
+            value = np.mean(value)
+            log_data[name] = value
+            self.max_rewards[prefix] = max(self.max_rewards[prefix], value)
+            log_data[prefix+'max_score'] = self.max_rewards[prefix]
+
+        return log_data
+
+    def undo_transform_action(self, action):
+        raw_shape = action.shape
+        if raw_shape[-1] == 20:
+            # dual arm
+            action = action.reshape(-1,2,10)
+
+        d_rot = action.shape[-1] - 4
+        pos = action[...,:3]
+        rot = action[...,3:3+d_rot]
+        gripper = action[...,[-1]]
+        rot = self.rotation_transformer.inverse(rot)
+        uaction = np.concatenate([
+            pos, rot, gripper
+        ], axis=-1)
+
+        if raw_shape[-1] == 20:
+            # dual arm
+            uaction = uaction.reshape(*raw_shape[:-1], 14)
+
+        return uaction

equidiff/equi_diffpo/env_runner/robomimic_lowdim_runner.py

@@ -0,0 +1,405 @@
+import os
+import wandb
+import numpy as np
+import torch
+import collections
+import pathlib
+import tqdm
+import h5py
+import dill
+import math
+import wandb.sdk.data_types.video as wv
+from equi_diffpo.gym_util.async_vector_env import AsyncVectorEnv
+# from equi_diffpo.gym_util.sync_vector_env import SyncVectorEnv
+from equi_diffpo.gym_util.multistep_wrapper import MultiStepWrapper
+from equi_diffpo.gym_util.video_recording_wrapper import VideoRecordingWrapper, VideoRecorder
+from equi_diffpo.model.common.rotation_transformer import RotationTransformer
+
+from equi_diffpo.policy.base_lowdim_policy import BaseLowdimPolicy
+from equi_diffpo.common.pytorch_util import dict_apply
+from equi_diffpo.env_runner.base_lowdim_runner import BaseLowdimRunner
+from equi_diffpo.env.robomimic.robomimic_lowdim_wrapper import RobomimicLowdimWrapper
+import robomimic.utils.file_utils as FileUtils
+import robomimic.utils.env_utils as EnvUtils
+import robomimic.utils.obs_utils as ObsUtils
+
+
+def create_env(env_meta, obs_keys, enable_render=True):
+    ObsUtils.initialize_obs_modality_mapping_from_dict(
+        {'low_dim': obs_keys})
+    env = EnvUtils.create_env_from_metadata(
+        env_meta=env_meta,
+        render=False, 
+        # only way to not show collision geometry
+        # is to enable render_offscreen
+        # which uses a lot of RAM.
+        render_offscreen=enable_render,
+        use_image_obs=enable_render, 
+    )
+    return env
+
+
+class RobomimicLowdimRunner(BaseLowdimRunner):
+    """
+    Robomimic envs already enforces number of steps.
+    """
+
+    def __init__(self, 
+            output_dir,
+            dataset_path,
+            obs_keys,
+            n_train=10,
+            n_train_vis=3,
+            train_start_idx=0,
+            n_test=22,
+            n_test_vis=6,
+            test_start_seed=10000,
+            max_steps=400,
+            n_obs_steps=2,
+            n_action_steps=8,
+            n_latency_steps=0,
+            render_hw=(256,256),
+            render_camera_name='agentview',
+            fps=10,
+            crf=22,
+            past_action=False,
+            abs_action=False,
+            tqdm_interval_sec=5.0,
+            n_envs=None
+        ):
+        """
+        Assuming:
+        n_obs_steps=2
+        n_latency_steps=3
+        n_action_steps=4
+        o: obs
+        i: inference
+        a: action
+        Batch t:
+        |o|o| | | | | | |
+        | |i|i|i| | | | |
+        | | | | |a|a|a|a|
+        Batch t+1
+        | | | | |o|o| | | | | | |
+        | | | | | |i|i|i| | | | |
+        | | | | | | | | |a|a|a|a|
+        """
+
+        super().__init__(output_dir)
+
+        if n_envs is None:
+            n_envs = n_train + n_test
+
+        # handle latency step
+        # to mimic latency, we request n_latency_steps additional steps 
+        # of past observations, and the discard the last n_latency_steps
+        env_n_obs_steps = n_obs_steps + n_latency_steps
+        env_n_action_steps = n_action_steps
+
+        # assert n_obs_steps <= n_action_steps
+        dataset_path = os.path.expanduser(dataset_path)
+        robosuite_fps = 20
+        steps_per_render = max(robosuite_fps // fps, 1)
+
+        # read from dataset
+        env_meta = FileUtils.get_env_metadata_from_dataset(
+            dataset_path)
+        rotation_transformer = None
+        if abs_action:
+            env_meta['env_kwargs']['controller_configs']['control_delta'] = False
+            rotation_transformer = RotationTransformer('axis_angle', 'rotation_6d')
+        
+        def env_fn():
+            robomimic_env = create_env(
+                env_meta=env_meta, 
+                obs_keys=obs_keys
+            )
+            # Robosuite's hard reset causes excessive memory consumption.
+            # Disabled to run more envs.
+            # https://github.com/ARISE-Initiative/robosuite/blob/92abf5595eddb3a845cd1093703e5a3ccd01e77e/robosuite/environments/base.py#L247-L248
+            robomimic_env.env.hard_reset = False
+            return MultiStepWrapper(
+                VideoRecordingWrapper(
+                    RobomimicLowdimWrapper(
+                        env=robomimic_env,
+                        obs_keys=obs_keys,
+                        init_state=None,
+                        render_hw=render_hw,
+                        render_camera_name=render_camera_name
+                    ),
+                    video_recoder=VideoRecorder.create_h264(
+                        fps=fps,
+                        codec='h264',
+                        input_pix_fmt='rgb24',
+                        crf=crf,
+                        thread_type='FRAME',
+                        thread_count=1
+                    ),
+                    file_path=None,
+                    steps_per_render=steps_per_render
+                ),
+                n_obs_steps=n_obs_steps,
+                n_action_steps=n_action_steps,
+                max_episode_steps=max_steps
+            )
+        
+        # For each process the OpenGL context can only be initialized once
+        # Since AsyncVectorEnv uses fork to create worker process,
+        # a separate env_fn that does not create OpenGL context (enable_render=False)
+        # is needed to initialize spaces.
+        def dummy_env_fn():
+            robomimic_env = create_env(
+                    env_meta=env_meta, 
+                    obs_keys=obs_keys,
+                    enable_render=False
+            )
+            return MultiStepWrapper(
+                VideoRecordingWrapper(
+                    RobomimicLowdimWrapper(
+                        env=robomimic_env,
+                        obs_keys=obs_keys,
+                        init_state=None,
+                        render_hw=render_hw,
+                        render_camera_name=render_camera_name
+                    ),
+                    video_recoder=VideoRecorder.create_h264(
+                        fps=fps,
+                        codec='h264',
+                        input_pix_fmt='rgb24',
+                        crf=crf,
+                        thread_type='FRAME',
+                        thread_count=1
+                    ),
+                    file_path=None,
+                    steps_per_render=steps_per_render
+                ),
+                n_obs_steps=n_obs_steps,
+                n_action_steps=n_action_steps,
+                max_episode_steps=max_steps
+            )
+
+        env_fns = [env_fn] * n_envs
+        env_seeds = list()
+        env_prefixs = list()
+        env_init_fn_dills = list()
+
+        # train
+        with h5py.File(dataset_path, 'r') as f:
+            for i in range(n_train):
+                train_idx = train_start_idx + i
+                enable_render = i < n_train_vis
+                init_state = f[f'data/demo_{train_idx}/states'][0]
+
+                def init_fn(env, init_state=init_state, 
+                    enable_render=enable_render):
+                    # setup rendering
+                    # video_wrapper
+                    assert isinstance(env.env, VideoRecordingWrapper)
+                    env.env.video_recoder.stop()
+                    env.env.file_path = None
+                    if enable_render:
+                        filename = pathlib.Path(output_dir).joinpath(
+                            'media', wv.util.generate_id() + ".mp4")
+                        filename.parent.mkdir(parents=False, exist_ok=True)
+                        filename = str(filename)
+                        env.env.file_path = filename
+
+                    # switch to init_state reset
+                    assert isinstance(env.env.env, RobomimicLowdimWrapper)
+                    env.env.env.init_state = init_state
+
+                env_seeds.append(train_idx)
+                env_prefixs.append('train/')
+                env_init_fn_dills.append(dill.dumps(init_fn))
+        
+        # test
+        for i in range(n_test):
+            seed = test_start_seed + i
+            enable_render = i < n_test_vis
+
+            def init_fn(env, seed=seed, 
+                enable_render=enable_render):
+                # setup rendering
+                # video_wrapper
+                assert isinstance(env.env, VideoRecordingWrapper)
+                env.env.video_recoder.stop()
+                env.env.file_path = None
+                if enable_render:
+                    filename = pathlib.Path(output_dir).joinpath(
+                        'media', wv.util.generate_id() + ".mp4")
+                    filename.parent.mkdir(parents=False, exist_ok=True)
+                    filename = str(filename)
+                    env.env.file_path = filename
+
+                # switch to seed reset
+                assert isinstance(env.env.env, RobomimicLowdimWrapper)
+                env.env.env.init_state = None
+                env.seed(seed)
+
+            env_seeds.append(seed)
+            env_prefixs.append('test/')
+            env_init_fn_dills.append(dill.dumps(init_fn))
+        
+        env = AsyncVectorEnv(env_fns, dummy_env_fn=dummy_env_fn)
+        # env = SyncVectorEnv(env_fns)
+
+        self.env_meta = env_meta
+        self.env = env
+        self.env_fns = env_fns
+        self.env_seeds = env_seeds
+        self.env_prefixs = env_prefixs
+        self.env_init_fn_dills = env_init_fn_dills
+        self.fps = fps
+        self.crf = crf
+        self.n_obs_steps = n_obs_steps
+        self.n_action_steps = n_action_steps
+        self.n_latency_steps = n_latency_steps
+        self.env_n_obs_steps = env_n_obs_steps
+        self.env_n_action_steps = env_n_action_steps
+        self.past_action = past_action
+        self.max_steps = max_steps
+        self.rotation_transformer = rotation_transformer
+        self.abs_action = abs_action
+        self.tqdm_interval_sec = tqdm_interval_sec
+
+    def run(self, policy: BaseLowdimPolicy):
+        device = policy.device
+        dtype = policy.dtype
+        env = self.env
+        
+        # plan for rollout
+        n_envs = len(self.env_fns)
+        n_inits = len(self.env_init_fn_dills)
+        n_chunks = math.ceil(n_inits / n_envs)
+
+        # allocate data
+        all_video_paths = [None] * n_inits
+        all_rewards = [None] * n_inits
+
+        for chunk_idx in range(n_chunks):
+            start = chunk_idx * n_envs
+            end = min(n_inits, start + n_envs)
+            this_global_slice = slice(start, end)
+            this_n_active_envs = end - start
+            this_local_slice = slice(0,this_n_active_envs)
+            
+            this_init_fns = self.env_init_fn_dills[this_global_slice]
+            n_diff = n_envs - len(this_init_fns)
+            if n_diff > 0:
+                this_init_fns.extend([self.env_init_fn_dills[0]]*n_diff)
+            assert len(this_init_fns) == n_envs
+
+            # init envs
+            env.call_each('run_dill_function', 
+                args_list=[(x,) for x in this_init_fns])
+
+            # start rollout
+            obs = env.reset()
+            past_action = None
+            policy.reset()
+
+            env_name = self.env_meta['env_name']
+            pbar = tqdm.tqdm(total=self.max_steps, desc=f"Eval {env_name}Lowdim {chunk_idx+1}/{n_chunks}", 
+                leave=False, mininterval=self.tqdm_interval_sec)
+
+            done = False
+            while not done:
+                # create obs dict
+                np_obs_dict = {
+                    # handle n_latency_steps by discarding the last n_latency_steps
+                    'obs': obs[:,:self.n_obs_steps].astype(np.float32)
+                }
+                if self.past_action and (past_action is not None):
+                    # TODO: not tested
+                    np_obs_dict['past_action'] = past_action[
+                        :,-(self.n_obs_steps-1):].astype(np.float32)
+                
+                # device transfer
+                obs_dict = dict_apply(np_obs_dict, 
+                    lambda x: torch.from_numpy(x).to(
+                        device=device))
+
+                # run policy
+                with torch.no_grad():
+                    action_dict = policy.predict_action(obs_dict)
+
+                # device_transfer
+                np_action_dict = dict_apply(action_dict,
+                    lambda x: x.detach().to('cpu').numpy())
+
+                # handle latency_steps, we discard the first n_latency_steps actions
+                # to simulate latency
+                action = np_action_dict['action'][:,self.n_latency_steps:]
+                if not np.all(np.isfinite(action)):
+                    print(action)
+                    raise RuntimeError("Nan or Inf action")
+                
+                # step env
+                env_action = action
+                if self.abs_action:
+                    env_action = self.undo_transform_action(action)
+
+                obs, reward, done, info = env.step(env_action)
+                done = np.all(done)
+                past_action = action
+
+                # update pbar
+                pbar.update(action.shape[1])
+            pbar.close()
+
+            # collect data for this round
+            all_video_paths[this_global_slice] = env.render()[this_local_slice]
+            all_rewards[this_global_slice] = env.call('get_attr', 'reward')[this_local_slice]
+
+        # log
+        max_rewards = collections.defaultdict(list)
+        log_data = dict()
+        # results reported in the paper are generated using the commented out line below
+        # which will only report and average metrics from first n_envs initial condition and seeds
+        # fortunately this won't invalidate our conclusion since
+        # 1. This bug only affects the variance of metrics, not their mean
+        # 2. All baseline methods are evaluated using the same code
+        # to completely reproduce reported numbers, uncomment this line:
+        # for i in range(len(self.env_fns)):
+        # and comment out this line
+        for i in range(n_inits):
+            seed = self.env_seeds[i]
+            prefix = self.env_prefixs[i]
+            max_reward = np.max(all_rewards[i])
+            max_rewards[prefix].append(max_reward)
+            log_data[prefix+f'sim_max_reward_{seed}'] = max_reward
+
+            # visualize sim
+            video_path = all_video_paths[i]
+            if video_path is not None:
+                sim_video = wandb.Video(video_path)
+                log_data[prefix+f'sim_video_{seed}'] = sim_video
+
+        # log aggregate metrics
+        for prefix, value in max_rewards.items():
+            name = prefix+'mean_score'
+            value = np.mean(value)
+            log_data[name] = value
+
+        return log_data
+
+    def undo_transform_action(self, action):
+        raw_shape = action.shape
+        if raw_shape[-1] == 20:
+            # dual arm
+            action = action.reshape(-1,2,10)
+
+        d_rot = action.shape[-1] - 4
+        pos = action[...,:3]
+        rot = action[...,3:3+d_rot]
+        gripper = action[...,[-1]]
+        rot = self.rotation_transformer.inverse(rot)
+        uaction = np.concatenate([
+            pos, rot, gripper
+        ], axis=-1)
+
+        if raw_shape[-1] == 20:
+            # dual arm
+            uaction = uaction.reshape(*raw_shape[:-1], 14)
+
+        return uaction

equidiff/equi_diffpo/policy/robomimic_image_policy.py

@@ -0,0 +1,142 @@
+from typing import Dict
+import torch
+from equi_diffpo.model.common.normalizer import LinearNormalizer
+from equi_diffpo.policy.base_image_policy import BaseImagePolicy
+from equi_diffpo.common.pytorch_util import dict_apply
+
+from robomimic.algo import algo_factory
+from robomimic.algo.algo import PolicyAlgo
+import robomimic.utils.obs_utils as ObsUtils
+from equi_diffpo.common.robomimic_config_util import get_robomimic_config
+
+class RobomimicImagePolicy(BaseImagePolicy):
+    def __init__(self, 
+            shape_meta: dict,
+            algo_name='bc_rnn',
+            obs_type='image',
+            task_name='square',
+            dataset_type='ph',
+            crop_shape=(76,76)
+        ):
+        super().__init__()
+
+        # parse shape_meta
+        action_shape = shape_meta['action']['shape']
+        assert len(action_shape) == 1
+        action_dim = action_shape[0]
+        obs_shape_meta = shape_meta['obs']
+        obs_config = {
+            'low_dim': [],
+            'rgb': [],
+            'depth': [],
+            'scan': []
+        }
+        obs_key_shapes = dict()
+        for key, attr in obs_shape_meta.items():
+            shape = attr['shape']
+            obs_key_shapes[key] = list(shape)
+
+            type = attr.get('type', 'low_dim')
+            if type == 'rgb':
+                obs_config['rgb'].append(key)
+            elif type == 'low_dim':
+                obs_config['low_dim'].append(key)
+            else:
+                raise RuntimeError(f"Unsupported obs type: {type}")
+
+        # get raw robomimic config
+        config = get_robomimic_config(
+            algo_name=algo_name,
+            hdf5_type=obs_type,
+            task_name=task_name,
+            dataset_type=dataset_type)
+
+        
+        with config.unlocked():
+            # set config with shape_meta
+            config.observation.modalities.obs = obs_config
+
+            if crop_shape is None:
+                for key, modality in config.observation.encoder.items():
+                    if modality.obs_randomizer_class == 'CropRandomizer':
+                        modality['obs_randomizer_class'] = None
+            else:
+                # set random crop parameter
+                ch, cw = crop_shape
+                for key, modality in config.observation.encoder.items():
+                    if modality.obs_randomizer_class == 'CropRandomizer':
+                        modality.obs_randomizer_kwargs.crop_height = ch
+                        modality.obs_randomizer_kwargs.crop_width = cw
+
+        # init global state
+        ObsUtils.initialize_obs_utils_with_config(config)
+
+        # load model
+        model: PolicyAlgo = algo_factory(
+                algo_name=config.algo_name,
+                config=config,
+                obs_key_shapes=obs_key_shapes,
+                ac_dim=action_dim,
+                device='cpu',
+            )
+
+        self.model = model
+        self.nets = model.nets
+        self.normalizer = LinearNormalizer()
+        self.config = config
+
+    def to(self,*args,**kwargs):
+        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
+        if device is not None:
+            self.model.device = device
+        super().to(*args,**kwargs)
+    
+    # =========== inference =============
+    def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        nobs_dict = self.normalizer(obs_dict)
+        robomimic_obs_dict = dict_apply(nobs_dict, lambda x: x[:,0,...])
+        naction = self.model.get_action(robomimic_obs_dict)
+        action = self.normalizer['action'].unnormalize(naction)
+        # (B, Da)
+        result = {
+            'action': action[:,None,:] # (B, 1, Da)
+        }
+        return result
+
+    def reset(self):
+        self.model.reset()
+
+    # =========== training ==============
+    def set_normalizer(self, normalizer: LinearNormalizer):
+        self.normalizer.load_state_dict(normalizer.state_dict())
+
+    def train_on_batch(self, batch, epoch, validate=False):
+        nobs = self.normalizer.normalize(batch['obs'])
+        nactions = self.normalizer['action'].normalize(batch['action'])
+        robomimic_batch = {
+            'obs': nobs,
+            'actions': nactions
+        }
+        input_batch = self.model.process_batch_for_training(
+            robomimic_batch)
+        info = self.model.train_on_batch(
+            batch=input_batch, epoch=epoch, validate=validate)
+        # keys: losses, predictions
+        return info
+    
+    def on_epoch_end(self, epoch):
+        self.model.on_epoch_end(epoch)
+
+    def get_optimizer(self):
+        return self.model.optimizers['policy']
+
+
+def test():
+    import os
+    from omegaconf import OmegaConf
+    cfg_path = os.path.expanduser('~/dev/diffusion_policy/diffusion_policy/config/task/lift_image.yaml')
+    cfg = OmegaConf.load(cfg_path)
+    shape_meta = cfg.shape_meta
+
+    policy = RobomimicImagePolicy(shape_meta=shape_meta)
+

equidiff/equi_diffpo/scripts/robomimic_dataset_action_comparison.py

@@ -0,0 +1,51 @@
+if __name__ == "__main__":
+    import sys
+    import os
+    import pathlib
+
+    ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent)
+    sys.path.append(ROOT_DIR)
+
+import os
+import click
+import pathlib
+import h5py
+import numpy as np
+from tqdm import tqdm
+from scipy.spatial.transform import Rotation
+
+def read_all_actions(hdf5_file, metric_skip_steps=1):
+    n_demos = len(hdf5_file['data'])
+    all_actions = list()
+    for i in tqdm(range(n_demos)):
+        actions = hdf5_file[f'data/demo_{i}/actions'][:]
+        all_actions.append(actions[metric_skip_steps:])
+    all_actions = np.concatenate(all_actions, axis=0)
+    return all_actions
+
+
+@click.command()
+@click.option('-i', '--input', required=True, help='input hdf5 path')
+@click.option('-o', '--output', required=True, help='output hdf5 path. Parent directory must exist')
+def main(input, output):
+    # process inputs
+    input = pathlib.Path(input).expanduser()
+    assert input.is_file()
+    output = pathlib.Path(output).expanduser()
+    assert output.is_file()
+
+    input_file = h5py.File(str(input), 'r')
+    output_file = h5py.File(str(output), 'r')
+
+    input_all_actions = read_all_actions(input_file)
+    output_all_actions = read_all_actions(output_file)
+    pos_dist = np.linalg.norm(input_all_actions[:,:3] - output_all_actions[:,:3], axis=-1)
+    rot_dist = (Rotation.from_rotvec(input_all_actions[:,3:6]
+        ) * Rotation.from_rotvec(output_all_actions[:,3:6]).inv()
+        ).magnitude()
+
+    print(f'max pos dist: {pos_dist.max()}')
+    print(f'max rot dist: {rot_dist.max()}')
+
+if __name__ == "__main__":
+    main()

equidiff/equi_diffpo/scripts/robomimic_dataset_conversion.py

@@ -0,0 +1,103 @@
+if __name__ == "__main__":
+    import sys
+    import os
+    import pathlib
+
+    ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent)
+    sys.path.append(ROOT_DIR)
+
+import multiprocessing
+import os
+import shutil
+import click
+import pathlib
+import h5py
+from tqdm import tqdm
+import collections
+import pickle
+from equi_diffpo.common.robomimic_util import RobomimicAbsoluteActionConverter
+
+def worker(x):
+    path, idx, do_eval = x
+    converter = RobomimicAbsoluteActionConverter(path)
+    if do_eval:
+        abs_actions, info = converter.convert_and_eval_idx(idx)
+    else:
+        abs_actions = converter.convert_idx(idx)
+        info = dict()
+    return abs_actions, info
+
+@click.command()
+@click.option('-i', '--input', required=True, help='input hdf5 path')
+@click.option('-o', '--output', required=True, help='output hdf5 path. Parent directory must exist')
+@click.option('-e', '--eval_dir', default=None, help='directory to output evaluation metrics')
+@click.option('-n', '--num_workers', default=None, type=int)
+def main(input, output, eval_dir, num_workers):
+    # process inputs
+    input = pathlib.Path(input).expanduser()
+    assert input.is_file()
+    output = pathlib.Path(output).expanduser()
+    assert output.parent.is_dir()
+    assert not output.is_dir()
+
+    do_eval = False
+    if eval_dir is not None:
+        eval_dir = pathlib.Path(eval_dir).expanduser()
+        assert eval_dir.parent.exists()
+        do_eval = True
+    
+    converter = RobomimicAbsoluteActionConverter(input)
+
+    # run
+    with multiprocessing.Pool(num_workers) as pool:
+        results = pool.map(worker, [(input, i, do_eval) for i in range(len(converter))])
+    
+    # save output
+    print('Copying hdf5')
+    shutil.copy(str(input), str(output))
+
+    # modify action
+    with h5py.File(output, 'r+') as out_file:
+        for i in tqdm(range(len(converter)), desc="Writing to output"):
+            abs_actions, info = results[i]
+            demo = out_file[f'data/demo_{i}']
+            demo['actions'][:] = abs_actions
+    
+    # save eval
+    if do_eval:
+        eval_dir.mkdir(parents=False, exist_ok=True)
+
+        print("Writing error_stats.pkl")
+        infos = [info for _, info in results]
+        pickle.dump(infos, eval_dir.joinpath('error_stats.pkl').open('wb'))
+
+        print("Generating visualization")
+        metrics = ['pos', 'rot']
+        metrics_dicts = dict()
+        for m in metrics:
+            metrics_dicts[m] = collections.defaultdict(list)
+
+        for i in range(len(infos)):
+            info = infos[i]
+            for k, v in info.items():
+                for m in metrics:
+                    metrics_dicts[m][k].append(v[m])
+
+        from matplotlib import pyplot as plt
+        plt.switch_backend('PDF')
+
+        fig, ax = plt.subplots(1, len(metrics))
+        for i in range(len(metrics)):
+            axis = ax[i]
+            data = metrics_dicts[metrics[i]]
+            for key, value in data.items():
+                axis.plot(value, label=key)
+            axis.legend()
+            axis.set_title(metrics[i])
+        fig.set_size_inches(10,4)
+        fig.savefig(str(eval_dir.joinpath('error_stats.pdf')))
+        fig.savefig(str(eval_dir.joinpath('error_stats.png')))
+
+
+if __name__ == "__main__":
+    main()

equidiff/equi_diffpo/scripts/robomimic_dataset_obs_conversion.py

@@ -0,0 +1,69 @@
+if __name__ == "__main__":
+    import sys
+    import os
+    import pathlib
+
+    ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent)
+    sys.path.append(ROOT_DIR)
+
+import multiprocessing
+import os
+import shutil
+import click
+import pathlib
+import h5py
+from tqdm import tqdm
+import numpy as np
+import collections
+import pickle
+from equi_diffpo.common.robomimic_util import RobomimicObsConverter
+
+multiprocessing.set_start_method('spawn', force=True)
+
+def worker(x):
+    path, idx = x
+    converter = RobomimicObsConverter(path)
+    obss = converter.convert_idx(idx)
+    return obss
+
+@click.command()
+@click.option('-i', '--input', required=True, help='input hdf5 path')
+@click.option('-o', '--output', required=True, help='output hdf5 path. Parent directory must exist')
+@click.option('-n', '--num_workers', default=None, type=int)
+def main(input, output, num_workers):
+    # process inputs
+    input = pathlib.Path(input).expanduser()
+    assert input.is_file()
+    output = pathlib.Path(output).expanduser()
+    assert output.parent.is_dir()
+    assert not output.is_dir()
+    
+    converter = RobomimicObsConverter(input)
+
+    # save output
+    print('Copying hdf5')
+    shutil.copy(str(input), str(output))
+
+    # run
+    idx = 0
+    while idx < len(converter):
+        with multiprocessing.Pool(num_workers) as pool:
+            end = min(idx + num_workers, len(converter))
+            results = pool.map(worker, [(input, i) for i in range(idx, end)])
+
+        # modify action
+        print('Writing {} to {}'.format(idx, end))
+        with h5py.File(output, 'r+') as out_file:
+            for i in tqdm(range(idx, end), desc="Writing to output"):
+                obss = results[i - idx]
+                demo = out_file[f'data/demo_{i}']
+                del demo['obs']
+                for k in obss:
+                    demo.create_dataset("obs/{}".format(k), data=np.array(obss[k]), compression="gzip")
+        
+        idx = end
+        del results
+
+
+if __name__ == "__main__":
+    main()