RLBench_IO_Test.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test the RLBench input output functionality for the internal tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rlbench.action_modes.action_mode import MoveArmThenGripper\n",
    "from rlbench.action_modes.arm_action_modes import EndEffectorPoseViaPlanning\n",
    "from rlbench.action_modes.gripper_action_modes import Discrete\n",
    "from rlbench.environment import Environment\n",
    "from rlbench.observation_config import ObservationConfig\n",
    "from rlbench.tasks import FS10_V1, ReachTarget, PickUpCup, PutGroceriesInCupboard, StackBlocks, PickAndLift, OpenDrawer, CloseDrawer, HangFrameOnHanger, EmptyDishwasher, BasketballInHoop\n",
    "import torch\n",
    "import numpy as np\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Launch the RLBench environment with all the observations enabled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs_config = ObservationConfig()\n",
    "obs_config.set_all(True)\n",
    "\n",
    "env = Environment(\n",
    "    action_mode=MoveArmThenGripper(\n",
    "        arm_action_mode=EndEffectorPoseViaPlanning(absolute_mode=True), gripper_action_mode=Discrete()),\n",
    "    obs_config=ObservationConfig(),\n",
    "    headless=False)\n",
    "env.launch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "20\n"
     ]
    }
   ],
   "source": [
    "import pprint as pp\n",
    "\n",
    "# Get the task\n",
    "task = env.get_task(ReachTarget)\n",
    "pp.pprint(task.variation_count())\n",
    "demo = task.get_demos(1, live_demos=True, random_selection=False)[0]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'touch the red ball with the panda gripper'\n"
     ]
    }
   ],
   "source": [
    "# Reset the task\n",
    "descriptions, obs = task.reset_to_demo(demo)\n",
    "instruction = descriptions[1] # Could make this random at some point\n",
    "pp.pprint(instruction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from data.utils import keypoint_discovery\n",
    "# pp.pprint(type(demo))\n",
    "# pp.pprint(len(demo))\n",
    "# keypoints = keypoint_discovery(demo)\n",
    "# pp.pprint(keypoints)\n",
    "# obs = demo[keypoints[0]]\n",
    "# type(obs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "array([ 4.20428336e-01, -1.77622497e-01,  7.97428846e-01,  2.23513280e-06,\n",
      "        1.00000000e+00,  1.28546322e-04, -1.44612830e-04,  1.00000000e+00])\n",
      "Step: 0 Reward: 1.0\n"
     ]
    }
   ],
   "source": [
    "from data.utils import keypoint_discovery\n",
    "keypoints = keypoint_discovery(demo)\n",
    "\n",
    "for i, keypoint in enumerate(keypoints):\n",
    "    obs = demo[keypoint]\n",
    "    action = np.concatenate([obs.gripper_pose, [obs.gripper_open]])\n",
    "    pp.pprint(action)\n",
    "    obs, reward, terminate = task.step(action)\n",
    "    print(f'Step: {i} Reward: {reward}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[CoppeliaSim:loadinfo]   done.\n"
     ]
    }
   ],
   "source": [
    "env.shutdown()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}