process_traj.ipynb

{
 "cells": [
  {
   "cell_type": "code",
<<<<<<< HEAD
   "execution_count": 8,
=======
   "execution_count": 1,
>>>>>>> origin/main
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import jax\n",
    "from jax import numpy as jp\n",
    "\n",
    "import mujoco\n",
    "from mujoco import mjx\n",
    "\n",
    "import h5py\n",
    "import pickle\n",
    "\n",
    "import mocap_preprocess as mp\n",
    "\n",
    "data_path = \"clips/all_snips.p\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(data_path, \"rb\") as f:\n",
    "    all_clips = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "snips_name_to_idx = {\n",
    "    name.split(\"/\")[-1].split(\".\")[0]: idx\n",
    "    for idx, name in enumerate(all_clips[\"snips_order\"])\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "132"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "snips_name_to_idx[\"Walk_49\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unexpected exception formatting exception. Falling back to standard exception\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Traceback (most recent call last):\n",
      "  File \"/home/talmolab/anaconda3/envs/vnl/lib/python3.11/site-packages/IPython/core/interactiveshell.py\", line 3553, in run_code\n",
      "    exec(code_obj, self.user_global_ns, self.user_ns)\n",
      "  File \"/tmp/ipykernel_79123/4228247411.py\", line 2, in <module>\n",
      "    mp.process(\n",
      "  File \"/home/talmolab/Desktop/Salk-Research/VNL-Brax-Imitation/mocap_preprocess.py\", line 64, in process\n",
      "    physics = mjcf.Physics.from_mjcf_model(walker.mjcf_model)\n",
      "                                           ^^^^^^^^^^^^^^^^^\n",
      "AttributeError: module 'walker' has no attribute 'mjcf_model'\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/talmolab/anaconda3/envs/vnl/lib/python3.11/site-packages/IPython/core/interactiveshell.py\", line 2144, in showtraceback\n",
      "    stb = self.InteractiveTB.structured_traceback(\n",
      "          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/talmolab/anaconda3/envs/vnl/lib/python3.11/site-packages/IPython/core/ultratb.py\", line 1435, in structured_traceback\n",
      "    return FormattedTB.structured_traceback(\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/talmolab/anaconda3/envs/vnl/lib/python3.11/site-packages/IPython/core/ultratb.py\", line 1326, in structured_traceback\n",
      "    return VerboseTB.structured_traceback(\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/talmolab/anaconda3/envs/vnl/lib/python3.11/site-packages/IPython/core/ultratb.py\", line 1173, in structured_traceback\n",
      "    formatted_exception = self.format_exception_as_a_whole(etype, evalue, etb, number_of_lines_of_context,\n",
      "                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/talmolab/anaconda3/envs/vnl/lib/python3.11/site-packages/IPython/core/ultratb.py\", line 1088, in format_exception_as_a_whole\n",
      "    frames.append(self.format_record(record))\n",
      "                  ^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/talmolab/anaconda3/envs/vnl/lib/python3.11/site-packages/IPython/core/ultratb.py\", line 970, in format_record\n",
      "    frame_info.lines, Colors, self.has_colors, lvals\n",
      "    ^^^^^^^^^^^^^^^^\n",
      "  File \"/home/talmolab/anaconda3/envs/vnl/lib/python3.11/site-packages/IPython/core/ultratb.py\", line 792, in lines\n",
      "    return self._sd.lines\n",
      "           ^^^^^^^^^^^^^^\n",
      "  File \"/home/talmolab/anaconda3/envs/vnl/lib/python3.11/site-packages/stack_data/utils.py\", line 145, in cached_property_wrapper\n",
      "    value = obj.__dict__[self.func.__name__] = self.func(obj)\n",
      "                                               ^^^^^^^^^^^^^^\n",
      "  File \"/home/talmolab/anaconda3/envs/vnl/lib/python3.11/site-packages/stack_data/core.py\", line 698, in lines\n",
      "    pieces = self.included_pieces\n",
      "             ^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/talmolab/anaconda3/envs/vnl/lib/python3.11/site-packages/stack_data/utils.py\", line 145, in cached_property_wrapper\n",
      "    value = obj.__dict__[self.func.__name__] = self.func(obj)\n",
      "                                               ^^^^^^^^^^^^^^\n",
      "  File \"/home/talmolab/anaconda3/envs/vnl/lib/python3.11/site-packages/stack_data/core.py\", line 649, in included_pieces\n",
      "    pos = scope_pieces.index(self.executing_piece)\n",
      "                             ^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/talmolab/anaconda3/envs/vnl/lib/python3.11/site-packages/stack_data/utils.py\", line 145, in cached_property_wrapper\n",
      "    value = obj.__dict__[self.func.__name__] = self.func(obj)\n",
      "                                               ^^^^^^^^^^^^^^\n",
      "  File \"/home/talmolab/anaconda3/envs/vnl/lib/python3.11/site-packages/stack_data/core.py\", line 628, in executing_piece\n",
      "    return only(\n",
      "           ^^^^^\n",
      "  File \"/home/talmolab/anaconda3/envs/vnl/lib/python3.11/site-packages/executing/executing.py\", line 164, in only\n",
      "    raise NotOneValueFound('Expected one value, found 0')\n",
      "executing.executing.NotOneValueFound: Expected one value, found 0\n"
     ]
    }
   ],
   "source": [
    "idx = snips_name_to_idx[\"Walk_49\"]\n",
    "mp.process(\n",
    "    data_path,\n",
    "    \"walk_49.p\",\n",
    "    start_step=idx * 250,\n",
    "    clip_length=250,\n",
    "    n_steps=250,\n",
    "    ref_steps=(1, 2, 3, 4, 5),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "____"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Contents of transform_snips_250.h5:\n",
      "clip_0 (Group)\n",
      "  props (Group)\n",
      "  walkers (Group)\n",
      "    walker_0 (Group)\n",
      "      angular_velocity (Dataset): (3, 256)\n",
      "      appendages (Dataset): (15, 256)\n",
      "      body_positions (Dataset): (54, 256)\n",
      "      body_quaternions (Dataset): (72, 256)\n",
      "      center_of_mass (Dataset): (3, 256)\n",
      "      end_effectors (Dataset): (12, 256)\n",
      "      joints (Dataset): (67, 256)\n",
      "      joints_velocity (Dataset): (67, 256)\n",
      "      markers (Dataset): (0,)\n",
      "      position (Dataset): (3, 256)\n",
      "      quaternion (Dataset): (4, 256)\n",
      "      scaling (Dataset): (0,)\n",
      "      velocity (Dataset): (3, 256)\n",
      "clip_250 (Group)\n",
      "  props (Group)\n",
      "  walkers (Group)\n",
      "    walker_0 (Group)\n",
      "      angular_velocity (Dataset): (3, 250)\n",
      "      appendages (Dataset): (15, 250)\n",
      "      body_positions (Dataset): (54, 250)\n",
      "      body_quaternions (Dataset): (72, 250)\n",
      "      center_of_mass (Dataset): (3, 250)\n",
      "      end_effectors (Dataset): (12, 250)\n",
      "      joints (Dataset): (67, 250)\n",
      "      joints_velocity (Dataset): (67, 250)\n",
      "      markers (Dataset): (0,)\n",
      "      position (Dataset): (3, 250)\n",
      "      quaternion (Dataset): (4, 250)\n",
      "      scaling (Dataset): (0,)\n",
      "      velocity (Dataset): (3, 250)\n"
     ]
    }
   ],
   "source": [
    "import h5py\n",
    "\n",
    "\n",
    "# Prints the h5 file hierarchy and the dataset shapes\n",
    "def print_hierarchy(h5file, indent=\"\"):\n",
    "    for key, item in h5file.items():\n",
    "        if isinstance(item, h5py.Group):\n",
    "            print(f\"{indent}{key} (Group)\")\n",
    "            print_hierarchy(item, indent + \"  \")\n",
    "        elif isinstance(item, h5py.Dataset):\n",
    "            print(f\"{indent}{key} (Dataset): {item.shape}\")\n",
    "\n",
    "\n",
    "filename = \"transform_snips_250.h5\"\n",
    "with h5py.File(filename, \"r\") as f:\n",
    "    print(f\"Contents of {filename}:\")\n",
    "    print_hierarchy(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "p = \"/home/charles/github/VNL-Brax-Imitation/transform_snips_250_clip_0.p\"\n",
    "\n",
    "with open(p, \"rb\") as f:\n",
    "    d = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(256, 3)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d.angular_velocity.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rl",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}