scripts/scripted_policy_movies.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import os\n",
    "import functools\n",
    "\n",
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "from tests.metaworld.envs.mujoco.sawyer_xyz.test_scripted_policies import ALL_ENVS, test_cases_latest_nonoise\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def trajectory_generator(env, policy, act_noise_pct, res=(640, 480), camera='corner'):\n",
    "    action_space_ptp = env.action_space.high - env.action_space.low\n",
    "\n",
    "    env.reset()\n",
    "    env.reset_model()\n",
    "    o = env.reset()\n",
    "\n",
    "    for _ in range(env.max_path_length):\n",
    "        a = policy.get_action(o)\n",
    "        a = np.random.normal(a, act_noise_pct * action_space_ptp)\n",
    "\n",
    "        o, r, done, info = env.step(a)\n",
    "        # Camera is one of ['corner', 'topview', 'behindGripper', 'gripperPOV']\n",
    "        yield r, done, info, env.sim.render(*res, mode='offscreen', camera_name=camera)[:,:,::-1]\n",
    "\n",
    "def writer_for(tag, fps, res):\n",
    "    if not os.path.exists('movies'):\n",
    "        os.mkdir('movies')\n",
    "    return cv2.VideoWriter(\n",
    "        f'movies/{tag}.avi',\n",
    "        cv2.VideoWriter_fourcc('M','J','P','G'),\n",
    "        fps,\n",
    "        res\n",
    "    )"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "resolution = (1920, 1080)\n",
    "camera = 'behindGripper' # one of ['corner', 'topview', 'behindGripper', 'gripperPOV']\n",
    "flip=True # if True, flips output image 180 degrees\n",
    "\n",
    "config = [\n",
    "    # env, action noise pct, cycles, quit on success\n",
    "    ('assembly-v2', np.zeros(4), 3, True),\n",
    "    ('basketball-v2', np.zeros(4), 3, True),\n",
    "    ('bin-picking-v2', np.zeros(4), 3, True),\n",
    "    ('box-close-v2', np.zeros(4), 3, True),\n",
    "    ('button-press-topdown-v2', np.zeros(4), 3, True),\n",
    "    ('button-press-topdown-wall-v2', np.zeros(4), 3, True),\n",
    "    ('button-press-v2', np.zeros(4), 3, True),\n",
    "    ('button-press-wall-v2', np.zeros(4), 3, True),\n",
    "    ('coffee-button-v2', np.zeros(4), 3, True),\n",
    "    ('coffee-pull-v2', np.zeros(4), 3, True),\n",
    "    ('coffee-push-v2', np.zeros(4), 3, True),\n",
    "    ('dial-turn-v2', np.zeros(4), 3, True),\n",
    "    ('disassemble-v2', np.zeros(4), 3, True),\n",
    "    ('door-close-v2', np.zeros(4), 3, True),\n",
    "    ('door-lock-v2', np.zeros(4), 3, True),\n",
    "    ('door-open-v2', np.zeros(4), 3, True),\n",
    "    ('door-unlock-v2', np.zeros(4), 3, True),\n",
    "    ('hand-insert-v2', np.zeros(4), 3, True),\n",
    "    ('drawer-close-v2', np.zeros(4), 3, True),\n",
    "    ('drawer-open-v2', np.zeros(4), 3, True),\n",
    "    ('faucet-open-v2', np.zeros(4), 3, True),\n",
    "    ('faucet-close-v2', np.zeros(4), 3, True),\n",
    "    ('hammer-v2', np.zeros(4), 3, True),\n",
    "    ('handle-press-side-v2', np.zeros(4), 3, True),\n",
    "    ('handle-press-v2', np.zeros(4), 3, True),\n",
    "    ('handle-pull-side-v2', np.zeros(4), 3, True),\n",
    "    ('handle-pull-v2', np.zeros(4), 3, True),\n",
    "    ('lever-pull-v2', np.zeros(4), 3, True),\n",
    "    ('peg-insert-side-v2', np.zeros(4), 3, True),\n",
    "    ('pick-place-wall-v2', np.zeros(4), 3, True),\n",
    "    ('pick-out-of-hole-v2', np.zeros(4), 3, True),\n",
    "    ('reach-v2', np.zeros(4), 3, True),\n",
    "    ('push-back-v2', np.zeros(4), 3, True),\n",
    "    ('push-v2', np.zeros(4), 3, True),\n",
    "    ('pick-place-v2', np.zeros(4), 3, True),\n",
    "    ('plate-slide-v2', np.zeros(4), 3, True),\n",
    "    ('plate-slide-side-v2', np.zeros(4), 3, True),\n",
    "    ('plate-slide-back-v2', np.zeros(4), 3, True),\n",
    "    ('plate-slide-back-side-v2', np.zeros(4), 3, True),\n",
    "    ('peg-insert-side-v2', np.zeros(4), 3, True),\n",
    "    ('peg-unplug-side-v2', np.zeros(4), 3, True),\n",
    "    ('soccer-v2', np.zeros(4), 3, True),\n",
    "    ('stick-push-v2', np.zeros(4), 3, True),\n",
    "    ('stick-pull-v2', np.zeros(4), 3, True),\n",
    "    ('push-wall-v2', np.zeros(4), 3, True),\n",
    "    ('push-v2', np.zeros(4), 3, True),\n",
    "    ('reach-wall-v2', np.zeros(4), 3, True),\n",
    "    ('reach-v2', np.zeros(4), 3, True),\n",
    "    ('shelf-place-v2', np.zeros(4), 3, True),\n",
    "    ('sweep-into-v2', np.zeros(4), 3, True),\n",
    "    ('sweep-v2', np.zeros(4), 3, True),\n",
    "    ('window-open-v2', np.zeros(4), 3, True),\n",
    "    ('window-close-v2', np.zeros(4), 3, True),\n",
    "]\n",
    "\n",
    "for env, noise, cycles, quit_on_success in config:\n",
    "    tag = env + '-noise-' + np.array2string(noise, precision=2, separator=',', suppress_small=True)\n",
    "\n",
    "    policy = functools.reduce(lambda a,b : a if a[0] == env else b, test_cases_latest_nonoise)[1]\n",
    "    env = ALL_ENVS[env]()\n",
    "    env._partially_observable = False\n",
    "    env._freeze_rand_vec = False\n",
    "    env._set_task_called = True\n",
    "\n",
    "    writer = writer_for(tag, env.metadata['video.frames_per_second'], resolution)\n",
    "    for _ in range(cycles):\n",
    "        for r, done, info, img in trajectory_generator(env, policy, noise, resolution, camera):\n",
    "            if flip: img = cv2.rotate(img, cv2.ROTATE_180)\n",
    "            writer.write(img)\n",
    "            if quit_on_success and info['success']:\n",
    "                break\n",
    "\n",
    "    writer.release()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}