quickstart.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "61a62ae2",
   "metadata": {},
   "source": [
    "# Quick Start on Evaluating HPTs\n",
    "\n",
    "This notebook shows how to build a policy and run it on a simulation benchmark\n",
    "\n",
    "### Setup\n",
    "1. pip install -e ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c529fa2",
   "metadata": {
    "metadata": {}
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import hydra\n",
    "import torch\n",
    "from hpt.utils import utils, model_utils\n",
    "from omegaconf import OmegaConf\n",
    "from hpt.models.policy import Policy\n",
    "import os\n",
    "from torch.utils import data\n",
    "from torch.utils.data import DataLoader, RandomSampler\n",
    "\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2f0e703",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Set up Config and Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bd940cc",
   "metadata": {
    "metadata": {}
   },
   "outputs": [],
   "source": [
    "from hydra import compose, initialize\n",
    "\n",
    "\n",
    "policy = Policy.from_pretrained(\"hf://liruiw/hpt-base\")\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "domain = \"mujoco_metaworld\"\n",
    "with initialize(version_base=\"1.2\", config_path=\"experiments/configs\"):\n",
    "    cfg = compose(config_name=\"config\", overrides=[f\"env={domain}\"])\n",
    "\n",
    "cfg.dataset.episode_cnt = 10 # modify\n",
    "dataset = hydra.utils.instantiate(\n",
    "        cfg.dataset, dataset_name=domain, env_rollout_fn=cfg.dataset_generator_func, **cfg.dataset\n",
    "    )\n",
    "normalizer = dataset.get_normalizer()\n",
    "\n",
    "####### set up model\n",
    "utils.update_network_dim(cfg, dataset, policy)\n",
    "policy.init_domain_stem(domain, cfg.stem)\n",
    "policy.init_domain_head(domain, normalizer, cfg.head)\n",
    "policy.finalize_modules()\n",
    "policy.print_model_stats()\n",
    "policy.to(device)\n",
    "print(\"policy action normalizer:\", policy.normalizer[domain].params_dict[\"action\"][\"input_stats\"].max)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a78065c",
   "metadata": {},
   "source": [
    "### Train one iteration\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6aa08d59",
   "metadata": {
    "metadata": {}
   },
   "outputs": [],
   "source": [
    " \n",
    "train_loader = data.DataLoader(dataset, **cfg.dataloader)\n",
    "batch = next(iter(train_loader))\n",
    "batch[\"data\"] = utils.dict_apply(batch[\"data\"], lambda x: x.to(device).float())\n",
    "output = policy.compute_loss(batch)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da9d74be",
   "metadata": {},
   "source": [
    "### Run Rollout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b8417f3",
   "metadata": {
    "metadata": {}
   },
   "outputs": [],
   "source": [
    "from env.mujoco.metaworld.envs.mujoco.sawyer_xyz.test_scripted_policies import ALL_ENVS\n",
    "from collections import OrderedDict\n",
    "RESOLUTION = (128, 128)\n",
    "camera_name=\"view_1\"\n",
    "env_name = \"reach-v2\"\n",
    "def get_observation_dict(o, img):\n",
    "    step_data = {\"state\": o, \"image\": img}\n",
    "    return OrderedDict(step_data)\n",
    "\n",
    "env = ALL_ENVS[env_name]()\n",
    "env._partially_observable = False\n",
    "env._freeze_rand_vec = False\n",
    "env._set_task_called = True\n",
    "img = env.sim.render(*RESOLUTION, mode=\"offscreen\", camera_name=camera_name)[:, :, ::-1].copy()\n",
    "o = env.reset()\n",
    "step_data = get_observation_dict(o, img)\n",
    "policy.reset()\n",
    "\n",
    "for _ in range(env.max_path_length):\n",
    "    a = policy.get_action(step_data)\n",
    "    o, r, done, info = env.step(a)\n",
    "    img = env.sim.render(*RESOLUTION, mode=\"offscreen\", camera_name=camera_name)[:, :, ::-1]\n",
    "    step_data = get_observation_dict(o, img)\n",
    "    print(a)\n",
    "    break\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d0847e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import time\n",
    "import collections\n",
    "class FPS:\n",
    "    def __init__(self, avarageof=50):\n",
    "        self.frametimestamps = collections.deque(maxlen=avarageof)\n",
    "\n",
    "    def __call__(self):\n",
    "        self.frametimestamps.append(time.time())\n",
    "        if len(self.frametimestamps) > 1:\n",
    "            return len(self.frametimestamps) / (self.frametimestamps[-1] - self.frametimestamps[0])\n",
    "        else:\n",
    "            return 0.0\n",
    "        \n",
    "fps_measure = FPS()\n",
    "for _ in range(50):\n",
    "    output = policy.get_action(step_data)\n",
    "    fps_measure()\n",
    "print(f\"FPS: {fps_measure():.3f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}