<jupyter_start><jupyter_text>A deep dive into robomimic datasetsThis notebook will provide examples on how to work with robomimic datasets through various python code examples. This notebook assumes that you have installed `robomimic` and `robosuite` (which should be on the `offline_study` branch).  Download datasetFirst, let's try downloading a simple dataset - we'll use the Lift (PH) dataset. Note that there are utility scripts such as `scripts/download_datasets.py` to do this for us, but for the purposes of this example, we'll use the python API.<jupyter_code>import os
import json
import h5py
import numpy as np

import robomimic
import robomimic.utils.file_utils as FileUtils

# the dataset registry can be found at robomimic/__init__.py
from robomimic import DATASET_REGISTRY

# set download folder and make it
download_folder = "/tmp/robomimic_ds_example"
os.makedirs(download_folder, exist_ok=True)

# download the dataset
task = "lift"
dataset_type = "ph"
hdf5_type = "low_dim"
FileUtils.download_url(
    url=DATASET_REGISTRY[task][dataset_type][hdf5_type]["url"], 
    download_dir=download_folder,
)

# enforce that the dataset exists
dataset_path = os.path.join(download_folder, "low_dim.hdf5")
assert os.path.exists(dataset_path)<jupyter_output>low_dim.hdf5: 18.6MB [00:00, 27.5MB/s]<jupyter_text>Read quantities from datasetNext, let's demonstrate how to read different quantities from the dataset. There are scripts such as `scripts/get_dataset_info.py` that can help you easily understand the contents of a dataset, but in this example, we'll break down how to do this directly.First, let's take a look at the number of demonstrations in the file.<jupyter_code># open file
f = h5py.File(dataset_path, "r")

# each demonstration is a group under "data"
demos = list(f["data"].keys())
num_demos = len(demos)

print("hdf5 file {} has {} demonstrations".format(dataset_path, num_demos))<jupyter_output>hdf5 file /tmp/robomimic_ds_example/low_dim.hdf5 has 200 demonstrations<jupyter_text>Next, let's list all of the demonstrations, along with the number of state-action pairs in each demonstration.<jupyter_code># each demonstration is named "demo_#" where # is a number.
# Let's put the demonstration list in increasing episode order
inds = np.argsort([int(elem[5:]) for elem in demos])
demos = [demos[i] for i in inds]

for ep in demos:
    num_actions = f["data/{}/actions".format(ep)].shape[0]
    print("{} has {} samples".format(ep, num_actions))<jupyter_output>demo_0 has 59 samples
demo_1 has 58 samples
demo_2 has 57 samples
demo_3 has 55 samples
demo_4 has 51 samples
demo_5 has 58 samples
demo_6 has 49 samples
demo_7 has 49 samples
demo_8 has 44 samples
demo_9 has 51 samples
demo_10 has 54 samples
demo_11 has 49 samples
demo_12 has 53 samples
demo_13 has 59 samples
demo_14 has 51 samples
demo_15 has 50 samples
demo_16 has 50 samples
demo_17 has 45 samples
demo_18 has 49 samples
demo_19 has 51 samples
demo_20 has 48 samples
demo_21 has 57 samples
demo_22 has 47 samples
demo_23 has 47 samples
demo_24 has 52 samples
demo_25 has 48 samples
demo_26 has 43 samples
demo_27 has 46 samples
demo_28 has 49 samples
demo_29 has 41 samples
demo_30 has 46 samples
demo_31 has 42 samples
demo_32 has 59 samples
demo_33 has 54 samples
demo_34 has 48 samples
demo_35 has 58 samples
demo_36 has 45 samples
demo_37 has 50 samples
demo_38 has 49 samples
demo_39 has 41 samples
demo_40 has 40 samples
demo_41 has 49 samples
demo_42 has 53 samples
demo_43 has 39 sample[...]<jupyter_text>Now, let's dig into a single trajectory to take a look at some of the quantities in each demonstration.<jupyter_code># look at first demonstration
demo_key = demos[0]
demo_grp = f["data/{}".format(demo_key)]

# Each observation is a dictionary that maps modalities to numpy arrays, and
# each action is a numpy array. Let's print the observations and actions for the 
# first 5 timesteps of this trajectory.
for t in range(5):
    print("timestep {}".format(t))
    obs_t = dict()
    # each observation modality is stored as a subgroup
    for k in demo_grp["obs"]:
        obs_t[k] = demo_grp["obs/{}".format(k)][t] # numpy array
    act_t = demo_grp["actions"][t]
    
    # pretty-print observation and action using json
    obs_t_pp = { k : obs_t[k].tolist() for k in obs_t }
    print("obs")
    print(json.dumps(obs_t_pp, indent=4))
    print("action")
    print(act_t)
# we can also grab multiple timesteps at once directly, or even the full trajectory at once
first_ten_actions = demo_grp["actions"][:10]
print("shape of first ten actions {}".format(first_ten_actions.shape))
all_actions = demo_grp["actions"][:]
print("shape of all actions {}".format(all_actions.shape))
# the trajectory also contains the next observations under "next_obs", 
# for convenient use in a batch (offline) RL pipeline. Let's verify
# that "next_obs" and "obs" are offset by 1.
for k in demo_grp["obs"]:
    # obs_{t+1} == next_obs_{t}
    assert(np.allclose(demo_grp["obs"][k][1:], demo_grp["next_obs"][k][:-1]))
print("success")
# we also have "done" and "reward" information stored in each trajectory.
# In this case, we have sparse rewards that indicate task completion at
# that timestep.
dones = demo_grp["dones"][:]
rewards = demo_grp["rewards"][:]
print("dones")
print(dones)
print("")
print("rewards")
print(rewards)
# each demonstration also contains metadata
num_samples = demo_grp.attrs["num_samples"] # number of samples in this trajectory
mujoco_xml_file = demo_grp.attrs["model_file"] # mujoco XML file for this demonstration
print(mujoco_xml_file)<jupyter_output><mujoco model="base">
    <compiler angle="radian" meshdir="meshes/" />
    <option impratio="20" density="1.2" viscosity="2e-05" cone="elliptic" />
    <size njmax="5000" nconmax="5000" />
    <visual>
        <map znear="0.001" />
    </visual>
    <asset>
        <texture type="skybox" builtin="gradient" rgb1="0.9 0.9 1" rgb2="0.2 0.3 0.4" width="256" height="1536" />
        <texture type="2d" name="texplane" file="/home/robot/installed_libraries/robosuite/robosuite/models/assets/arenas/../textures/wood-tiles.png" />
        <texture type="cube" name="tex-ceramic" file="/home/robot/installed_libraries/robosuite/robosuite/models/assets/arenas/../textures/ceramic.png" />
        <texture type="cube" name="tex-steel-brushed" file="/home/robot/installed_libraries/robosuite/robosuite/models/assets/arenas/../textures/steel-brushed.png" />
        <texture type="2d" name="tex-cream-plaster" file="/home/robot/installed_libraries/robosuite/robosuite/models/assets/arenas/../textures/cream-pl[...]<jupyter_text>Finally, let's take a look at some global metadata present in the file. The hdf5 file stores environment metadata which is a convenient way to understand which simulation environment (task) the dataset was collected on.<jupyter_code>env_meta = json.loads(f["data"].attrs["env_args"])
# note: we could also have used the following function:
# env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=dataset_path)
print("==== Env Meta ====")
print(json.dumps(env_meta, indent=4))
print("")<jupyter_output>==== Env Meta ====
{
    "env_name": "Lift",
    "type": 1,
    "env_kwargs": {
        "has_renderer": false,
        "has_offscreen_renderer": false,
        "ignore_done": true,
        "use_object_obs": true,
        "use_camera_obs": false,
        "control_freq": 20,
        "controller_configs": {
            "type": "OSC_POSE",
            "input_max": 1,
            "input_min": -1,
            "output_max": [
                0.05,
                0.05,
                0.05,
                0.5,
                0.5,
                0.5
            ],
            "output_min": [
                -0.05,
                -0.05,
                -0.05,
                -0.5,
                -0.5,
                -0.5
            ],
            "kp": 150,
            "damping": 1,
            "impedance_mode": "fixed",
            "kp_limits": [
                0,
                300
            ],
            "damping_limits": [
                0,
                10
            ],
   [...]<jupyter_text>Visualizing demonstration trajectoriesFinally, let's play some of these demonstrations back in the simulation environment to easily visualize the data that was collected. It turns out that the environment metadata stored in the hdf5 allows us to easily create a simulation environment that is consistent with the way the dataset was collected!<jupyter_code>import robomimic.utils.env_utils as EnvUtils

# create simulation environment from environment metedata
env = EnvUtils.create_env_from_metadata(
    env_meta=env_meta, 
    render=False,            # no on-screen rendering
    render_offscreen=True,   # off-screen rendering to support rendering video frames
)
import robomimic.utils.obs_utils as ObsUtils

# We normally need to make sure robomimic knows which observations are images (for the
# data processing pipeline). This is usually inferred from your training config, but
# since we are just playing back demonstrations, we just need to initialize robomimic
# with a dummy spec.
dummy_spec = dict(
    obs=dict(
            low_dim=["robot0_eef_pos"],
            rgb=[],
        ),
)
ObsUtils.initialize_obs_utils_with_obs_specs(obs_modality_specs=dummy_spec)
import imageio

# prepare to write playback trajectories to video
video_path = os.path.join(download_folder, "playback.mp4")
video_writer = imageio.get_writer(video_path, fps=20)
def playback_trajectory(demo_key):
    """
    Simple helper function to playback the trajectory stored under the hdf5 group @demo_key and
    write frames rendered from the simulation to the active @video_writer.
    """
    
    # robosuite datasets store the ground-truth simulator states under the "states" key.
    # We will use the first one, alone with the model xml, to reset the environment to
    # the initial configuration before playing back actions.
    init_state = f["data/{}/states".format(demo_key)][0]
    model_xml = f["data/{}".format(demo_key)].attrs["model_file"]
    initial_state_dict = dict(states=init_state, model=model_xml)
    
    # reset to initial state
    env.reset_to(initial_state_dict)
    
    # playback actions one by one, and render frames
    actions = f["data/{}/actions".format(demo_key)][:]
    for t in range(actions.shape[0]):
        env.step(actions[t])
        video_img = env.render(mode="rgb_array", height=512, width=512, camera_name="agentview")
        video_writer.append_data(video_img)
# playback the first 5 demos
for ep in demos[:5]:
    print("Playing back demo key: {}".format(ep))
    playback_trajectory(ep)

# done writing video
video_writer.close()
# view the trajectories!
from IPython.display import Video
Video(video_path, embed=True)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10065686002405402, 'fraction_numerical': 0.02821722638541956, 'mean_word_length': 3.795918367346939, 'pattern_counts': {'":': 20, '<': 33, '<?xml version=': 0, '>': 32, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 5}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '581476', 'license': 'MIT', 'n_tokens_mistral': 3597, 'n_tokens_neox': 3350, 'n_words': 1090, 'path': 'examples/notebooks/datasets.ipynb', 'repo_name': 'ellislm/robomimic', 'size': '10809'}