docs/source/notebooks/plangym_introduction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is an introductory tutorial to the main features of plangym."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with states"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `reset` and `step` return the environment state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The main difference with the `gym` API is that environment state is considered as important as observations, rewards and terminal flags. This is why plangym incorporates them to the tuples that the environment returns after calling `step` and `reset`:\n",
    "\n",
    "- The `reset` method will return a tuple of (state, observation) unless you pass `return_state=False` as an argument.\n",
    "\n",
    "- When `step` is called passing the environment state as an argument it will return a tuple containing `(state, obs, reward, end, info)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import plangym\n",
    "\n",
    "env = plangym.make(\"CartPole-v0\")\n",
    "action = env.action_space.sample()\n",
    "\n",
    "state, obs = env.reset()\n",
    "state, obs, reward, end, info = env.step(action, state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, if you don't provide the environment state when calling `step`, the returned tuple will match the standard `gym` interface:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = plangym.make(\"CartPole-v0\")\n",
    "action = env.action_space.sample()\n",
    "\n",
    "obs = env.reset(return_state=False)\n",
    "obs, reward, end, info = env.step(action)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Accessing and modifying the environment state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can get a copy of the environment's state calling `env.get_state()`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.03145539,  0.17749025,  0.01348916, -0.25611924])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "state = env.get_state()\n",
    "state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And set the environment state using `env.set_state(state)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.set_state(state)\n",
    "assert (state == env.get_state()).all()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step vectorization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All plangym environments offer a `step_batch` method that allows vectorized steps of batches of states and actions. \n",
    "\n",
    "Calling `step_batch` with a list of states and actions will return a tuple of lists containing the step data for each of the states and actions provided."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(list, list)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "states = [state.copy() for _ in range(10)]\n",
    "actions = [env.action_space.sample() for _ in range(10)]\n",
    "\n",
    "data = env.step_batch(states=states, actions=actions)\n",
    "new_states, observs, rewards, ends, infos = data\n",
    "type(new_states), type(observs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parallel step vectorization using multiprocessing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Passing the argument `n_workers` to `plangym.make` will return an environment that steps a batch of actions and states in parallel using multiprocessing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(plangym.parallel.ParallelEnv, list, list)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env = plangym.make(\"CartPole-v0\", n_workers=2)\n",
    "states = [state.copy() for _ in range(10)]\n",
    "actions = [env.action_space.sample() for _ in range(10)]\n",
    "\n",
    "data = env.step_batch(states=states, actions=actions)\n",
    "new_states, observs, rewards, ends, infos = data\n",
    "type(env), type(new_states), type(observs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step vectorization using ray"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is possible to use ray actors to step the environment in parallel when calling `step_batch` by passing `ray=True` to `plangym.make`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-12-13 10:01:47,772\tINFO services.py:1247 -- View the Ray dashboard at \u001b[1m\u001b[32mhttp://127.0.0.1:8265\u001b[39m\u001b[22m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(plangym.ray.RayEnv, list, list)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ray\n",
    "ray.init()\n",
    "\n",
    "env = plangym.make(\"CartPole-v0\", n_workers=2, ray=True)\n",
    "states = [state.copy() for _ in range(10)]\n",
    "actions = [env.action_space.sample() for _ in range(10)]\n",
    "\n",
    "data = env.step_batch(states=states, actions=actions)\n",
    "new_states, observs, rewards, ends, infos = data\n",
    "type(env), type(new_states), type(observs)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}