zarr2hdf5.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# diffusion policy zarr fromat to robomimic hdf5 format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import zarr \n",
    "import h5py\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data\n",
      "meta\n",
      "meta: episode_ends (206,)\n",
      "data: action (25650, 2)\n",
      "data: img (25650, 96, 96, 3)\n",
      "data: keypoint (25650, 9, 2)\n",
      "data: n_contacts (25650, 1)\n",
      "data: state (25650, 5)\n"
     ]
    }
   ],
   "source": [
    "dataset_path = \"data/pusht/pusht_cchi_v7_replay.zarr\"\n",
    "dataset_zarr = zarr.open(dataset_path, 'r')\n",
    "\n",
    "for key in dataset_zarr.keys():\n",
    "    print(key )\n",
    "\n",
    "for key in dataset_zarr['meta']:\n",
    "    print('meta:', key, dataset_zarr['meta'][key].shape)\n",
    "\n",
    "for key in dataset_zarr['data'].keys():\n",
    "    print('data:', key, dataset_zarr['data'][key].shape) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "episode_ends=dataset_zarr['meta']['episode_ends'][:]\n",
    "action=dataset_zarr['data']['action'][:]\n",
    "img=dataset_zarr['data']['img'][:].astype(np.uint8)\n",
    "keypoint=dataset_zarr['data']['keypoint'][:]\n",
    "n_contacts=dataset_zarr['data']['n_contacts'][:]\n",
    "state=dataset_zarr['data']['state'][:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "episode_starts = np.concatenate(([0], episode_ends[:-1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([   0,  161,  279,  420,  579,  738,  895,  964, 1133, 1213]),\n",
       " array([ 161,  279,  420,  579,  738,  895,  964, 1133, 1213, 1347]))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "episode_starts[:10], episode_ends[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([25059, 25151, 25318, 25468, 25601]),\n",
       " array([25151, 25318, 25468, 25601, 25650]))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "episode_starts[-5:], episode_ends[-5:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_episode(episode_idx):\n",
    "    start = episode_starts[episode_idx]\n",
    "    end = episode_ends[episode_idx]\n",
    "    actions=action[start:end]\n",
    "    dactions = np.diff(actions, axis=0)\n",
    "    dactions = np.concatenate([dactions, np.zeros((1, 2))], axis=0)\n",
    "    return {\n",
    "        'action': action[start:end],\n",
    "        'daction': dactions,\n",
    "        'img': img[start:end],\n",
    "        'keypoint': keypoint[start:end],\n",
    "        'n_contacts': n_contacts[start:end],\n",
    "        'state': state[start:end],\n",
    "        'agent_pos': state[start:end, :2],\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((161, 2), (161, 2), (161, 96, 96, 3))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "episode = get_episode(0)\n",
    "episode['action'].shape, episode['daction'].shape, episode['img'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "206"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n_episode=len(episode_starts)\n",
    "n_episode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "hdf5_file_name='pusht_v7_zarr_206.hdf5'\n",
    "\n",
    "env_args = \"\" \n",
    "f_sub = h5py.File(hdf5_file_name, \"w\")\n",
    "f_sub.create_group(\"data\")\n",
    "f_sub.create_group(\"mask\")\n",
    "f_sub[\"data\"].attrs[\"env_args\"] = env_args\n",
    "f_sub.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "saving demo /data/demo_0 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_1 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_2 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_3 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_4 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_5 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_6 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_7 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_8 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_9 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_10 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_11 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_12 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_13 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_14 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_15 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_16 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_17 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_18 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_19 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_20 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_21 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_22 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_23 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_24 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_25 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_26 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_27 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_28 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_29 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_30 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_31 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_32 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_33 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_34 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_35 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_36 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_37 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_38 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_39 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_40 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_41 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_42 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_43 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_44 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_45 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_46 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_47 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_48 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_49 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_50 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_51 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_52 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_53 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_54 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_55 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_56 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_57 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_58 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_59 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_60 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_61 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_62 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_63 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_64 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_65 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_66 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_67 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_68 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_69 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_70 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_71 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_72 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_73 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_74 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_75 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_76 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_77 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_78 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_79 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_80 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_81 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_82 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_83 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_84 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_85 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_86 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_87 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_88 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_89 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_90 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_91 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_92 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_93 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_94 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_95 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_96 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_97 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_98 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_99 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_100 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_101 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_102 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_103 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_104 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_105 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_106 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_107 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_108 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_109 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_110 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_111 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_112 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_113 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_114 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_115 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_116 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_117 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_118 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_119 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_120 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_121 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_122 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_123 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_124 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_125 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_126 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_127 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_128 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_129 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_130 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_131 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_132 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_133 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_134 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_135 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_136 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_137 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_138 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_139 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_140 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_141 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_142 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_143 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_144 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_145 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_146 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_147 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_148 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_149 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_150 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_151 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_152 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_153 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_154 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_155 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_156 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_157 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_158 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_159 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_160 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_161 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_162 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_163 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_164 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_165 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_166 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_167 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_168 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_169 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_170 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_171 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_172 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_173 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_174 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_175 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_176 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_177 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_178 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_179 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_180 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_181 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_182 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_183 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_184 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_185 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_186 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_187 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_188 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_189 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_190 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_191 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_192 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_193 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_194 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_195 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_196 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_197 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_198 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_199 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_200 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_201 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_202 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_203 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_204 to pusht_v7_zarr_206.hdf5\n",
      "saving demo /data/demo_205 to pusht_v7_zarr_206.hdf5\n"
     ]
    }
   ],
   "source": [
    "for demo_no in range(n_episode):\n",
    "    episode = get_episode(demo_no)\n",
    "\n",
    "    demo_group=f\"/data/demo_{demo_no}\"\n",
    "    print(f'saving demo {demo_group} to {hdf5_file_name}')\n",
    "    with h5py.File(hdf5_file_name, 'a') as hf:\n",
    "        group = hf.create_group(demo_group) \n",
    "        group.attrs['num_samples'] = episode['action'].shape[0]\n",
    "        group.create_dataset('obs/img', data=episode['img'])      \n",
    "        group.create_dataset('obs/state', data=episode['state'])\n",
    "        group.create_dataset('obs/keypoint', data=episode['keypoint'])\n",
    "        group.create_dataset('obs/n_contacts', data=episode['n_contacts'])\n",
    "        group.create_dataset('obs/agent_pos', data=episode['agent_pos'])\n",
    "        group.create_dataset('daction', data=episode['daction'])\n",
    "        group.create_dataset('action', data=   episode['action'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### load hdf5 file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<KeysViewHDF5 ['data', 'mask']>\n",
      "<KeysViewHDF5 ['demo_0', 'demo_1', 'demo_10', 'demo_100', 'demo_101', 'demo_102', 'demo_103', 'demo_104', 'demo_105', 'demo_106', 'demo_107', 'demo_108', 'demo_109', 'demo_11', 'demo_110', 'demo_111', 'demo_112', 'demo_113', 'demo_114', 'demo_115', 'demo_116', 'demo_117', 'demo_118', 'demo_119', 'demo_12', 'demo_120', 'demo_121', 'demo_122', 'demo_123', 'demo_124', 'demo_125', 'demo_126', 'demo_127', 'demo_128', 'demo_129', 'demo_13', 'demo_130', 'demo_131', 'demo_132', 'demo_133', 'demo_134', 'demo_135', 'demo_136', 'demo_137', 'demo_138', 'demo_139', 'demo_14', 'demo_140', 'demo_141', 'demo_142', 'demo_143', 'demo_144', 'demo_145', 'demo_146', 'demo_147', 'demo_148', 'demo_149', 'demo_15', 'demo_150', 'demo_151', 'demo_152', 'demo_153', 'demo_154', 'demo_155', 'demo_156', 'demo_157', 'demo_158', 'demo_159', 'demo_16', 'demo_160', 'demo_161', 'demo_162', 'demo_163', 'demo_164', 'demo_165', 'demo_166', 'demo_167', 'demo_168', 'demo_169', 'demo_17', 'demo_170', 'demo_171', 'demo_172', 'demo_173', 'demo_174', 'demo_175', 'demo_176', 'demo_177', 'demo_178', 'demo_179', 'demo_18', 'demo_180', 'demo_181', 'demo_182', 'demo_183', 'demo_184', 'demo_185', 'demo_186', 'demo_187', 'demo_188', 'demo_189', 'demo_19', 'demo_190', 'demo_191', 'demo_192', 'demo_193', 'demo_194', 'demo_195', 'demo_196', 'demo_197', 'demo_198', 'demo_199', 'demo_2', 'demo_20', 'demo_200', 'demo_201', 'demo_202', 'demo_203', 'demo_204', 'demo_205', 'demo_21', 'demo_22', 'demo_23', 'demo_24', 'demo_25', 'demo_26', 'demo_27', 'demo_28', 'demo_29', 'demo_3', 'demo_30', 'demo_31', 'demo_32', 'demo_33', 'demo_34', 'demo_35', 'demo_36', 'demo_37', 'demo_38', 'demo_39', 'demo_4', 'demo_40', 'demo_41', 'demo_42', 'demo_43', 'demo_44', 'demo_45', 'demo_46', 'demo_47', 'demo_48', 'demo_49', 'demo_5', 'demo_50', 'demo_51', 'demo_52', 'demo_53', 'demo_54', 'demo_55', 'demo_56', 'demo_57', 'demo_58', 'demo_59', 'demo_6', 'demo_60', 'demo_61', 'demo_62', 'demo_63', 'demo_64', 'demo_65', 'demo_66', 'demo_67', 'demo_68', 'demo_69', 'demo_7', 'demo_70', 'demo_71', 'demo_72', 'demo_73', 'demo_74', 'demo_75', 'demo_76', 'demo_77', 'demo_78', 'demo_79', 'demo_8', 'demo_80', 'demo_81', 'demo_82', 'demo_83', 'demo_84', 'demo_85', 'demo_86', 'demo_87', 'demo_88', 'demo_89', 'demo_9', 'demo_90', 'demo_91', 'demo_92', 'demo_93', 'demo_94', 'demo_95', 'demo_96', 'demo_97', 'demo_98', 'demo_99']>\n",
      "<KeysViewHDF5 ['action', 'daction', 'obs']>\n",
      "<KeysViewHDF5 ['agent_pos', 'img', 'keypoint', 'n_contacts', 'state']>\n",
      "(118, 96, 96, 3)\n",
      "(118, 5)\n"
     ]
    }
   ],
   "source": [
    "hdf5_file_name='pusht_v7_zarr_206.hdf5'\n",
    "\n",
    "f=h5py.File(hdf5_file_name, 'r')\n",
    "print(f.keys())\n",
    "print(f['data'].keys())\n",
    "print(f['data']['demo_1'].keys())\n",
    "print(f['data']['demo_1']['obs'].keys())\n",
    "print(f['data']['demo_1']['obs']['img'].shape)\n",
    "print(f['data']['demo_1']['obs']['state'].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "206"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(f['data'].keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(118, 96, 96, 3)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "img=f['data']['demo_1']['obs']['img'][:] \n",
    "img.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(65, 255)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "img0=img[10].astype(np.uint8)\n",
    "np.min(img0), np.max(img0)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f8f548bf370>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAaEAAAGgCAYAAAAD9NhnAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/av/WaAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAem0lEQVR4nO3df3BU9b3/8VcIsiSYrAg3m6wGDDNxUKIjJcoYGKFfJZ2KnTrcWhUUrJ1RCmhipvKjaI1cSID7vVxsLVgYv0CvpTCOdKSddkrUGuXSFsRGabgT2ppirhC22DQbhSZCPt8/KMeczS822c3nnOT5mNmZ/Zxz9uwnH5l9ec7nnPdJMcYYAQBgwTDbHQAADF2EEADAGkIIAGANIQQAsIYQAgBYQwgBAKwhhAAA1hBCAABrCCEAgDWEEADAmqSF0KZNm5SXl6eRI0dqypQpevvtt5P1VQAAnxqejJ3u3r1bpaWl2rRpk6ZNm6Yf/vCH+vKXv6yjR49q3LhxPX62vb1dJ06cUEZGhlJSUpLRPQBAEhlj1NLSonA4rGHDejnWMUlwyy23mIULF7qWTZw40SxfvrzXzzY0NBhJvHjx4sXL56+GhoZef/MTfiTU1tamw4cPa/ny5a7lxcXFOnDgQKftW1tb1dra6rTNP4t6NzQ0KDMz09kGAOBdgUDAeR+NRpWbm6uMjIxeP5fwEDp9+rTOnz+vUCjkWh4KhdTY2Nhp+8rKSj377LOdlmdmZhJCAOATHUPookuZUknKnFBXX26M6bJDK1asUFlZmdO+mKAdj5BOnDjh+syoUaOS0GMAwKX69NNPXe1wOOy8j+fAIeEhNHbsWKWmpnY66olEIp2OjqQL6dlVggIABr+EX6I9YsQITZkyRVVVVa7lVVVVKioqSvTXAQB8LCmn48rKyvTggw+qsLBQt956q7Zs2aIPP/xQCxcuTMbXAQB8KikhdO+99+rjjz/WqlWrdPLkSRUUFOgXv/iFxo8ff8n7MMY4V8rFzgFlZWUltL8AgPhEIhFX++Lvdez73qSYeLYeANFoVMFgUKdOnXKujotGo65tCCEAsCs2hC7+XksXfrNDoZCam5tdy7tC7TgAgDWEEADAGkIIAGANIQQAsIYQAgBYk7SyPf2VkpLCoxwAwCc6/l7H89vNkRAAwBpCCABgDSEEALDGs3NCHcv2AAC8ra9lezgSAgBYQwgBAKwhhAAA1hBCAABrCCEAgDWEEADAGkIIAGCNZ+8TonYcAPgHteMAAL5DCAEArCGEAADWeHZOiNpxAOAf1I4DAPgOIQQAsIYQAgBYQwgBAKwhhAAA1hBCAABrCCEAgDWevU+I2nEA4B/UjgMA+A4hBACwxrOn4yjbAwD+QdkeAIDvEEIAAGsIIQCANYQQAMAaQggAYA0hBACwhhACAFjj2fuEKNsDAP5B2R4AgO8QQgAAawghAIA1np0TonYcAPgHteMAAL5DCAEArCGEAADWEEIAAGsIIQCANYQQAMAaQggAYI1n7xOidhwA+Ae14wAAvkMIAQCsIYQAANZ4dk6I2nEA4B/UjgMA+A4hBACwhhACAFhDCAEArCGEAADWEEIAAGs8e4k2ZXsAwD8o2wMA8B1CCABgTVwhVFlZqZtvvlkZGRnKysrS3Xffrbq6Otc2xhiVl5crHA4rLS1NM2fOVG1tbUI7DQAYHOIKoerqai1evFi//e1vVVVVpXPnzqm4uFiffvqps8369eu1YcMGPf/88zp06JCys7M1a9YstbS0xNWxi2V7KN0DAN7X8Tc7nt/tFNOPX/m//vWvysrKUnV1tW677TYZYxQOh1VaWqply5ZJklpbWxUKhbRu3To9+uijve4zGo0qGAzq1KlTyszMdJZ1lJWV1dcuAwASIBKJuNoXf6+lC7/ZoVBIzc3NruVd6decUHNzsyTpyiuvlCTV19ersbFRxcXFzjaBQEAzZszQgQMHutxHa2urotGo6wUAGBr6HELGGJWVlWn69OkqKCiQJDU2NkqSQqGQa9tQKOSsi1VZWalgMOi8cnNz+9olAIDP9DmElixZovfff18/+clPOq2LvUbcGNPtdeMrVqxQc3Oz82poaOhrlwAAPtOnm1Ufe+wx7d27V2+99ZauvvpqZ3l2drakC0dEOTk5zvJIJNLp6OiiQCCgQCDQl24AAHwuriMhY4yWLFmiPXv26I033lBeXp5rfV5enrKzs1VVVeUsa2trU3V1tYqKihLTYwDAoBHXkdDixYu1c+dOvfrqq8rIyHDmeYLBoNLS0pSSkqLS0lJVVFQoPz9f+fn5qqioUHp6uubOnZuUPwAA4F9xhdDmzZslSTNnznQt37Ztmx566CFJ0tKlS3X27FktWrRITU1Nmjp1qvbt26eMjIy4OkbtOADwj77WjuvXfULJcPE+oUgk4lxffvFS8Iu4TwgA7Iq9TygYDDrvo9GosrKykn+fEAAA/UEIAQCs8ezzhKgbBwD+0fH3Op7fbo6EAADWEEIAAGsIIQCANYQQAMAaQggAYA0hBACwxrOXaFO2BwD8o69lezgSAgBYQwgBAKwhhAAA1nh2ToiyPQDgH5TtAQD4DiEEALCGEAIAWEMIAQCsIYQAANYQQgAAawghAIA1nr1PiNpxAOAf1I4DAPgOIQQAsIYQAgBY49k5IWrHAYB/UDsOAOA7hBAAwBpCCABgDSEEALCGEAIAWEMIAQCs8ewl2pTtAQD/oGwPAMB3CCEAgDWEEADAGs/OCVG2BwD8g7I9AADfIYQAANYQQgAAawghAIA1hBAAwBpCCABgDSEEALDGs/cJUTsOAPyD2nEAAN8hhAAA1hBCAABrPDsnRO04APAPascBAHyHEAIAWEMIAQCsIYQAANYQQgAAawghAIA1hBAAwBrP3idE7TgA8A9qxwEAfIcQAgBY49nTcZTtAQD/oGwPAMB3CCEAgDWEEADAGkIIAGANIQQAsIYQAgBY068QqqysVEpKikpLS51lxhiVl5crHA4rLS1NM2fOVG1tbX/7CQAYhPp8n9ChQ4e0ZcsW3Xjjja7l69ev14YNG7R9+3Zde+21Wr16tWbNmqW6ujplZGRc8v4p24Nke6XlFVf7XzP+1VJPAP8b0LI9n3zyiebNm6etW7dq9OjRznJjjDZu3KiVK1dqzpw5Kigo0I4dO3TmzBnt3LmzL18FABjE+hRCixcv1uzZs3XHHXe4ltfX16uxsVHFxcXOskAgoBkzZujAgQNd7qu1tVXRaNT1AgAMDXGfjtu1a5feffddHTp0qNO6xsZGSVIoFHItD4VCOn78eJf7q6ys1LPPPhtvNwAAg0BcIdTQ0KCSkhLt27dPI0eO7Ha72POBxphuzxGuWLFCZWVlTjsajSo3N5facUi455qe63H99ubtrvZDwYeS1xlgkOlr7bi4Qujw4cOKRCKaMmWKs+z8+fN666239Pzzz6uurk7ShSOinJwcZ5tIJNLp6OiiQCCgQCAQTzcAAINEXHNCt99+u44cOaKamhrnVVhYqHnz5qmmpkYTJkxQdna2qqqqnM+0tbWpurpaRUVFCe88AMDf4joSysjIUEFBgWvZqFGjNGbMGGd5aWmpKioqlJ+fr/z8fFVUVCg9PV1z585NXK8BAINCwp8ntHTpUp09e1aLFi1SU1OTpk6dqn379sV1jxCQCL3NAcVqbm92tZkjApKv3yH05ptvutopKSkqLy9XeXl5f3cNABjkqB0HALCGEAIAWJPwOaFEoXYc+iLeeaCexM4RUWsO6N6A1o4DACARCCEAgDWEEADAGs/OCVE7DpcikXNAvfnfc//rane8j4h7iDDU9bV2HEdCAABrCCEAgDWePR0HXIqS0SWu9pa/b3HenzVnk/rdHS/hpsQP0DccCQEArCGEAADWEEIAAGs8OydE2R70xSNXPOK87zg/JCV3jogSPxjqKNsDAPAdQggAYA0hBACwxrNzQpTtQX91nB+SvFPiR+I+Igw+lO0BAPgOIQQAsIYQAgBY49k5ISDReqozJw3sfUTMEQEXcCQEALCGEAIAWEMIAQCs8eycELXjkGyx9xFRaw7oO2rHAQB8hxACAFhDCAEArPHsnBC14zDQqDUH9B214wAAvkMIAQCsIYQAANZ4dk4ISLbY+4Bi7xOyqadac8wPYTDhSAgAYA0hBACwxrOn4yjbg0Tz8um33nQ8PUeJH3gRZXsAAL5DCAEArCGEAADWeHZOiLI9SDQ/zQH1hBI/8CLK9gAAfIcQAgBYQwgBAKzx7JwQ0Bcd7wUayDmgVKW62iNSRrjaA/mocOaI4CccCQEArCGEAADWEEIAAGs8OydE7ThcCq/Ug1syekmP62P7NZBzRNSaw0CgdhwAwHcIIQCANYQQAMAaz84JUTsOl8JmPbiS0SWXvO0jVzziaj/X9Fyiu9Ot2FpzQDJQOw4A4DuEEADAGkIIAGCNZ+eEAMnufUCx9eB6uxcoHrHzSYm8jyiZ/QYSjSMhAIA1hBAAwBpCCABgjWfnhKgdNzR5pRacNLBzKbH3EfVnjog5INhA7TgAgO8QQgAAazx7Os4rZXv+z6IPu103PucyV3vb0znJ7s6g55cyPMkWT5kfL/UbQxdlewAAvkMIAQCsiTuEPvroIz3wwAMaM2aM0tPTddNNN+nw4cPOemOMysvLFQ6HlZaWppkzZ6q2tjahnQYADA5xzQk1NTVp2rRp+uIXv6hf/vKXysrK0p///GddccUVzjbr16/Xhg0btH37dl177bVavXq1Zs2apbq6OmVkZCS6/wkXOweU8S/jut32b+fc7Yf/7aSr/f+YI+pksJbhSTbmfTBYxRVC69atU25urrZt2+Ysu+aaa5z3xhht3LhRK1eu1Jw5cyRJO3bsUCgU0s6dO/Xoo48mptcAgEEhrtNxe/fuVWFhoe655x5lZWVp8uTJ2rp1q7O+vr5ejY2NKi4udpYFAgHNmDFDBw4c6HKfra2tikajrhcAYGiIK4Q++OADbd68Wfn5+frVr36lhQsX6vHHH9ePfvQjSVJjY6MkKRQKuT4XCoWcdbEqKysVDAadV25ubl/+DgCAD8V1Oq69vV2FhYWqqKiQJE2ePFm1tbXavHmz5s+f72wXW7LBGNNtGYcVK1aorKzMaUejUeXm5g5Y2Z545oB685eTn/W3O4NSx3mgoVKGBxhqBqRsT05Ojq6//nrXsuuuu04ffnjhhzw7O1uSOh31RCKRTkdHFwUCAWVmZrpeAIChIa4QmjZtmurq6lzLjh07pvHjx0uS8vLylJ2draqqKmd9W1ubqqurVVRUlIDuAgAGk7hOxz3xxBMqKipSRUWFvv71r+vgwYPasmWLtmy5cIolJSVFpaWlqqioUH5+vvLz81VRUaH09HTNnTs3KX8AAMC/4gqhm2++WT/96U+1YsUKrVq1Snl5edq4caPmzZvnbLN06VKdPXtWixYtUlNTk6ZOnap9+/bFfY/QQNWO++ptl7vab/xP0r9yyBmhEVa+l3trgIHT19pxcRcwveuuu3TXXXd1uz4lJUXl5eUqLy+Pd9cAgCGG2nEAAGsIIQCANZ59ntBAKbnvSlf71TjuG2r5q3vbNzb1/R6jwSw15fOabbHzND09J+eS9h1HPbijfzruav/olV+52muXuZ/hAyD5OBICAFhDCAEArCGEAADWeHZOaKBqx8WKndd5bneTq/1qdUu32yJ+/Z0jmnj0Vld7+f7ua9NlXnFlt+sA9M+A1I4DACCRCCEAgDWEEADAGs/OCQ1U7bjelNw7usc2EqvtFfcjPz7+uOcn7X6WddLVzr5q/CV/V/v58672qdPu+b/QWP5bA5eqr7XjOBICAFhDCAEArPHs6TgMHsvXXfojvWNPp2VflbxTYpdnXuFq/+eLL7valPEBko8jIQCANYQQAMAaQggAYI1n54Rsle1B716t+m9X+zfv1va4fTyXTQPwJ8r2AAB8hxACAFhDCAEArPHsnJBXyvYMBf/+w12u9sd/77lUztjQVa72YJnzGT78MttdAHyLsj0AAN8hhAAA1hBCAABrPDsnhP7prV5b6vDP/9P/S+wcz6jB8QiDc+fOudqnT33U4/ZZPLoBGHAcCQEArCGEAADWEEIAAGs8OydE7Ti32EdPxz77JtZguXenN40fHe923a1fmORqf/uh2cnuDjBkUTsOAOA7hBAAwBpCCABgjWfnhKgd5xbq5R6WwTIH1NMcT1fWLnskST0BEA9qxwEAfIcQAgBY49nTcfCvc5995mqfjpzodtvYxyfEnlaMPT3H6TdgcOFICABgDSEEALCGEAIAWOPZOSHK9vTsjulTXO2aP550tRP5qOpPWprd7ejfe9w+fVSGqz1YLh8H0D3K9gAAfIcQAgBYQwgBAKzx7JwQZXt6dsc095zQa/vdj/PubR6m6eOI8771H2d73PaJb97javdWQij20eKZV1zZ4/Y9if07YvfNfUOAN1C2BwDgO4QQAMAaQggAYI1n54TQP709EuHJR+5z3o8ZnZns7gBAlzgSAgBYQwgBAKwhhAAA1nh2TojacfHx0v0yo9JH2u4CgAFG7TgAgO8QQgAAawghAIA1KcZjBdqi0aiCwaBOnTqlzMxMZ1lHWVlZNrqGPnrqP7Y578eGwgndt2l1/9uIrXMHIDkikYirffH3Wrrwmx0KhdTc3Oxa3hWOhAAA1hBCAABrCCEAgDWevU8Ig8e5c58lbd+nTjclbd8Ako8jIQCANYQQAMAaz56Oo2wPAPgHZXsAAL5DCAEArIkrhM6dO6ennnpKeXl5SktL04QJE7Rq1Sq1t7c72xhjVF5ernA4rLS0NM2cOVO1tbUJ7zgAwP/imhNat26dXnjhBe3YsUOTJk3SO++8o2984xsKBoMqKSmRJK1fv14bNmzQ9u3bde2112r16tWaNWuW6urqlJGRccnfZYyRxyoKoY9Wf/ubzvvvbtjuWpeVc3W/9j02y10G6LX9h533d0yf0q99A7h0HX+v4/ntjutI6De/+Y2++tWvavbs2brmmmv0ta99TcXFxXrnnXecL964caNWrlypOXPmqKCgQDt27NCZM2e0c+fOeL4KADAExBVC06dP1+uvv65jx45Jkt577z3t379fd955pySpvr5ejY2NKi4udj4TCAQ0Y8YMHThwoMt9tra2KhqNul4AgKEhrtNxy5YtU3NzsyZOnKjU1FSdP39ea9as0f333y9JamxslCSFQiHX50KhkI4fP97lPisrK/Xss8/2pe8AAJ+LK4R2796tl156STt37tSkSZNUU1Oj0tJShcNhLViwwNku9hpxY0y3142vWLFCZWVlTjsajSo3NzeebsHjhqemOu/b288ndt+XXeZqv/bfzAkBfhJXCD355JNavny57rvvPknSDTfcoOPHj6uyslILFixQdna2pAtHRDk5Oc7nIpFIp6OjiwKBgAKBQF/7DwDwsbjmhM6cOaNhw9wfSU1NdS7RzsvLU3Z2tqqqqpz1bW1tqq6uVlFRUQK6CwAYTOI6EvrKV76iNWvWaNy4cZo0aZJ+//vfa8OGDXr44YclXTgNV1paqoqKCuXn5ys/P18VFRVKT0/X3Llzk/IHAAD8K64Q+v73v6+nn35aixYtUiQSUTgc1qOPPqrvfve7zjZLly7V2bNntWjRIjU1NWnq1Knat29fXPcISdSOAwA/6WvtuBTjsTtCo9GogsGgIpGI82zy5uZm1zZZWVk2uoYEWL5ui6udfdX4hO6/8aPPr8Jcu+yRhO4bwOcikYirHQwGnffRaFRZWVlqbm52fse7Q+04AIA1hBAAwBrPPk+I2nGD09OPPehq/9v3/8vV7u/pucszPj8lEPvo79DY0f3aN4DuDUjtOAAAEokQAgBYQwgBAKzx7JwQBqdR6WlJ3f/lmVc47//zxZdd67hkG/AejoQAANYQQgAAawghAIA1np0TonYcAPhHX2vHcSQEALCGEAIAWOPZ03GU7RkaYi+bfuo/trnaY0PhPu97+PDLet8IQEJQtgcA4DuEEADAGkIIAGCNZ+eEMDSdO/dZwvZl5D4vHftUV8r4APZxJAQAsIYQAgBYQwgBAKzx7JwQZXvQlU9amt3t6N+73fbpx+a72qPSRyajSwBE2R4AgA8RQgAAawghAIA1np0Tonbc0LT62990tZ/6vy+62k988x5XOzR2dNL7BKB31I4DAPgOIQQAsIYQAgBY49k5IQxNw1NTXW3quwGDG0dCAABrCCEAgDWEEADAGs/OCVE7DgD8g9pxAADfIYQAANYQQgAAazw7J0TtOADwD2rHAQB8hxACAFhDCAEArCGEAADWEEIAAGsIIQCANZ69RLtj2Z5PP/3UtS4SidjoEgDgn2J/l4PBoPOesj0AAF8ghAAA1hBCAABrPDsnFAgEFAgEJEnhcNhybwAAHXWcA5Lk/F7Hvu8NR0IAAGsIIQCANZ47HXex+mo0GnWWtba22uoOAOASdDwFd/H3+1KqaXsuhFpaWiRJubm5lnsCAOiPlpaWTnNHsVKMxx7a097erhMnTsgYo3HjxqmhoUGZmZm2u+UL0WhUubm5jFkcGLP4MWbxG2pjZoxRS0uLwuGwhg3redbHc0dCw4YN09VXX+0czmVmZg6J/2iJxJjFjzGLH2MWv6E0Zr0dAV3EhQkAAGsIIQCANZ4NoUAgoGeeeSaum56GOsYsfoxZ/Biz+DFm3fPchQkAgKHDs0dCAIDBjxACAFhDCAEArCGEAADWEEIAAGs8G0KbNm1SXl6eRo4cqSlTpujtt9+23SXPqKys1M0336yMjAxlZWXp7rvvVl1dnWsbY4zKy8sVDoeVlpammTNnqra21lKPvaWyslIpKSkqLS11ljFeXfvoo4/0wAMPaMyYMUpPT9dNN92kw4cPO+sZN7dz587pqaeeUl5entLS0jRhwgStWrVK7e3tzjaMWQzjQbt27TKXXXaZ2bp1qzl69KgpKSkxo0aNMsePH7fdNU/40pe+ZLZt22b+8Ic/mJqaGjN79mwzbtw488knnzjbrF271mRkZJhXXnnFHDlyxNx7770mJyfHRKNRiz237+DBg+aaa64xN954oykpKXGWM16d/e1vfzPjx483Dz30kPnd735n6uvrzWuvvWb+9Kc/Odswbm6rV682Y8aMMT//+c9NfX29efnll83ll19uNm7c6GzDmLl5MoRuueUWs3DhQteyiRMnmuXLl1vqkbdFIhEjyVRXVxtjjGlvbzfZ2dlm7dq1zjb/+Mc/TDAYNC+88IKtblrX0tJi8vPzTVVVlZkxY4YTQoxX15YtW2amT5/e7XrGrbPZs2ebhx9+2LVszpw55oEHHjDGMGZd8dzpuLa2Nh0+fFjFxcWu5cXFxTpw4IClXnlbc3OzJOnKK6+UJNXX16uxsdE1hoFAQDNmzBjSY7h48WLNnj1bd9xxh2s549W1vXv3qrCwUPfcc4+ysrI0efJkbd261VnPuHU2ffp0vf766zp27Jgk6b333tP+/ft15513SmLMuuK5KtqnT5/W+fPnFQqFXMtDoZAaGxst9cq7jDEqKyvT9OnTVVBQIEnOOHU1hsePHx/wPnrBrl279O677+rQoUOd1jFeXfvggw+0efNmlZWV6Tvf+Y4OHjyoxx9/XIFAQPPnz2fcurBs2TI1Nzdr4sSJSk1N1fnz57VmzRrdf//9kvi31hXPhdBFKSkprrYxptMySEuWLNH777+v/fv3d1rHGF7Q0NCgkpIS7du3TyNHjux2O8bLrb29XYWFhaqoqJAkTZ48WbW1tdq8ebPmz5/vbMe4fW737t166aWXtHPnTk2aNEk1NTUqLS1VOBzWggULnO0Ys8957nTc2LFjlZqa2umoJxKJdPq/h6Huscce0969e/XrX/9aV199tbM8OztbkhjDfzp8+LAikYimTJmi4cOHa/jw4aqurtb3vvc9DR8+3BkTxsstJydH119/vWvZddddpw8//FAS/8668uSTT2r58uW67777dMMNN+jBBx/UE088ocrKSkmMWVc8F0IjRozQlClTVFVV5VpeVVWloqIiS73yFmOMlixZoj179uiNN95QXl6ea31eXp6ys7NdY9jW1qbq6uohOYa33367jhw5opqaGudVWFioefPmqaamRhMmTGC8ujBt2rROl/4fO3ZM48ePl8S/s66cOXOm05NEU1NTnUu0GbMuWLwoolsXL9F+8cUXzdGjR01paakZNWqU+ctf/mK7a57wrW99ywSDQfPmm2+akydPOq8zZ84426xdu9YEg0GzZ88ec+TIEXP//fcP6ctAY3W8Os4YxqsrBw8eNMOHDzdr1qwxf/zjH82Pf/xjk56ebl566SVnG8bNbcGCBeaqq65yLtHes2ePGTt2rFm6dKmzDWPm5skQMsaYH/zgB2b8+PFmxIgR5gtf+IJz+TGMkdTla9u2bc427e3t5plnnjHZ2dkmEAiY2267zRw5csRepz0mNoQYr6797Gc/MwUFBSYQCJiJEyeaLVu2uNYzbm7RaNSUlJSYcePGmZEjR5oJEyaYlStXmtbWVmcbxsyN5wkBAKzx3JwQAGDoIIQAANYQQgAAawghAIA1hBAAwBpCCABgDSEEALCGEAIAWEMIAQCsIYQAANYQQgAAa/4/wQOotNEUTlUAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(img0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import imageio \n",
    "imageio.mimsave('pusht_v7_zarr_206.mp4', img,  fps=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}