### Install DM-Reverb with TensorFlow Support

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

Installs the DM-Reverb library with TensorFlow support. Necessary for using Reverb, a component for large-scale data storage and retrieval.

```bash
!pip install dm-reverb[tensorflow]
```

--------------------------------

### Install Dependencies

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Installs the necessary Python packages for the project, including rlds, dm-reverb, flax, and jax.

```python
!pip install rlds dm-reverb[tensorflow]
```

```python
!pip install flax jax==0.4.14
```

--------------------------------

### Install Required Libraries

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_example_for_running_inference_using_RT_1_X_TF_using_tensorflow_datasets.ipynb

Installs the necessary libraries for running RT-1-X inference, including rlds, tf_agents, dm-reverb, apache_beam, and tfp-nightly. Note: tfp-nightly is used due to a specific issue.

```python
# Install required library

# Using tfp-nightly due to https://github.com/tensorflow/probability/issues/1752
!pip install rlds tf_agents dm-reverb[tensorflow] apache_beam tfp-nightly
```

--------------------------------

### Set up Training Functions

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Initializes the Adam optimizer with a specified learning rate and epsilon. This setup is a prerequisite for creating the training state.

```python
# @title Set up the functions for training

optimizer = optax.adam(learning_rate=1e-4, eps=1e-7)

# Create the train state.
# input: batch, rng, ds_info
```

--------------------------------

### Install RLDS with TensorFlow Support

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

Installs the RLDS library with TensorFlow support. Required for loading and processing RLDS datasets.

```bash
!pip install rlds[tensorflow]
```

--------------------------------

### Set up Sharding and Data Parallel Mesh

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Initializes sharding and data parallel mesh configurations. This is a boilerplate setup for distributed training environments.

```python
# @title Set up sharding and data parallel mesh

# Actual global batch size is 1024. Use a smaller batch size for this colab

```

--------------------------------

### Install TensorFlow Datasets and Apache Beam

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

Installs the `tfds-nightly` package for the latest datasets and `apache_beam` for potential distributed dataset generation. Run this before downloading datasets.

```python
!pip install tfds-nightly   # to get most up-to-date registered datasets
!pip install apache_beam
```

--------------------------------

### Download RT-1-X JAX Checkpoint

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/README.md

Use this command to download the JAX checkpoint for RT-1-X models. Ensure gsutil is installed and configured.

```bash
gsutil -m cp -r gs://gdm-robotics-open-x-embodiment/open_x_embodiment_and_rt_x_oss/rt_1_x_jax .
```

--------------------------------

### Create Dataset for Episode Retrieval

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_example_for_running_inference_using_RT_1_X_TF_using_tensorflow_datasets.ipynb

Creates a TensorFlow Dataset object from a specified builder directory to obtain an episode. This example loads the first training episode.

```python
# Create a dataset object to obtain episode from

builder = tfds.builder_from_directory(builder_dir='gs://gresearch/robotics/bridge/0.1.0/')
ds = builder.as_dataset(split='train[:1]')

ds_iterator = iter(ds)
```

--------------------------------

### Download Checkpoint Folder

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_example_for_running_inference_using_RT_1_X_TF_using_tensorflow_datasets.ipynb

Downloads a zipped checkpoint folder from a Google Cloud Storage bucket using gsutil. Ensure gsutil is installed and configured.

```bash
# Download zipped checkpoint folder

!gsutil -m cp -r gs://gdm-robotics-open-x-embodiment/open_x_embodiment_and_rt_x_oss/rt_1_x_tf_trained_for_002272480_step.zip .
```

--------------------------------

### Run the Training Loop

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Executes the training loop for a specified number of steps. It handles data loading, state updates, and metric logging. Ensure the state and RNG are correctly resharded before starting.

```python
# @title Run the train loop

num_train_steps = 1_000  # 1k for example, actual should be > 1M
log_loss_every_steps = 1


# The state should be resharded since we may have loaded pretrained weights
# that need to be converted to jax.Arrays.
state_repl = reshard(state, shardings=replicate_sharding)
# The RNG must be replicated.
rng_repl = reshard(rng, shardings=replicate_sharding)

for step in range(num_train_steps):
  is_last_step = step == num_train_steps

  rng_repl = jax.random.fold_in(rng_repl, step)

  batch = next(train_iter)
  batch = jax.tree_map(_form_gda, batch, global_data_shape)

  state_repl, metrics_update = jitted_train_step(
      state=state_repl, batch=batch, rng=rng_repl
  )

  if step % log_loss_every_steps == 0 or is_last_step:
    metrics_update = jax.device_get(metrics_update)
    print(f"Metrics: step={step}, {metrics_update}")
```

--------------------------------

### Iterate and Verify Datasets with TensorFlow Datasets

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

Use this code to iterate through a list of dataset names, construct their URIs, build dataset objects using tfds.builder_from_directory, and print the size of the first split. Ensure TensorFlow Datasets is installed and the dataset paths are accessible.

```python
for name in DATASETS:
  uri = dataset2path(name)
  b = tfds.builder_from_directory(builder_dir=uri)
  split = list(b.info.splits.keys())[0]
  b.as_dataset(split=split)
  print('Dataset %s has size %s'%(uri, b.info.dataset_size))
```

--------------------------------

### Load and Process RLDS Episode Dataset

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

This code demonstrates loading a subset of a dataset and processing it into individual steps. It resizes observations and concatenates action components. Ensure `tensorflow` and `tensorflow_datasets` are installed.

```python
import tensorflow as tf
import tensorflow_datasets as tfds

# load raw dataset --> replace this with tfds.load(<dataset_name>) on your
# local machine!
dataset = 'kuka'
b = tfds.builder_from_directory(builder_dir=dataset2path(dataset))
ds = b.as_dataset(split='train[:10]')

def episode2steps(episode):
  return episode['steps']

def step_map_fn(step):
  return {
      'observation': {
          'image': tf.image.resize(step['observation']['image'], (128, 128)),
      },
      'action': tf.concat([
          step['action']['world_vector'],
          step['action']['rotation_delta'],
          step['action']['gripper_closedness_action'],
      ], axis=-1)
  }

# convert RLDS episode dataset to individual steps & reformat
ds = ds.map(
    episode2steps, num_parallel_calls=tf.data.AUTOTUNE).flat_map(lambda x: x)
ds = ds.map(step_map_fn, num_parallel_calls=tf.data.AUTOTUNE)

# shuffle, repeat, pre-fetch, batch
ds = ds.cache()         # optionally keep full dataset in memory
ds = ds.shuffle(100)    # set shuffle buffer size
ds = ds.repeat()        # ensure that data never runs out
```

--------------------------------

### RoboTurk: Map Observation

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Configures observation mapping for RoboTurk, setting 'front_rgb' as the source image feature and 'image' as the target. Use this for image input setup in RoboTurk.

```python
roboturk_map_observation = functools.partial(
    map_observation,
    from_image_feature_names=('front_rgb',),
    to_image_feature_names=('image',)
)
```

--------------------------------

### Download Multiple Datasets with TFDS

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

This code snippet downloads a specified list of datasets using `tensorflow_datasets`. It iterates through a list of dataset names and loads each one into a specified directory. Ensure `tqdm` is installed for progress visualization.

```python
DATASET_NAMES = ['fractal20220817_data', 'kuka', 'bridge', 'taco_play', 'jaco_play', 'berkeley_cable_routing', 'roboturk', 'nyu_door_opening_surprising_effectiveness', 'viola', 'berkeley_autolab_ur5', 'toto', 'language_table', 'columbia_cairlab_pusht_real', 'stanford_kuka_multimodal_dataset_converted_externally_to_rlds', 'nyu_rot_dataset_converted_externally_to_rlds', 'stanford_hydra_dataset_converted_externally_to_rlds', 'austin_buds_dataset_converted_externally_to_rlds', 'nyu_franka_play_dataset_converted_externally_to_rlds', 'maniskill_dataset_converted_externally_to_rlds', 'furniture_bench_dataset_converted_externally_to_rlds', 'cmu_franka_exploration_dataset_converted_externally_to_rlds', 'ucsd_kitchen_dataset_converted_externally_to_rlds', 'ucsd_pick_and_place_dataset_converted_externally_to_rlds', 'austin_sailor_dataset_converted_externally_to_rlds', 'austin_sirius_dataset_converted_externally_to_rlds', 'bc_z', 'usc_cloth_sim_converted_externally_to_rlds', 'utokyo_pr2_opening_fridge_converted_externally_to_rlds', 'utokyo_pr2_tabletop_manipulation_converted_externally_to_rlds', 'utokyo_saytap_converted_externally_to_rlds', 'utokyo_xarm_pick_and_place_converted_externally_to_rlds', 'utokyo_xarm_bimanual_converted_externally_to_rlds', 'robo_net', 'berkeley_mvp_converted_externally_to_rlds', 'berkeley_rpt_converted_externally_to_rlds', 'kaist_nonprehensile_converted_externally_to_rlds', 'stanford_mask_vit_converted_externally_to_rlds', 'tokyo_u_lsmo_converted_externally_to_rlds', 'dlr_sara_pour_converted_externally_to_rlds', 'dlr_sara_grid_clamp_converted_externally_to_rlds', 'dlr_edan_shared_control_converted_externally_to_rlds', 'asu_table_top_converted_externally_to_rlds', 'stanford_robocook_converted_externally_to_rlds', 'eth_agent_affordances', 'imperialcollege_sawyer_wrist_cam', 'iamlab_cmu_pickup_insert_converted_externally_to_rlds', 'uiuc_d3field', 'utaustin_mutex', 'berkeley_fanuc_manipulation', 'cmu_food_manipulation', 'cmu_play_fusion', 'cmu_stretch', 'berkeley_gnm_recon', 'berkeley_gnm_cory_hall', 'berkeley_gnm_sac_son']
DOWNLOAD_DIR = '~/tensorflow_datasets'

print(f"Downloading {len(DATASET_NAMES)} datasets to {DOWNLOAD_DIR}.")
for dataset_name in tqdm.tqdm(DATASET_NAMES):
  _ = tfds.load(dataset_name, data_dir=DOWNLOAD_DIR)
```

--------------------------------

### Load and Display Dataset Samples

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

Loads a specified dataset, checks for the display key in observations, and displays the first 10 episodes as a GIF. Requires TensorFlow Datasets and PIL.Image.

```python
dataset = 'fractal20220817_data' # @param ['fractal20220817_data', 'kuka', 'bridge', 'taco_play', 'jaco_play', 'berkeley_cable_routing', 'roboturk', 'nyu_door_opening_surprising_effectiveness', 'viola', 'berkeley_autolab_ur5', 'toto', 'language_table', 'columbia_cairlab_pusht_real', 'stanford_kuka_multimodal_dataset_converted_externally_to_rlds', 'nyu_rot_dataset_converted_externally_to_rlds', 'stanford_hydra_dataset_converted_externally_to_rlds', 'austin_buds_dataset_converted_externally_to_rlds', 'nyu_franka_play_dataset_converted_externally_to_rlds', 'maniskill_dataset_converted_externally_to_rlds', 'furniture_bench_dataset_converted_externally_to_rlds', 'cmu_franka_exploration_dataset_converted_externally_to_rlds', 'ucsd_kitchen_dataset_converted_externally_to_rlds', 'ucsd_pick_and_place_dataset_converted_externally_to_rlds', 'austin_sailor_dataset_converted_externally_to_rlds', 'austin_sirius_dataset_converted_externally_to_rlds', 'bc_z', 'usc_cloth_sim_converted_externally_to_rlds', 'utokyo_pr2_opening_fridge_converted_externally_to_rlds', 'utokyo_pr2_tabletop_manipulation_converted_externally_to_rlds', 'utokyo_saytap_converted_externally_to_rlds', 'utokyo_xarm_pick_and_place_converted_externally_to_rlds', 'utokyo_xarm_bimanual_converted_externally_to_rlds', 'robo_net', 'berkeley_mvp_converted_externally_to_rlds', 'berkeley_rpt_converted_externally_to_rlds', 'kaist_nonprehensile_converted_externally_to_rlds', 'stanford_mask_vit_converted_externally_to_rlds', 'tokyo_u_lsmo_converted_externally_to_rlds', 'dlr_sara_pour_converted_externally_to_rlds', 'dlr_sara_grid_clamp_converted_externally_to_rlds', 'dlr_edan_shared_control_converted_externally_to_rlds', 'asu_table_top_converted_externally_to_rlds', 'stanford_robocook_converted_externally_to_rlds', 'eth_agent_affordances', 'imperialcollege_sawyer_wrist_cam', 'iamlab_cmu_pickup_insert_converted_externally_to_rlds', 'uiuc_d3field', 'utaustin_mutex', 'berkeley_fanuc_manipulation', 'cmu_food_manipulation', 'cmu_play_fusion', 'cmu_stretch', 'berkeley_gnm_recon', 'berkeley_gnm_cory_hall', 'berkeley_gnm_sac_son']
display_key = 'image'

b = tfds.builder_from_directory(builder_dir=dataset2path(dataset))
if display_key not in b.info.features['steps']['observation']:
  raise ValueError(
      f"The key {display_key} was not found in this dataset.\n"
      + "Please choose a different image key to display for this dataset.\n"
      + "Here is the observation spec:\n"
      + str(b.info.features['steps']['observation']))

ds = b.as_dataset(split='train[:10]').shuffle(10)   # take only first 10 episodes
episode = next(iter(ds))
images = [step['observation'][display_key] for step in episode['steps']]
images = [Image.fromarray(image.numpy()) for image in images]
display.Image(as_gif(images))
```

--------------------------------

### Transformation Definitions

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Defines transformations for RLDS datasets. Refer to the dataset colab for usage examples.

```python
# @title Transformation definitions

# For an example usage of the code in this code cell, please take a look at the
# dataset colab at the link below:

```

--------------------------------

### Get Next Trajectory

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

Retrieves the next trajectory from the iterator. This is useful for inspecting individual trajectories or processing them one by one.

```python
trajectory = next(trajectory_iter)
```

--------------------------------

### Prepare and Batch Training Data

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Prepares datasets by shuffling and sampling from multiple sources with specified weights, then batches the data for training. Ensure 'DATASET_NAME_TO_TRAJECTORY_DATASET' is defined and imported.

```python
# @title Batch, and sample one training sample

BATCH_SIZE = 6

# Larger shuffle buffer leads to better performance, but consumes more RAM
datasets = []
weights = []

for name, dataset in DATASET_NAME_TO_TRAJECTORY_DATASET.items():

  datasets.append(dataset.shuffle(10))
  weights.append(float(DATASET_NAME_TO_WEIGHTS[name]))

dataset = tf.data.Dataset.sample_from_datasets(datasets, weights=weights)

# Larger shuffle buffer leads to better performance, but consumes more RAM
dataset = dataset.shuffle(1)

dataset = dataset.batch(BATCH_SIZE)

trajectory_dataset_iter = iter(dataset)

sample = next(trajectory_dataset_iter)
```

--------------------------------

### Get Sample Image Shape

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Retrieves and displays the shape of the image tensor from the sampled batch. Requires the 'sample' object to be previously defined.

```python
sample[rlds.OBSERVATION]['image'].shape
```

--------------------------------

### Initialize Model and Run Forward Pass

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Initializes the model with random weights and performs a forward pass using sample observation and action data. This is useful for verifying model structure and initial behavior.

```python
obs = {
    "image": jnp.ones((1, 15, 300, 300, 3)),
    "natural_language_embedding": jnp.ones((1, 15, 512)),
}
act = {
    "world_vector": jnp.ones((1, 15, 3)),
    "rotation_delta": jnp.ones((1, 15, 3)),
    "gripper_closedness_action": jnp.ones((1, 15, 1)),
    "base_displacement_vertical_rotation": jnp.ones((1, 15, 1)),
    "base_displacement_vector": jnp.ones((1, 15, 2)),
    "terminate_episode": jnp.ones((1, 15, 3), dtype=jnp.int32),
}

variables = rt1x_model.init(
    {
        "params": jax.random.PRNGKey(0),
        "random": jax.random.PRNGKey(0),
    },
    obs,
    act,
    train=False,
)
model_output = rt1x_model.apply(
    variables,
    obs,
    act,
    train=False,
    rngs={"random": jax.random.PRNGKey(0)},
)
```

--------------------------------

### Import Libraries and Define Helper Function

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_example_for_running_inference_using_RT_1_X_TF_using_tensorflow_datasets.ipynb

Imports essential TensorFlow, TF-Agents, and other utility libraries. Includes a helper function `as_gif` to render images as a GIF.

```python
import tensorflow as tf
import tensorflow_datasets as tfds
import rlds
from PIL import Image
import numpy as np
from tf_agents.policies import py_tf_eager_policy
import tf_agents
from tf_agents.trajectories import time_step as ts
from IPython import display
from collections import defaultdict
import matplotlib.pyplot as plt
import tensorflow_hub as hub


def as_gif(images):
  # Render the images as the gif:
  images[0].save('/tmp/temp.gif', save_all=True, append_images=images[1:], duration=1000, loop=0)
  gif_bytes = open('/tmp/temp.gif','rb').read()
  return gif_bytes
```

--------------------------------

### Create and JIT Compile Training State

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Defines and compiles a function to create the initial training state for the agent. This function is JIT compiled for performance.

```python
agent_create_train_state = functools.partial(
    create_train_state, model=rt1x_model, optimizer=optimizer
)
create_train_state_jit = jax.jit(
    agent_create_train_state,
    out_shardings=replicate_sharding,
)

global_data_shape = jax.tree_map(
    lambda x: (global_batch_size,) + x.shape[1:], sample_batch
)

local_devices = mesh.local_devices
local_device_count = jax.local_device_count()

def _put_to_devices(x):
  per_device_arrays = np.split(x, local_device_count, axis=0)
  return jax.device_put(per_device_arrays, local_devices)

def _form_gda(local_data, global_shape):
  arrays = _put_to_devices(local_data)
  return jax.make_array_from_single_device_arrays(
      global_shape, sharding, arrays
  )

rng = jax.random.PRNGKey(0)

sample_batch = jax.tree_map(_form_gda, sample_batch, global_data_shape)
rng, agent_rng = jax.random.split(rng)
state = create_train_state_jit(
    batch=sample_batch, rng=agent_rng
)

# Create the train step.
agent_train = functools.partial(train, model=rt1x_model, optimizer=optimizer)
jitted_train_step = jax.jit(
    agent_train,
    out_shardings=(replicate_sharding, replicate_sharding),
)
```

--------------------------------

### Get next batch from combined dataset

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

Retrieves the next batch of data from the combined trajectory dataset iterator. This batch will contain trajectories from both RoboNet and MT-Opt.

```python
example = next(combined_dataset_it)
```

--------------------------------

### Display Sample Image

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Displays a single image from the sampled batch. Requires the 'sample' object to be previously defined.

```python
Image.fromarray(sample[rlds.OBSERVATION]['image'].numpy()[0][-1])
```

--------------------------------

### Perform Inference with Dummy Input

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_example_for_running_inference_using_RT_1_X_TF_using_tensorflow_datasets.ipynb

Performs a single step of inference using a dummy observation and time step. This initializes the policy state and generates an action based on the dummy input.

```python
# Perform one step of inference using dummy input

# Obtain a dummy observation, where the features are all 0
observation = tf_agents.specs.zero_spec_nest(tf_agents.specs.from_spec(tfa_policy.time_step_spec.observation))

# Construct a tf_agents time_step from the dummy observation
tfa_time_step = ts.transition(observation, reward=np.zeros((), dtype=np.float32))

# Initialize the state of the policy
policy_state = tfa_policy.get_initial_state(batch_size=1)

# Run inference using the policy
action = tfa_policy.action(tfa_time_step, policy_state)
```

--------------------------------

### Create Reverb Pattern Dataset

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Constructs a Reverb PatternDataset from a steps dataset, applying transformations and respecting episode boundaries.

```python
config = create_structured_writer_config('temp', self.pattern)

    # Further transform each step if the `step_map_fn` is provided.
    if self.step_map_fn:
      steps_dataset = steps_dataset.map(self.step_map_fn)
    pattern_dataset = reverb.PatternDataset(
        input_dataset=steps_dataset,
        configs=[config],
        respect_episode_boundaries=True,
        is_end_of_episode=lambda x: x[rlds_types.IS_LAST])
    return pattern_dataset
```

--------------------------------

### Get Trajectory Dataset

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Constructs a trajectory dataset from a directory of episodic data. It handles dataset loading, initial zero-step padding, and trajectory transformation using a specified step mapping function and pattern builder.

```python
def get_trajectory_dataset(builder_dir: str, step_map_fn, trajectory_length: int, split='train[:10]'):
  dataset_builder = tfds.builder_from_directory(builder_dir=builder_dir)
  dataset_builder_episodic_dataset = dataset_builder.as_dataset(split=split)
  dataset_builder_episodic_dataset = dataset_builder_episodic_dataset.map(
      functools.partial(pad_initial_zero_episode, num_zero_step=trajectory_length-1), num_parallel_calls=tf.data.AUTOTUNE)
  rlds_spec = RLDSSpec(
      observation_info=dataset_builder.info.features[rlds.STEPS][rlds.OBSERVATION],
      action_info=dataset_builder.info.features[rlds.STEPS][rlds.ACTION],
  )
  trajectory_transform = TrajectoryTransformBuilder(rlds_spec,
                                                    step_map_fn=step_map_fn,
                                                    pattern_fn=n_step_pattern_builder(trajectory_length)).build(validate_expected_tensor_spec=False)
  trajectory_dataset = trajectory_transform.transform_episodic_rlds_dataset(dataset_builder_episodic_dataset)
  return trajectory_dataset
```

--------------------------------

### Load and Process Second Dataset for Alignment

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

This code loads a second dataset and applies a modified `step_map_fn` to reformat its observations and actions. This is useful for aligning datasets with different specifications before combining them. Ensure `tensorflow` and `tensorflow_datasets` are installed.

```python
# Load second dataset --> replace this with tfds.load(<dataset_name>) on your
# local machine!
dataset = 'utaustin_mutex'
b = tfds.builder_from_directory(builder_dir=dataset2path(dataset))
ds2 = b.as_dataset(split='train[:10]')

def step_map_fn_mutex(step):
  # reformat to align specs of both datasets
  return {
      'observation': {
          'image': tf.image.resize(step['observation']['image'], (128, 128)),
      },
      'action': step['action'],
  }

ds2 = ds2.map(
    episode2steps, num_parallel_calls=tf.data.AUTOTUNE).flat_map(lambda x: x)
ds2 = ds2.map(step_map_fn_mutex, num_parallel_calls=tf.data.AUTOTUNE)

# shuffle, repeat, pre-fetch, batch
ds2 = ds2.cache()         # optionally keep full dataset in memory
ds2 = ds2.shuffle(100)    # set shuffle buffer size
ds2 = ds2.repeat()        # ensure that data never runs out
```

--------------------------------

### Instantiate Trajectory Datasets

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Creates trajectory datasets from the defined configurations using a dictionary comprehension. This processes the keyword arguments to generate usable dataset objects.

```python
DATASET_NAME_TO_TRAJECTORY_DATASET = {k: get_trajectory_dataset(**v) for k, v in DATASET_NAME_TO_TRAJECTORY_DATASET_KWARGS.items()}

```

--------------------------------

### Create Train State Function

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Initializes the training state for the agent, including model parameters and optimizer state. Requires the model, a batch of data, a random number generator, and an optimizer.

```python
def create_train_state(model, batch, rng, optimizer):
    """Creates the train state and initial metrics for agent."""
    obs_input = batch["observation"]
    act_input = batch["action"]

    rng, rng2, rng3 = jax.random.split(rng, 3)
    variables = model.init(
        {"params": rng, "random": rng3},
        obs=obs_input,
        act=act_input,
        train=False,
    )

    params = flax.core.unfreeze(variables["params"])
    batch_stats = flax.core.unfreeze(variables["batch_stats"])

    train_state = TrainState(
        step=0,
        params=flax.core.unfreeze(params),
        opt_state=optimizer.init(params),
        batch_stats=batch_stats,
    )
    return train_state
```

--------------------------------

### Display first trajectory observation from batch

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

Displays the first observation (image) from the current batch, which corresponds to a RoboNet trajectory. Assumes the batch has been retrieved.

```python
# First element of the batch returns a robot_net trajectory
Image.fromarray(example['observation'].numpy()[0][0])
```

--------------------------------

### Organize Action Data for Plotting

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_example_for_running_inference_using_RT_1_X_TF_using_tensorflow_datasets.ipynb

Prepares ground truth and predicted action values into dictionaries keyed by action name and dimension, suitable for plotting. Assumes 'gt_actions' and 'predicted_actions' are populated from the previous step.

```python
action_name_to_values_over_time = defaultdict(list)
predicted_action_name_to_values_over_time = defaultdict(list)
figure_layout = ['terminate_episode_0', 'terminate_episode_1',
        'terminate_episode_2', 'world_vector_0', 'world_vector_1',
        'world_vector_2', 'rotation_delta_0', 'rotation_delta_1',
        'rotation_delta_2', 'gripper_closedness_action_0']
action_order = ['terminate_episode', 'world_vector', 'rotation_delta', 'gripper_closedness_action']

for i, action in enumerate(gt_actions):

  for action_name in action_order:

    for action_sub_dimension in range(action[action_name].shape[0]):

      # print(action_name, action_sub_dimension)
      title = f'{action_name}_{action_sub_dimension}'

      action_name_to_values_over_time[title].append(action[action_name][action_sub_dimension])
      predicted_action_name_to_values_over_time[title].append(predicted_actions[i][action_name][action_sub_dimension])
```

--------------------------------

### Create Zero Episode Dataset

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

Creates a dataset of zero-valued episodes based on a given RLDS Spec. It handles the creation of steps within episodes and removes a 'fake' key if present.

```python
def zero_episode_dataset_from_spec(rlds_spec: RLDS_SPEC):
  """Creates a zero valued dataset of episodes for the given RLDS Spec."""

  def add_steps(episode, step_spec):
    episode[rlds_types.STEPS] = transformations.zero_dataset_like(
        tf.data.DatasetSpec(step_spec))
    if 'fake' in episode:
      del episode['fake']
    return episode

  episode_without_steps_spec = {
      k: v
      for k, v in rlds_spec.episode_tensor_spec().items()
      if k != rlds_types.STEPS
  }

  if episode_without_steps_spec:
    episodes_dataset = transformations.zero_dataset_like(
        tf.data.DatasetSpec(episode_without_steps_spec))
  else:
    episodes_dataset = tf.data.Dataset.from_tensors({'fake': ''})

  episodes_dataset_with_steps = episodes_dataset.map(
      lambda episode: add_steps(episode, rlds_spec.step_tensor_spec()))
  return episodes_dataset_with_steps
```

--------------------------------

### Prepare and Batch Dataset

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Prepares a TensorFlow dataset for model input and batches it for training. Ensures the dataset is prefetched for performance.

```python
train_dataset = tf.data.Dataset.sample_from_datasets(datasets, weights=weights)

train_dataset = prepare_for_model_input(
    train_dataset, target_height=300, target_width=300, training=True
)

# Creating mesh and shardings.
num_devices = len(jax.devices())
mesh = jax.sharding.Mesh(
    mesh_utils.create_device_mesh((num_devices,)), ("data",)
)

# Data parallel mesh.
sharding = jax.sharding.NamedSharding(mesh, P("data"))
replicate_sharding = NamedSharding(mesh, P())

global_batch_size = jax.device_count() * PER_DEVICE_BATCH_SIZE
local_batch_size = jax.local_device_count() * PER_DEVICE_BATCH_SIZE
train_dataset = train_dataset.batch(local_batch_size, drop_remainder=True)

train_dataset = train_dataset.prefetch(tf.data.AUTOTUNE)

train_iter = train_dataset.as_numpy_iterator()
sample_batch = jax.tree_map(lambda x: x, next(train_iter))

print(f"Local batch size: {local_batch_size}")
print(f"Global batch size: {global_batch_size}")
print(f"Devices: {jax.devices()}")
print(f"Sample batch keys: {sample_batch.keys()}")
```

--------------------------------

### Pad Initial Zero Steps for Dataset

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Pads the beginning of a tf.data.Dataset with a specified number of zero steps. This is useful for ensuring trajectories have a minimum length or for handling edge cases in trajectory construction.

```python
def pad_initial_zero_steps(
    steps: tf.data.Dataset, num_zero_step: int
) -> tf.data.Dataset:
  zero_steps = steps.take(1)
  zero_steps = zero_steps.map(lambda x: tf.nest.map_structure(tf.zeros_like, x), num_parallel_calls=tf.data.AUTOTUNE)
  zero_steps = zero_steps.repeat(num_zero_step)
  return rlds.transformations.concatenate(zero_steps, steps)
```

--------------------------------

### Create Pattern Dataset with Step Mapping

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

Internal method to create a reverb.PatternDataset. It applies an optional `step_map_fn` to each step before creating the pattern dataset.

```python
def _create_pattern_dataset(
      self, steps_dataset: tf.data.Dataset) -> tf.data.Dataset:
    """Create PatternDataset from the `steps_dataset`."""
    config = create_structured_writer_config('temp', self.pattern)

    # Further transform each step if the `step_map_fn` is provided.
    if self.step_map_fn:
      steps_dataset = steps_dataset.map(self.step_map_fn)
    pattern_dataset = reverb.PatternDataset(

```

--------------------------------

### Convert Steps to List

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_example_for_running_inference_using_RT_1_X_TF_using_tensorflow_datasets.ipynb

Converts the steps from a dataset episode into a Python list for easier manipulation.

```python
steps = list(steps)
```

--------------------------------

### RoboTurk: Map Action

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Maps RoboTurk actions to the model's expected format, including 'world_vector', 'gripper_closedness_action', 'rotation_delta', and 'terminate_episode'. Use this for RoboTurk task integration.

```python
def roboturk_map_action(to_step: rlds.Step, from_step: rlds.Step):
  to_step[rlds.ACTION]['world_vector'] = from_step[rlds.ACTION]['world_vector']
  to_step[rlds.ACTION]['gripper_closedness_action'] = from_step[rlds.ACTION][
      'gripper_closedness_action'
  ]
  to_step[rlds.ACTION]['rotation_delta'] = from_step[rlds.ACTION]['rotation_delta']
  to_step[rlds.ACTION]['terminate_episode'] = terminate_bool_to_act(
      from_step[rlds.ACTION]['terminate_episode']
  )
```

--------------------------------

### Create Trajectories of Length 3

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

This snippet demonstrates the structure of a trajectory data sample, showing tensors for actions, gripper states, and terminal/last flags. It's useful for understanding the format of data used in robotic control experiments.

```python
{
  'actions': <tf.Tensor: shape=(3, 7), dtype=float32, numpy=array([[ 0.65729046, -0.21117291,  0.18413243,  0.70710677,  0.70710677,
          0.        ,  0.        ],
        [ 0.65729046, -0.21117291,  0.18413243,  0.70710677,  0.70710677,
          0.        ,  0.        ]], dtype=float32)>,
  'close_gripper': <tf.Tensor: shape=(3,), dtype=bool, numpy=array([False, False, False])>},
 'is_terminal': <tf.Tensor: shape=(3,), dtype=bool, numpy=array([False, False, False])>,
 'is_last': <tf.Tensor: shape=(3,), dtype=bool, numpy=array([False, False, False])>
}
```

--------------------------------

### NYU VINN: Map Action

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Maps NYU VINN actions to the model's expected format by scaling 'world_vector' and 'rotation_delta' and mapping 'gripper_closedness_action' and 'terminate_episode'. Use this for NYU VINN task integration.

```python
def nyu_door_opening_surprising_effectiveness_map_action(to_step: rlds.Step, from_step: rlds.Step):
  """Maps dataset action to action expected by the model."""

  # The world vector as existed in the dataset on disk ranges from -0.07 to 0.07
  # We scale by 20.0 so that the action spans the limit of the world_vector
  # action, from -2.0 to 2.0.
  to_step['action']['world_vector'] = from_step['action']['world_vector'] * 20.0

  # Similarly, the rotation_delta in the dataset on disk ranges from -0.07 to
  # 0.07.
  # We scale by 15.0 so that the rotation_delta almost spans the limit of
  # rotation_delta, from -pi/2 to pi/2.
  to_step['action']['rotation_delta'] = (
      from_step['action']['rotation_delta'] * 15.0
  )

  to_step['action']['gripper_closedness_action'] = (
      from_step['action']['gripper_closedness_action']
  )

  to_step['action']['terminate_episode'] = terminate_bool_to_act(
      from_step['action']['terminate_episode']
  )
```

--------------------------------

### Model Configuration Class (EfficientNet-B0 Default)

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Defines configuration parameters for the entire model, including block structure and training settings. This class provides a default configuration for EfficientNet-B0.

```python
class ModelConfig(object):
  """Class that contains configuration parameters for the model."""

  def __init__(
      self,
      width_coefficient: float = 1.0,
      depth_coefficient: float = 1.0,
      resolution: int = 224,
      dropout_rate: float = 0.2,
      blocks: Tuple[BlockConfig, ...] = (
          # (input_filters, output_filters, kernel_size, num_repeat,
          #  expand_ratio, strides, se_ratio)
          # pylint: disable=bad-whitespace
          BlockConfig(32, 16, 3, 1, 1, (1, 1), 0.25),
          BlockConfig(16, 24, 3, 2, 6, (2, 2), 0.25),
          BlockConfig(24, 40, 5, 2, 6, (2, 2), 0.25),
          BlockConfig(40, 80, 3, 3, 6, (2, 2), 0.25),
          BlockConfig(80, 112, 5, 3, 6, (1, 1), 0.25),
          BlockConfig(112, 192, 5, 4, 6, (2, 2), 0.25),
          BlockConfig(192, 320, 3, 1, 6, (1, 1), 0.25),
          # pylint: enable=bad-whitespace
      ),
      stem_base_filters: int = 32,
      top_base_filters: int = 1280,
      activation: str = 'swish',
      batch_norm: str = 'default',
      bn_momentum: float = 0.99,
      bn_epsilon: float = 1e-3,
      # While the original implementation used a weight decay of 1e-5,
      # tf.nn.l2_loss divides it by 2, so we halve this to compensate in Keras
      weight_decay: float = 5e-6,
      drop_connect_rate: float = 0.2,
      depth_divisor: int = 8,
      min_depth: Optional[int] = None,
      use_se: bool = True,
      input_channels: int = 3,
      num_classes: int = 1000,
      model_name: str = 'efficientnet',
      rescale_input: bool = True,
      data_format: str = 'channels_last',
      final_projection_size: int = 0,
      classifier_head: bool = True,
      dtype: jnp.dtype = jnp.float32,
  ):
    """Default Config for Efficientnet-B0."""
    for arg in locals().items():
      setattr(self, *arg)
```

--------------------------------

### Create and combine trajectory datasets

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

Transforms episodic datasets from RoboNet and MT-Opt into trajectory datasets and then combines them using tf.data.Dataset.sample_from_datasets. Batches the combined dataset for further processing.

```python
# Create trajectory datasets for the two normalized representations:
rrobo_net_trajectory_dataset = robo_net_trajectory_transform.transform_episodic_rlds_dataset(robo_net_builder_episodic_dataset)
mt_opt_trajectory_dataset = mt_opt_trajectory_transform.transform_episodic_rlds_dataset(mt_opt_episodic_dataset)


combined_dataset = tf.data.Dataset.sample_from_datasets([robo_net_trajectory_dataset, mt_opt_trajectory_dataset])
combined_dataset = combined_dataset.batch(2)
combined_dataset_it = iter(combined_dataset)
```

--------------------------------

### Model Configuration

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Defines model configurations, including width, depth, resolution, and dropout. Used for EfficientNet-B3.

```python
EN_MODEL_CONFIGS = {
    # (width, depth, resolution, dropout)
    'efficientnet-b3': ModelConfig(1.2, 1.4, 300, 0.3),
}
```

--------------------------------

### Instantiate RT-1 Model

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Initializes the RT-1 model with specified parameters for image tokens, action tokens, layer size, and vocabulary size. This configuration is used for training or inference.

```python
SEQUENCE_LENGTH = 15
NUM_ACTION_TOKENS = 11
LAYER_SIZE = 256
VOCAB_SIZE = 512
NUM_IMAGE_TOKENS = 81

rt1x_model = RT1(
    num_image_tokens=NUM_IMAGE_TOKENS,
    num_action_tokens=NUM_ACTION_TOKENS,
    layer_size=LAYER_SIZE,
    vocab_size=VOCAB_SIZE,
    # Use token learner to reduce tokens per image to 81.
    use_token_learner=True,
    # RT-1-X uses (-2.0, 2.0) instead of (-1.0, 1.0).
    world_vector_range=(-2.0, 2.0)
)

```

--------------------------------

### Extract Steps from Dataset Episode

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_example_for_running_inference_using_RT_1_X_TF_using_tensorflow_datasets.ipynb

Obtain the steps from a single episode from the dataset. This is a prerequisite for further processing.

```python
episode = next(ds_iterator)
steps = episode[rlds.STEPS]
```

--------------------------------

### Iterate Through Episode Steps

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

This snippet shows how to iterate through the steps of an episode and print each element. It's useful for inspecting the structure of episode data.

```python
for elem in next(iter(episode['steps'])).items():
  print(elem)
```

--------------------------------

### Create Structured Writer Config

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb

Creates a Reverb structured writer configuration for a given table name and pattern. It initializes the config with an empty list of conditions.

```python
def create_structured_writer_config(table_name: str,
                                    pattern: reverb.structured_writer.Pattern) -> Any:
  config = reverb.structured_writer.create_config(
      pattern=pattern, table=table_name, conditions=[])
  return config
```

--------------------------------

### Define Trajectory Dataset Configurations

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Defines keyword arguments for creating trajectory datasets, including builder directories, trajectory lengths, and mapping functions for observations and actions. This configuration is used to instantiate different robotics datasets.

```python
DATASET_NAME_TO_TRAJECTORY_DATASET_KWARGS = {
    # Roboturk
    'roboturk_map_action': {
        'builder_dir': 'gs://gresearch/robotics/roboturk_map_action/0.1.0',
        'trajectory_length': 15,
        'step_map_fn':functools.partial(step_map_fn,
                                        map_observation=map_observation,
                                        map_action=roboturk_map_action)
    },
    # NYU VINN
    'nyu_door_opening_surprising_effectiveness': {
        'builder_dir': 'gs://gresearch/robotics/nyu_door_opening_surprising_effectiveness/0.1.0',
        'trajectory_length': 15,
        'step_map_fn':functools.partial(step_map_fn,
                                        map_observation=map_observation,
                                        map_action=nyu_door_opening_surprising_effectiveness_map_action)
    },
    # Austin VIOLA
    'viola': {
        'builder_dir': 'gs://gresearch/robotics/viola/0.1.0',
        'trajectory_length': 15,
        'step_map_fn':functools.partial(step_map_fn,
                                        map_observation=viola_map_observation,
                                        map_action=viola_map_action)
    },
    # Berkeley Autolab UR5
    'berkeley_autolab_ur5': {
        'builder_dir': 'gs://gresearch/robotics/berkeley_autolab_ur5/0.1.0',
        'trajectory_length': 15,
        'step_map_fn':functools.partial(step_map_fn,
                                        map_observation=map_observation,
                                        map_action=berkeley_autolab_ur5_map_action)
    },
    # TODO: (add Language Table)
    'toto': {
        'builder_dir': 'gs://gresearch/robotics/toto/0.1.0',
        'trajectory_length': 15,
        'step_map_fn':functools.partial(step_map_fn,
                                        map_observation=map_observation,
                                        map_action=toto_map_action)
    }
}

```

--------------------------------

### Training Step Function

Source: https://github.com/google-deepmind/open_x_embodiment/blob/main/colabs/Minimal_Training_Example.ipynb

Performs a single training step using the provided batch, state, model, optimizer, and random number generator. Returns the updated state and metrics.

```python
def train(batch, state, model, optimizer, rng):
  """Performs a single training step."""
  rng, loss_rng = jax.random.split(rng)

  def loss_fn(params):
    variables = {"params": params, "batch_stats": state.batch_stats}
    per_example_loss, new_variables = rt1_loss(
        model, batch=batch, variables=variables, rng=loss_rng
    )
    loss = jnp.mean(per_example_loss)
    return loss, new_variables["batch_stats"]

  grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
  (loss, new_batch_stats), grad = grad_fn(state.params)

  loss = jnp.mean(loss)

  updates, new_opt_state = optimizer.update(
      grad, state.opt_state, state.params
  )

  new_params = optax.apply_updates(state.params, updates)
  new_state = state.replace(
      step=state.step + 1,
      params=flax.core.unfreeze(new_params),
      opt_state=flax.core.unfreeze(new_opt_state),
      batch_stats=flax.core.unfreeze(new_batch_stats),
  )

  metrics_update = {
      "loss": loss,
  }
  return new_state, metrics_update
```