Course_OpenAIBaselines_Unit3.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using OpenAI with ROS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"img/robotignite_logo_text.png\" width=\"700\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unit 3: How to create a new Robot Environment for a new Robot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p style=\"background:green;color:white;\">SUMMARY</p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Estimated time to completion: <b>2 hours</b>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "In this unit, you will learn how to create a Robot Environment for a Moving Cube with a single disk in the roll axis using the **OpenAI ROS** structure. We will call it **MyCubeSingleDiskEnv**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p style=\"background:green;color:white;\">END OF SUMMARY</p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use the power of OpenAI through simple Robot-Gazebo-Env"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, you now understand how the **CartPole Environments structure** works. Now, we need to create one from scratch for the One-Disk-Moving-Cube. There are three main steps:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* We need to create a **Robot Environment** that has the basic functions needed to use the RoboCube.\n",
    "* We need to decide how we move the RoboCube and how to get the sensor data.\n",
    "* We need to create functions that allow environments that inherit from our class to retrieve sensor data and access the cube's functionality, without knowing all the ROS related stuff."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this created, you will have a **Robot Environment** that can be used by any **Task Environment** that you create afterwards."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  0. Create a package to place all your code in"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will first create a package that will hold all the code that you generate in these next chapters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<table style=\"float:left;background: #407EAF\">\n",
    "<tr>\n",
    "<th>\n",
    "<p class=\"transparent\">Execute in WebShell #1</p>\n",
    "</th>\n",
    "</tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "roscd; cd ../src; "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "catkin_create_pkg my_moving_cube_pkg rospy openai_ros"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "cd my_moving_cube_pkg/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "mkdir scripts; cd scripts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "touch my_cube_single_disk_env.py;chmod +x my_cube_single_disk_env.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  1. Create a basic Robot Environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will start from an empty template that we give you here, which has all the functions needed for the **RobotGazeboEnv**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p style=\"background:#407EAF;color:white;\">**template_my_robot_env.py**</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from openai_ros import robot_gazebo_env\n",
    "\n",
    "\n",
    "class MyRobotEnv(robot_gazebo_env.RobotGazeboEnv):\n",
    "    \"\"\"Superclass for all Robot environments.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        \"\"\"Initializes a new Robot environment.\n",
    "        \"\"\"\n",
    "        # Variables that we give through the constructor.\n",
    "\n",
    "        # Internal Vars\n",
    "        self.controllers_list = ['my_robot_controller1','my_robot_controller2', ..., 'my_robot_controllerX']\n",
    "\n",
    "        self.robot_name_space = \"my_robot_namespace\"\n",
    "\n",
    "        reset_controls_bool = True or False\n",
    "        \n",
    "        # We launch the init function of the Parent Class robot_gazebo_env.RobotGazeboEnv\n",
    "        \n",
    "        super(MyRobotEnv, self).__init__(controllers_list=self.controllers_list,\n",
    "                                                robot_name_space=self.robot_name_space,\n",
    "                                                reset_controls=reset_controls_bool)\n",
    "\n",
    "    # Methods needed by the RobotGazeboEnv\n",
    "    # ----------------------------\n",
    "    \n",
    "    \n",
    "\n",
    "    def _check_all_systems_ready(self):\n",
    "        \"\"\"\n",
    "        Checks that all the sensors, publishers and other simulation systems are\n",
    "        operational.\n",
    "        \"\"\"\n",
    "        # TODO\n",
    "        return True\n",
    "    \n",
    "    # Methods that the TaskEnvironment will need to define here as virtual\n",
    "    # because they will be used in RobotGazeboEnv GrandParentClass and defined in the\n",
    "    # TaskEnvironment.\n",
    "    # ----------------------------\n",
    "    def _set_init_pose(self):\n",
    "        \"\"\"Sets the Robot in its init pose\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "    \n",
    "    \n",
    "    def _init_env_variables(self):\n",
    "        \"\"\"Inits variables needed to be initialised each time we reset at the start\n",
    "        of an episode.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def _compute_reward(self, observations, done):\n",
    "        \"\"\"Calculates the reward to give based on the observations given.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def _set_action(self, action):\n",
    "        \"\"\"Applies the given action to the simulation.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def _get_obs(self):\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def _is_done(self, observations):\n",
    "        \"\"\"Checks if episode done based on observations given.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "        \n",
    "    # Methods that the TaskEnvironment will need.\n",
    "    # ----------------------------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see here that the template is divided into four parts:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* 1.1- Initialization of the class **MyRobotEnv**, or whatever name you choose.\n",
    "* 1.2- Definition of the methods needed by the Gazebo Environment, which were declared virtually inside RobotGazeboEnv.\n",
    "* 1.3- Virtual definition of methods that the Task Environment will need to define here as virtual because they will be used in the RobotGazeboEnv GrandParentClass and defined in the Task Environment.\n",
    "* 1.4- Definition of methods that the Task Environment will need."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.1- Initialization of the class MyRobotEnv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from openai_ros import robot_gazebo_env\n",
    "\n",
    "class MyRobotEnv(robot_gazebo_env.RobotGazeboEnv):\n",
    "    \"\"\"Superclass for all CubeSingleDisk environments.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, init_roll_vel):\n",
    "        \"\"\"Initializes a new CubeSingleDisk environment.\n",
    "        \"\"\"\n",
    "        # Variables that we give through the constructor.\n",
    "        self.init_roll_vel = init_roll_vel\n",
    "\n",
    "        self.controllers_list = ['my_robot_controller1','my_robot_controller2', ..., 'my_robot_controllerX']\n",
    "\n",
    "        self.robot_name_space = \"my_robot_namespace\"\n",
    "\n",
    "        reset_controls_bool = True or False\n",
    "        \n",
    "        # We launch the init function of the Parent Class robot_gazebo_env.RobotGazeboEnv\n",
    "        \n",
    "        super(CubeSingleDiskEnv, self).__init__(controllers_list=self.controllers_list,\n",
    "                                                robot_name_space=self.robot_name_space,\n",
    "                                                reset_controls=reset_controls_bool)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we import the **robot_gazebo_env** from the python module folder openai_ros. Inside **robot_gazebo_env**, we find the class **RobotGazeboEnv**. In the class definition of **MyRobotEnv**, we give inheritance to **robot_gazebo_env.RobotGazeboEnv**. That means to look for the python module file **robot_gazebo_env**, and inside, get the class **RobotGazeboEnv**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from openai_ros import robot_gazebo_env\n",
    "\n",
    "class MyRobotEnv(robot_gazebo_env.RobotGazeboEnv):"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, you have to change the **MyRobotEnv** to the class name you prefer; in this case, let's call it **MyCubeSingleDiskEnv**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class MyCubeSingleDiskEnv(robot_gazebo_env.RobotGazeboEnv):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def __init__(self, init_roll_vel):\n",
    "    # Variables that we give through the constructor.\n",
    "    self.init_roll_vel = init_roll_vel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case, as you can see, we added a variable in the init function. These variables are the ones you want the **Task Environment** to pass to this new Robot Environment. We need it to pass the speed at the start of each episode to set the roll disk. For this scenario, it will most likely always be **0.0**, but it could change depending on the **Task Environment**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we define some variables that need to be passed to the **RobotGazeboEnv** super constructor method __init__. These are: **controllers_list, robot_name_space, and reset_controls**. These variables are used by the RobotGazeboEnv to know which controllers to reset each time a learning episode starts."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can make the controllers available in the simulation by just executing the following commands in the web shell:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<table style=\"float:left;background: #407EAF\">\n",
    "<tr>\n",
    "<th>\n",
    "<p class=\"transparent\">Execute in WebShell #1</p>\n",
    "</th>\n",
    "</tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "rosservice call /moving_cube/controller_manager/list_controllers \"{}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<table class=\"transparent float_l\">\n",
    "<tr>\n",
    "<th>\n",
    "<p style=\"background: #407EAF\">WebShell #1 Output</p>\n",
    "</th>\n",
    "</tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "controller:\n",
    "  -\n",
    "    name: \"joint_state_controller\"\n",
    "    state: \"running\"\n",
    "    type: \"joint_state_controller/JointStateController\"\n",
    "    claimed_resources:\n",
    "      -\n",
    "        hardware_interface: \"hardware_interface::JointStateInterface\"\n",
    "        resources: []\n",
    "  -\n",
    "    name: \"inertia_wheel_roll_joint_velocity_controller\"\n",
    "    state: \"running\"\n",
    "    type: \"effort_controllers/JointVelocityController\"\n",
    "    claimed_resources:\n",
    "      -\n",
    "        hardware_interface: \"hardware_interface::EffortJointInterface\"\n",
    "        resources: [inertia_wheel_roll_joint]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, the list of controllers should look something like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "self.controllers_list = ['joint_state_controller','inertia_wheel_roll_joint_velocity_controller']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And to get the namespace used in the controllers, just execute the following command to see all the controllers' namespace: "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<table style=\"float:left;background: #407EAF\">\n",
    "<tr>\n",
    "<th>\n",
    "<p class=\"transparent\">Execute in WebShell #1</p>\n",
    "</th>\n",
    "</tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "rostopic list | grep controller"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<table class=\"transparent float_l\">\n",
    "<tr>\n",
    "<th>\n",
    "<p style=\"background: #407EAF\">WebShell #1 Output</p>\n",
    "</th>\n",
    "</tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "/moving_cube/inertia_wheel_roll_joint_velocity_controller/command\n",
    "/moving_cube/inertia_wheel_roll_joint_velocity_controller/pid/parameter_descriptions\n",
    "/moving_cube/inertia_wheel_roll_joint_velocity_controller/pid/parameter_updates\n",
    "/moving_cube/inertia_wheel_roll_joint_velocity_controller/state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So. you should have the following variable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "self.robot_name_space = \"moving_cube\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we decide if we want to reset the controllers or not:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "reset_controls_bool = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally, you pass it to the creation **__init__** function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "super(MyCubeSingleDiskEnv, self).__init__(controllers_list=self.controllers_list,\n",
    "                                                robot_name_space=self.robot_name_space,\n",
    "                                                reset_controls=reset_controls_bool)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.2- Definition of the Methods needed by the RobotGazeboEnv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are methods that RobotGazeboEnv needs to have defined somewhere. And the place to be defined is in this RobotEnvironment because it has access to the robot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def _check_all_systems_ready(self):\n",
    "    \"\"\"\n",
    "    Checks that all the sensors, publishers and other simulation systems are\n",
    "    operational.\n",
    "    \"\"\"\n",
    "    # TODO\n",
    "    return True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have to return to the __init__ method and create all the subscribers and publishers for the One Disk Cube so that the check_all_systems_ready can work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "    def __init__(self):\n",
    "        \"\"\"Initializes a new CubeSingleDisk environment.\n",
    "\n",
    "        Args:\n",
    "        \"\"\"\n",
    "        # Variables that we give through the constructor.\n",
    "        # None\n",
    "\n",
    "        \n",
    "        self.controllers_list = ['joint_state_controller',\n",
    "                                 'inertia_wheel_roll_joint_velocity_controller'\n",
    "                                 ]\n",
    "\n",
    "        self.robot_name_space = \"moving_cube\"\n",
    "\n",
    "        # We launch the init function of the Parent Class robot_gazebo_env.RobotGazeboEnv\n",
    "        super(MyCubeSingleDiskEnv, self).__init__(controllers_list=self.controllers_list,\n",
    "                                                robot_name_space=self.robot_name_space,\n",
    "                                                reset_controls=True)\n",
    "\n",
    "\n",
    "\n",
    "        \"\"\"\n",
    "        To check any topic we need to have the simulations running, we need to do two things:\n",
    "        1) Unpause the simulation: without that the stream of data doesn't flow. This is for simulations\n",
    "        that are pause for whatever reason\n",
    "        2) If the simulation was running already for some reason, we need to reset the controllers.\n",
    "        This has to do with the fact that some plugins with tf don't understand the reset of the simulation\n",
    "        and need to be reset to work properly.\n",
    "        \"\"\"\n",
    "        self.gazebo.unpauseSim()\n",
    "        self.controllers_object.reset_controllers()\n",
    "        self._check_all_sensors_ready()\n",
    "\n",
    "        # We Start all the ROS related Subscribers and publishers\n",
    "        rospy.Subscriber(\"/moving_cube/joint_states\", JointState, self._joints_callback)\n",
    "        rospy.Subscriber(\"/moving_cube/odom\", Odometry, self._odom_callback)\n",
    "\n",
    "        self._roll_vel_pub = rospy.Publisher('/moving_cube/inertia_wheel_roll_joint_velocity_controller/command',\n",
    "                                             Float64, queue_size=1)\n",
    "\n",
    "        self._check_publishers_connection()\n",
    "\n",
    "        self.gazebo.pauseSim()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we create the **Subscribers** and **Publishers**. The important part is the **unpause**, **check_all_sensors_ready**, and **pause**. These are key to being able to reset the controllers and read the sensors. We use the objects created in the RobotGazeboEnv parent class so that we have access to it without having to know how it works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "self.gazebo.unpauseSim()\n",
    "self.controllers_object.reset_controllers()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we unpause the simulation. This allows us to reset controllers and make the first test to see if the sensors are working. Checking this is key for Ai learning because we need a reliable sensor and controller communication."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "self._check_all_sensors_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note** that inside this **MyCubeSingleDiskEnv** we use **_check_all_sensors_ready** which is an internal function, while the **RobotGazeboEnv** parent class will call the **_check_all_systems_ready**. We could also just useone function, but its separeted here to show the diference in who uses what function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have to define the method inside this class, like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "    def _check_all_systems_ready(self):\n",
    "        \"\"\"\n",
    "        Checks that all the sensors, publishers and other simulation systems are\n",
    "        operational.\n",
    "        \"\"\"\n",
    "        self._check_all_sensors_ready()\n",
    "        return True\n",
    "\n",
    "\n",
    "    # CubeSingleDiskEnv virtual methods\n",
    "    # ----------------------------\n",
    "\n",
    "    def _check_all_sensors_ready(self):\n",
    "        self._check_joint_states_ready()\n",
    "        self._check_odom_ready()\n",
    "        rospy.logdebug(\"ALL SENSORS READY\")\n",
    "\n",
    "    def _check_joint_states_ready(self):\n",
    "        self.joints = None\n",
    "        while self.joints is None and not rospy.is_shutdown():\n",
    "            try:\n",
    "                self.joints = rospy.wait_for_message(\"/moving_cube/joint_states\", JointState, timeout=1.0)\n",
    "                rospy.logdebug(\"Current moving_cube/joint_states READY=>\" + str(self.joints))\n",
    "\n",
    "            except:\n",
    "                rospy.logerr(\"Current moving_cube/joint_states not ready yet, retrying for getting joint_states\")\n",
    "        return self.joints\n",
    "\n",
    "    def _check_odom_ready(self):\n",
    "        self.odom = None\n",
    "        while self.odom is None and not rospy.is_shutdown():\n",
    "            try:\n",
    "                self.odom = rospy.wait_for_message(\"/moving_cube/odom\", Odometry, timeout=1.0)\n",
    "                rospy.logdebug(\"Current /moving_cube/odom READY=>\" + str(self.odom))\n",
    "\n",
    "            except:\n",
    "                rospy.logerr(\"Current /moving_cube/odom not ready yet, retrying for getting odom\")\n",
    "\n",
    "        return self.odom"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the case of the OneDiskCube, for sensors, we only have the odometry that tells us where the Cube Body is in the simulated world (**/moving_cube/odom**), and how the disk joint (speed, position, efforts) is through the **/moving_cube/joint_states**. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we have this, we know that ROS can establish a connection to these topics and they're ready, so we can declare the susbcribers now. So, let's go back to the __init__ method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "rospy.Subscriber(\"/moving_cube/joint_states\", JointState, self._joints_callback)\n",
    "rospy.Subscriber(\"/moving_cube/odom\", Odometry, self._odom_callback)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We add the necessary imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import rospy\n",
    "from sensor_msgs.msg import JointState\n",
    "from nav_msgs.msg import Odometry"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we also have to declare their callbacks, which will start to store the sensor data in **self.joints** and **self.odom**. Now we will always have the most updated sensor data, even when we pause the simulation, having updated observations for the learning algorithms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "    def _joints_callback(self, data):\n",
    "        self.joints = data\n",
    "    \n",
    "    def _odom_callback(self, data):\n",
    "        self.odom = data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We add the publisher now:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "self._roll_vel_pub = rospy.Publisher('/moving_cube/inertia_wheel_roll_joint_velocity_controller/command',\n",
    "                                             Float64, queue_size=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the necessary imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from std_msgs.msg import Float64"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we define the **_check_publishers_connection**, to check that our publisher is ready to receive the speed commands and doesn't lose any messages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "    def _check_publishers_connection(self):\n",
    "        \"\"\"\n",
    "        Checks that all the publishers are working\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        rate = rospy.Rate(10)  # 10hz\n",
    "        while self._roll_vel_pub.get_num_connections() == 0 and not rospy.is_shutdown():\n",
    "            rospy.logdebug(\"No susbribers to _roll_vel_pub yet so we wait and try again\")\n",
    "            try:\n",
    "                rate.sleep()\n",
    "            except rospy.ROSInterruptException:\n",
    "                # This is to avoid error when world is rested, time when backwards.\n",
    "                pass\n",
    "        rospy.logdebug(\"_roll_vel_pub Publisher Connected\")\n",
    "\n",
    "        rospy.logdebug(\"All Publishers READY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we have to define the method that will move the disks, publishing in our **self._roll_vel_pub** publisher, called **move_joints**. This method will be used internally in **MyCubeSingleDiskEnv** and also in the **Task Environment**. We go to this section of the template:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "    # Methods that the TaskEnvironment will need.\n",
    "    # ----------------------------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We place our code here because it's a function that will be used by the **TaskEnv**, and this way, we keep the code organized:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "    def move_joints(self, roll_speed):\n",
    "        joint_speed_value = Float64()\n",
    "        joint_speed_value.data = roll_speed\n",
    "        rospy.logdebug(\"Single Disk Roll Velocity>>\" + str(joint_speed_value))\n",
    "        self._roll_vel_pub.publish(joint_speed_value)\n",
    "        self.wait_until_roll_is_in_vel(joint_speed_value.data)\n",
    "    \n",
    "    def wait_until_roll_is_in_vel(self, velocity):\n",
    "    \n",
    "        rate = rospy.Rate(10)\n",
    "        start_wait_time = rospy.get_rostime().to_sec()\n",
    "        end_wait_time = 0.0\n",
    "        epsilon = 0.1\n",
    "        v_plus = velocity + epsilon\n",
    "        v_minus = velocity - epsilon\n",
    "        while not rospy.is_shutdown():\n",
    "            joint_data = self._check_joint_states_ready()\n",
    "            roll_vel = joint_data.velocity[0]\n",
    "            rospy.logdebug(\"VEL=\" + str(roll_vel) + \", ?RANGE=[\" + str(v_minus) + \",\"+str(v_plus)+\"]\")\n",
    "            are_close = (roll_vel <= v_plus) and (roll_vel > v_minus)\n",
    "            if are_close:\n",
    "                rospy.logdebug(\"Reached Velocity!\")\n",
    "                end_wait_time = rospy.get_rostime().to_sec()\n",
    "                break\n",
    "            rospy.logdebug(\"Not there yet, keep waiting...\")\n",
    "            rate.sleep()\n",
    "        delta_time = end_wait_time- start_wait_time\n",
    "        rospy.logdebug(\"[Wait Time=\" + str(delta_time)+\"]\")\n",
    "        return delta_time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It essentially gets a given roll speed and publishes that speed through the ROS publisher **self._roll_vel_pub**. The last part is also vital because it guarantees that all the actions are executed and aren't overrun by the next one. This is the method **wait_until_roll_is_in_vel**. This method will wait until the roll disk wheel reaches the desired speed, with a certain error. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.3- Virtual definition of methods that the Task Environment will need to define"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These methods are basically virtual methods that allow them to be used in **RobotGazeboEnv** and **MyCubeSingleDiskEnv**, but have to be defined upstream in the **TaskEnvironment**. We do this because these methods relate to how the Task is to be learned. So, it has to be defined on that level."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "    # Methods that the TaskEnvironment will need to define here as virtual\n",
    "    # because they will be used in RobotGazeboEnv GrandParentClass and defined in the\n",
    "    # TaskEnvironment.\n",
    "    # ----------------------------\n",
    "    def _set_init_pose(self):\n",
    "        \"\"\"Sets the Robot in its init pose\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "    \n",
    "    def _init_env_variables(self):\n",
    "        \"\"\"Inits variables needed to be initialised each time we reset at the start\n",
    "        of an episode.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def _compute_reward(self, observations, done):\n",
    "        \"\"\"Calculates the reward to give based on the observations given.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def _set_action(self, action):\n",
    "        \"\"\"Applies the given action to the simulation.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def _get_obs(self):\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def _is_done(self, observations):\n",
    "        \"\"\"Checks if episode done based on observations given.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.4- Definition of Methods that the TaskEnvironment will need"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These methods will be used by the **TaskEnvironment**, but need to be defined in this RobotEvironment in order to have the necessary access to publishers and subscribers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def move_joints(self, roll_speed):\n",
    "    joint_speed_value = Float64()\n",
    "    joint_speed_value.data = roll_speed\n",
    "    rospy.logdebug(\"Single Disk Roll Velocity>>\" + str(joint_speed_value))\n",
    "    self._roll_vel_pub.publish(joint_speed_value)\n",
    "    self.wait_until_roll_is_in_vel(joint_speed_value.data)\n",
    "\n",
    "def wait_until_roll_is_in_vel(self, velocity):\n",
    "\n",
    "    rate = rospy.Rate(10)\n",
    "    start_wait_time = rospy.get_rostime().to_sec()\n",
    "    end_wait_time = 0.0\n",
    "    epsilon = 0.1\n",
    "    v_plus = velocity + epsilon\n",
    "    v_minus = velocity - epsilon\n",
    "    while not rospy.is_shutdown():\n",
    "        joint_data = self._check_joint_states_ready()\n",
    "        roll_vel = joint_data.velocity[0]\n",
    "        rospy.logdebug(\"VEL=\" + str(roll_vel) + \", ?RANGE=[\" + str(v_minus) + \",\"+str(v_plus)+\"]\")\n",
    "        are_close = (roll_vel <= v_plus) and (roll_vel > v_minus)\n",
    "        if are_close:\n",
    "            rospy.logdebug(\"Reached Velocity!\")\n",
    "            end_wait_time = rospy.get_rostime().to_sec()\n",
    "            break\n",
    "        rospy.logdebug(\"Not there yet, keep waiting...\")\n",
    "        rate.sleep()\n",
    "    delta_time = end_wait_time- start_wait_time\n",
    "    rospy.logdebug(\"[Wait Time=\" + str(delta_time)+\"]\")\n",
    "    return delta_time\n",
    "\n",
    "def get_joints(self):\n",
    "    return self.joints\n",
    "\n",
    "def get_odom(self):\n",
    "    return self.odom"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define only RAW data methods to access the sensor data **get_joints** and **get_odom**. The **move_joints** are also used by the **TaskEnv**, but we defined them already in the **_set_init_pose** method."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And don't forget the imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Final Result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p style=\"background:#407EAF;color:white;\">**my_cube_single_disk_env.py**</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy\n",
    "import rospy\n",
    "from openai_ros import robot_gazebo_env\n",
    "from std_msgs.msg import Float64\n",
    "from sensor_msgs.msg import JointState\n",
    "from nav_msgs.msg import Odometry\n",
    "\n",
    "\n",
    "\n",
    "class MyCubeSingleDiskEnv(robot_gazebo_env.RobotGazeboEnv):\n",
    "    \"\"\"Superclass for all CubeSingleDisk environments.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        \"\"\"Initializes a new CubeSingleDisk environment.\n",
    "\n",
    "        Args:\n",
    "        \"\"\"\n",
    "        # Variables that we give through the constructor.\n",
    "        # None in this case\n",
    "\n",
    "        # Internal Vars\n",
    "        self.controllers_list = ['joint_state_controller',\n",
    "                                 'inertia_wheel_roll_joint_velocity_controller'\n",
    "                                 ]\n",
    "\n",
    "        self.robot_name_space = \"moving_cube\"\n",
    "\n",
    "        # We launch the init function of the Parent Class robot_gazebo_env.RobotGazeboEnv\n",
    "        super(MyCubeSingleDiskEnv, self).__init__(controllers_list=self.controllers_list,\n",
    "                                                robot_name_space=self.robot_name_space,\n",
    "                                                reset_controls=True)\n",
    "\n",
    "\n",
    "\n",
    "        \"\"\"\n",
    "        To check any topic we need to have the simulations running, we need to do two things:\n",
    "        1) Unpause the simulation: without that th stream of data doesnt flow. This is for simulations\n",
    "        that are pause for whatever the reason\n",
    "        2) If the simulation was running already for some reason, we need to reset the controlers.\n",
    "        This has to do with the fact that some plugins with tf, dont understand the reset of the simulation\n",
    "        and need to be reseted to work properly.\n",
    "        \"\"\"\n",
    "        self.gazebo.unpauseSim()\n",
    "        self.controllers_object.reset_controllers()\n",
    "        self._check_all_sensors_ready()\n",
    "\n",
    "        # We Start all the ROS related Subscribers and publishers\n",
    "        rospy.Subscriber(\"/moving_cube/joint_states\", JointState, self._joints_callback)\n",
    "        rospy.Subscriber(\"/moving_cube/odom\", Odometry, self._odom_callback)\n",
    "\n",
    "        self._roll_vel_pub = rospy.Publisher('/moving_cube/inertia_wheel_roll_joint_velocity_controller/command',\n",
    "                                             Float64, queue_size=1)\n",
    "\n",
    "        self._check_publishers_connection()\n",
    "\n",
    "        self.gazebo.pauseSim()\n",
    "\n",
    "    # Methods needed by the RobotGazeboEnv\n",
    "    # ----------------------------\n",
    "    \n",
    "\n",
    "    def _check_all_systems_ready(self):\n",
    "        \"\"\"\n",
    "        Checks that all the sensors, publishers and other simulation systems are\n",
    "        operational.\n",
    "        \"\"\"\n",
    "        self._check_all_sensors_ready()\n",
    "        return True\n",
    "\n",
    "\n",
    "    # CubeSingleDiskEnv virtual methods\n",
    "    # ----------------------------\n",
    "\n",
    "    def _check_all_sensors_ready(self):\n",
    "        self._check_joint_states_ready()\n",
    "        self._check_odom_ready()\n",
    "        rospy.logdebug(\"ALL SENSORS READY\")\n",
    "\n",
    "    def _check_joint_states_ready(self):\n",
    "        self.joints = None\n",
    "        while self.joints is None and not rospy.is_shutdown():\n",
    "            try:\n",
    "                self.joints = rospy.wait_for_message(\"/moving_cube/joint_states\", JointState, timeout=1.0)\n",
    "                rospy.logdebug(\"Current moving_cube/joint_states READY=>\" + str(self.joints))\n",
    "\n",
    "            except:\n",
    "                rospy.logerr(\"Current moving_cube/joint_states not ready yet, retrying for getting joint_states\")\n",
    "        return self.joints\n",
    "\n",
    "    def _check_odom_ready(self):\n",
    "        self.odom = None\n",
    "        while self.odom is None and not rospy.is_shutdown():\n",
    "            try:\n",
    "                self.odom = rospy.wait_for_message(\"/moving_cube/odom\", Odometry, timeout=1.0)\n",
    "                rospy.logdebug(\"Current /moving_cube/odom READY=>\" + str(self.odom))\n",
    "\n",
    "            except:\n",
    "                rospy.logerr(\"Current /moving_cube/odom not ready yet, retrying for getting odom\")\n",
    "\n",
    "        return self.odom\n",
    "        \n",
    "    def _joints_callback(self, data):\n",
    "        self.joints = data\n",
    "    \n",
    "    def _odom_callback(self, data):\n",
    "        self.odom = data\n",
    "        \n",
    "    def _check_publishers_connection(self):\n",
    "        \"\"\"\n",
    "        Checks that all the publishers are working\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        rate = rospy.Rate(10)  # 10hz\n",
    "        while self._roll_vel_pub.get_num_connections() == 0 and not rospy.is_shutdown():\n",
    "            rospy.logdebug(\"No susbribers to _roll_vel_pub yet so we wait and try again\")\n",
    "            try:\n",
    "                rate.sleep()\n",
    "            except rospy.ROSInterruptException:\n",
    "                # This is to avoid error when world is rested, time when backwards.\n",
    "                pass\n",
    "        rospy.logdebug(\"_roll_vel_pub Publisher Connected\")\n",
    "\n",
    "        rospy.logdebug(\"All Publishers READY\")\n",
    "    \n",
    "    # Methods that the TaskEnvironment will need to define here as virtual\n",
    "    # because they will be used in RobotGazeboEnv GrandParentClass and defined in the\n",
    "    # TaskEnvironment.\n",
    "    # ----------------------------\n",
    "    def _set_init_pose(self):\n",
    "        \"\"\"Sets the Robot in its init pose\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "    \n",
    "    def _init_env_variables(self):\n",
    "        \"\"\"Inits variables needed to be initialised each time we reset at the start\n",
    "        of an episode.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def _compute_reward(self, observations, done):\n",
    "        \"\"\"Calculates the reward to give based on the observations given.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def _set_action(self, action):\n",
    "        \"\"\"Applies the given action to the simulation.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def _get_obs(self):\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def _is_done(self, observations):\n",
    "        \"\"\"Checks if episode done based on observations given.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "        \n",
    "    # Methods that the TaskEnvironment will need.\n",
    "    # ----------------------------\n",
    "    def move_joints(self, roll_speed):\n",
    "        joint_speed_value = Float64()\n",
    "        joint_speed_value.data = roll_speed\n",
    "        rospy.logdebug(\"Single Disk Roll Velocity>>\" + str(joint_speed_value))\n",
    "        self._roll_vel_pub.publish(joint_speed_value)\n",
    "        self.wait_until_roll_is_in_vel(joint_speed_value.data)\n",
    "    \n",
    "    def wait_until_roll_is_in_vel(self, velocity):\n",
    "    \n",
    "        rate = rospy.Rate(10)\n",
    "        start_wait_time = rospy.get_rostime().to_sec()\n",
    "        end_wait_time = 0.0\n",
    "        epsilon = 0.1\n",
    "        v_plus = velocity + epsilon\n",
    "        v_minus = velocity - epsilon\n",
    "        while not rospy.is_shutdown():\n",
    "            joint_data = self._check_joint_states_ready()\n",
    "            roll_vel = joint_data.velocity[0]\n",
    "            rospy.logdebug(\"VEL=\" + str(roll_vel) + \", ?RANGE=[\" + str(v_minus) + \",\"+str(v_plus)+\"]\")\n",
    "            are_close = (roll_vel <= v_plus) and (roll_vel > v_minus)\n",
    "            if are_close:\n",
    "                rospy.logdebug(\"Reached Velocity!\")\n",
    "                end_wait_time = rospy.get_rostime().to_sec()\n",
    "                break\n",
    "            rospy.logdebug(\"Not there yet, keep waiting...\")\n",
    "            rate.sleep()\n",
    "        delta_time = end_wait_time- start_wait_time\n",
    "        rospy.logdebug(\"[Wait Time=\" + str(delta_time)+\"]\")\n",
    "        return delta_time\n",
    "        \n",
    "\n",
    "    def get_joints(self):\n",
    "        return self.joints\n",
    "    \n",
    "    def get_odom(self):\n",
    "        return self.odom"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p style=\"background:#407EAF;color:white;\">**END my_cube_single_disk_env.py**</p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NEXT STEP: How to create the One Disk Moving Cube Learning env for walking on the Y-axis."
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [Root]",
   "language": "python",
   "name": "Python [Root]"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}