test_dlimp.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/yagmurlu/micromamba/envs/octo/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "2024-06-05 14:13:42.656938: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2024-06-05 14:13:42.657071: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2024-06-05 14:13:42.693888: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2024-06-05 14:13:44.699990: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import jax\n",
    "import tensorflow_datasets as tfds\n",
    "import tensorflow as tf\n",
    "import tqdm\n",
    "import mediapy\n",
    "import numpy as np\n",
    "import dlimp as dl\n",
    "from functools import partial\n",
    "from fnmatch import fnmatch\n",
    "from octo.utils.spec import ModuleSpec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_by_language_key(traj, *, language_key_template):\n",
    "    match_keys = [key for key in traj.keys() if fnmatch(key, language_key_template)]\n",
    "    if len(match_keys) == 0:\n",
    "        raise ValueError(f\"No matching key found for {language_key_template}. Keys: {traj.keys()}\")\n",
    "    \n",
    "    labels = tf.stack([traj[key] for key in match_keys], axis=0)\n",
    "    # if _any_ label in _any_ step is not empty, return True\n",
    "    return tf.math.reduce_any(labels != \"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-06-05 14:13:59.203440: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2256] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.\n",
      "Skipping registering GPU devices...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'episode_metadata': {'file_path': <tf.Tensor: shape=(), dtype=string, numpy=b'/home/marcelr/BridgeData/raw/bridge_data_v1/berkeley/toykitchen1/flip_pot_upright_in_sink_distractors/2021-06-02_17-24-16/raw/traj_group0/traj2'>, 'has_depth_0': <tf.Tensor: shape=(), dtype=bool, numpy=False>, 'has_image_0': <tf.Tensor: shape=(), dtype=bool, numpy=True>, 'has_image_1': <tf.Tensor: shape=(), dtype=bool, numpy=True>, 'has_image_2': <tf.Tensor: shape=(), dtype=bool, numpy=True>, 'has_image_3': <tf.Tensor: shape=(), dtype=bool, numpy=False>, 'has_language': <tf.Tensor: shape=(), dtype=bool, numpy=False>, 'traj_length': <tf.Tensor: shape=(), dtype=float64, numpy=20.0>}, 'steps': <_VariantDataset element_spec={'action': TensorSpec(shape=(7,), dtype=tf.float64, name=None), 'delta_robot_transform': TensorSpec(shape=(4, 4), dtype=tf.float64, name=None), 'discount': TensorSpec(shape=(), dtype=tf.float64, name=None), 'groundtruth': TensorSpec(shape=(), dtype=tf.string, name=None), 'groundtruth_1': TensorSpec(shape=(), dtype=tf.string, name=None), 'groundtruth_10': TensorSpec(shape=(), dtype=tf.string, name=None), 'groundtruth_11': TensorSpec(shape=(), dtype=tf.string, name=None), 'groundtruth_2': TensorSpec(shape=(), dtype=tf.string, name=None), 'groundtruth_3': TensorSpec(shape=(), dtype=tf.string, name=None), 'groundtruth_4': TensorSpec(shape=(), dtype=tf.string, name=None), 'groundtruth_5': TensorSpec(shape=(), dtype=tf.string, name=None), 'groundtruth_6': TensorSpec(shape=(), dtype=tf.string, name=None), 'groundtruth_7': TensorSpec(shape=(), dtype=tf.string, name=None), 'groundtruth_8': TensorSpec(shape=(), dtype=tf.string, name=None), 'groundtruth_9': TensorSpec(shape=(), dtype=tf.string, name=None), 'is_first': TensorSpec(shape=(), dtype=tf.bool, name=None), 'is_last': TensorSpec(shape=(), dtype=tf.bool, name=None), 'is_terminal': TensorSpec(shape=(), dtype=tf.bool, name=None), 'language_embedding': TensorSpec(shape=(1, 512), dtype=tf.float32, name=None), 'language_instruction': TensorSpec(shape=(), dtype=tf.string, name=None), 'language_instruction_1': TensorSpec(shape=(), dtype=tf.string, name=None), 'language_instruction_10': TensorSpec(shape=(), dtype=tf.string, name=None), 'language_instruction_11': TensorSpec(shape=(), dtype=tf.string, name=None), 'language_instruction_2': TensorSpec(shape=(), dtype=tf.string, name=None), 'language_instruction_3': TensorSpec(shape=(), dtype=tf.string, name=None), 'language_instruction_4': TensorSpec(shape=(), dtype=tf.string, name=None), 'language_instruction_5': TensorSpec(shape=(), dtype=tf.string, name=None), 'language_instruction_6': TensorSpec(shape=(), dtype=tf.string, name=None), 'language_instruction_7': TensorSpec(shape=(), dtype=tf.string, name=None), 'language_instruction_8': TensorSpec(shape=(), dtype=tf.string, name=None), 'language_instruction_9': TensorSpec(shape=(), dtype=tf.string, name=None), 'new_robot_transform': TensorSpec(shape=(4, 4), dtype=tf.float64, name=None), 'observation': {'depth_0': TensorSpec(shape=(480, 640, 3), dtype=tf.uint8, name=None), 'desired_state': TensorSpec(shape=(7,), dtype=tf.float64, name=None), 'full_state': TensorSpec(shape=(7,), dtype=tf.float64, name=None), 'image_0': TensorSpec(shape=(480, 640, 3), dtype=tf.uint8, name=None), 'image_1': TensorSpec(shape=(480, 640, 3), dtype=tf.uint8, name=None), 'image_2': TensorSpec(shape=(480, 640, 3), dtype=tf.uint8, name=None), 'image_3': TensorSpec(shape=(480, 640, 3), dtype=tf.uint8, name=None), 'state': TensorSpec(shape=(7,), dtype=tf.float64, name=None)}, 'reward': TensorSpec(shape=(), dtype=tf.float64, name=None)}>}\n"
     ]
    }
   ],
   "source": [
    "builder = tfds.builder_from_directory(builder_dir='/home/marcelr/tensorflow_datasets/bridge/1.0.0/')\n",
    "\n",
    "ds = builder.as_dataset(split='train[:50]')\n",
    "# ds = dl.DLataset.from_rlds(builder, split=\"train[:50]\", shuffle=False)\n",
    "\n",
    "# filter_functions = (ModuleSpec.create(\n",
    "#     partial(\n",
    "#         filter_by_language_key,\n",
    "#         language_key_template=\"groundtruth*\"\n",
    "#     )\n",
    "# ), ModuleSpec.create(\n",
    "#     partial(\n",
    "#         filter_by_language_key,\n",
    "#         language_key_template=\"language_instruction*\"\n",
    "#     )\n",
    "# ))\n",
    "\n",
    "# for filter_fcn_spec in filter_functions:\n",
    "#     ds = ds.filter(ModuleSpec.instantiate(filter_fcn_spec))\n",
    "\n",
    "# filter_functions = ((\n",
    "#     partial(\n",
    "#         filter_by_language_key,\n",
    "#         language_key_template=\"groundtruth*\"\n",
    "#     )\n",
    "# ), (\n",
    "#     partial(\n",
    "#         filter_by_language_key,\n",
    "#         language_key_template=\"language_instruction*\"\n",
    "#     )\n",
    "# ))\n",
    "\n",
    "# for filter_fcn_spec in filter_functions:\n",
    "#     ds = ds.filter((filter_fcn_spec))\n",
    "\n",
    "# sample episode + resize to 256x256 (default third-person cam resolution)\n",
    "it = iter(ds)\n",
    "episode = next(it)\n",
    "print(episode)\n",
    "\n",
    "# images = [cv2.resize(np.array(step['observation']['image_0']), (256, 256)) for step in steps]\n",
    "# mediapy.show_video(images, fps=10)"
   ]
  }
 ],
 "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}