tutorials/demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "308f4833-c7ec-4040-b724-df01d437ce44",
   "metadata": {},
   "source": [
    "[<img align=\"left\" src=\"https://colab.research.google.com/assets/colab-badge.svg\">](https://colab.research.google.com/github/pytorch/rl/blob/main/tutorials/demo.ipynb)\n",
    "\n",
    "# TorchRL Demo\n",
    "\n",
    "___\n",
    "This demo was presented at ICML 2022 on the industry demo day.\n",
    "\n",
    "It gives a good overview of TorchRL functionalities.\n",
    "\n",
    "Feel free to reach out to vmoens@fb.com or submit issues if you have questions or comments about it.\n",
    "___\n",
    "\n",
    "TorchRL is an open-source Reinforcement Learning (RL) library for PyTorch.\n",
    "\n",
    "https://github.com/pytorch/rl\n",
    "\n",
    "The PyTorch ecosystem team (Meta) has decided to invest in that library to provide a leading platform to develop RL solutions in research settings.\n",
    "\n",
    "It provides pytorch and **python-first**, low and high level **abstractions** for RL that are intended to be efficient, documented and properly tested. The code is aimed at supporting research in RL. Most of it is written in python in a highly modular way, such that researchers can easily swap components, transform them or write new ones with little effort.\n",
    "\n",
    "This repo attempts to align with the existing pytorch ecosystem libraries in that it has a dataset pillar (torchrl/envs), transforms, models, data utilities (e.g. collectors and containers), etc. TorchRL aims at having as few dependencies as possible (python standard library, numpy and pytorch). Common environment libraries (e.g. OpenAI gym) are only optional.\n",
    "\n",
    "Content:\n",
    "\n",
    "```\n",
    "torchrl\n",
    "│\n",
    "└───collectors\n",
    "│       collectors.py\n",
    "│   \n",
    "└───data\n",
    "│   │   tensor_specs.py\n",
    "│   └───postprocs\n",
    "│   │       postprocs.py\n",
    "│   └───replay_buffers\n",
    "│   │       replay_buffers.py\n",
    "│   │       storages.py\n",
    "│   └───tensordict\n",
    "│           memmap.py\n",
    "│           metatensor.py\n",
    "│           tensordict.py\n",
    "└───envs\n",
    "│   │   common.py\n",
    "│   │   env_creator.py\n",
    "│   │   gym_like.py\n",
    "│   │   vec_env.py\n",
    "│   └───libs\n",
    "│   │       dm_control.py\n",
    "│   │       gym.py\n",
    "│   └───transforms\n",
    "│           functional.py\n",
    "│           transforms.py\n",
    "└───modules\n",
    "│   └───distributions\n",
    "│   │       continuous.py\n",
    "│   │       discrete.py\n",
    "│   └───models\n",
    "│   │       models.py\n",
    "│   │       exploration.py\n",
    "│   └───tensordict_module\n",
    "│           actors.py\n",
    "│           common.py\n",
    "│           exploration.py\n",
    "│           probabilistic.py\n",
    "│           sequence.py\n",
    "└───objectives\n",
    "│   │   common.py\n",
    "│   │   ddpg.py\n",
    "│   │   dqn.py\n",
    "│   │   functional.py\n",
    "│   │   ppo.py\n",
    "│   │   redq.py\n",
    "│   │   reinforce.py\n",
    "│   │   sac.py\n",
    "│   │   utils.py\n",
    "│   └───value\n",
    "│           advantages.py\n",
    "│           functional.py\n",
    "│           pg.py\n",
    "│           returns.py\n",
    "│           utils.py\n",
    "│           vtrace.py\n",
    "└───record\n",
    "└───trainers\n",
    "    │   loggers.py\n",
    "    │   trainers.py\n",
    "    └───helpers\n",
    "            collectors.py\n",
    "            envs.py\n",
    "            losses.py\n",
    "            models.py\n",
    "            recorder.py\n",
    "            replay_buffer.py\n",
    "            trainers.py\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d38be66d-a807-4a53-956f-d53eb99cc801",
   "metadata": {},
   "source": [
    "Unlike other domains, RL is less about media than _algorithms_. As such, it is harder to make truly independent components.\n",
    "\n",
    "What TorchRL is not:\n",
    "- a collection of algorithms: we do not intend to provide SOTA implementations of RL algorithms, but we provide these algorithms only as examples of how to use the library.\n",
    "- a research framework\n",
    "\n",
    "TorchRL has very few core dependencies, mostly PyTorch and functorch. All other dependencies (gym, torchvision, wandb / tensorboard) are optional."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ce9584b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install functorch\n",
    "!pip install \"gym[classic_control]\"\n",
    "!pip install torchrl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1af6372c-de8b-4435-92b4-53f95f4e5db5",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Data\n",
    "### TensorDict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8a66580d-4ab5-4c64-a3c0-7af171603bbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchrl.data import TensorDict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "99bb0b39-fc81-4fcf-8de4-b115b1a49724",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorDict(\n",
      "    fields={\n",
      "        key 1: Tensor(torch.Size([5, 3]), dtype=torch.float32),\n",
      "        key 2: Tensor(torch.Size([5, 5, 6]), dtype=torch.bool)},\n",
      "    batch_size=torch.Size([5]),\n",
      "    device=cpu,\n",
      "    is_shared=False)\n"
     ]
    }
   ],
   "source": [
    "# Creating a TensorDict\n",
    "batch_size = 5\n",
    "tensordict = TensorDict(source={\n",
    "    \"key 1\": torch.zeros(batch_size, 3),\n",
    "    \"key 2\": torch.zeros(batch_size, 5, 6, dtype=torch.bool)\n",
    "}, batch_size = [batch_size])\n",
    "print(tensordict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8ef9a4f1-4682-41cb-b024-37770740d389",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorDict(\n",
       "    fields={\n",
       "        key 1: Tensor(torch.Size([3]), dtype=torch.float32),\n",
       "        key 2: Tensor(torch.Size([5, 6]), dtype=torch.bool)},\n",
       "    batch_size=torch.Size([]),\n",
       "    device=cpu,\n",
       "    is_shared=False)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# indexing\n",
    "tensordict[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7e0c8f4a-29c8-4194-acf8-fce33eb21d3e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# querying keys\n",
    "tensordict[\"key 1\"] is tensordict.get(\"key 1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "40aab986-6949-4b21-990c-1afa56eea914",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([2, 5]),\n",
       " tensor([[[0.],\n",
       "          [0.],\n",
       "          [0.],\n",
       "          [0.],\n",
       "          [0.]],\n",
       " \n",
       "         [[1.],\n",
       "          [1.],\n",
       "          [1.],\n",
       "          [1.],\n",
       "          [1.]]]))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Stacking tensordicts\n",
    "\n",
    "tensordict1 = TensorDict(source={\n",
    "    \"key 1\": torch.zeros(batch_size, 1),\n",
    "    \"key 2\": torch.zeros(batch_size, 5, 6, dtype=torch.bool)\n",
    "}, batch_size = [batch_size])\n",
    "\n",
    "tensordict2 = TensorDict(source={\n",
    "    \"key 1\": torch.ones(batch_size, 1),\n",
    "    \"key 2\": torch.ones(batch_size, 5, 6, dtype=torch.bool)\n",
    "}, batch_size = [batch_size])\n",
    "\n",
    "tensordict = torch.stack([tensordict1, tensordict2], 0)\n",
    "tensordict.batch_size, tensordict[\"key 1\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "117ad9cb-9eec-4e50-9f26-03fe1f58ebf2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "view(-1):  torch.Size([10]) torch.Size([10, 1])\n",
      "to device:  LazyStackedTensorDict(\n",
      "    fields={\n",
      "        key 1: Tensor(torch.Size([2, 5, 1]), dtype=torch.float32),\n",
      "        key 2: Tensor(torch.Size([2, 5, 5, 6]), dtype=torch.bool)},\n",
      "    batch_size=torch.Size([2, 5]),\n",
      "    device=cpu,\n",
      "    is_shared=False)\n",
      "share memory:  LazyStackedTensorDict(\n",
      "    fields={\n",
      "        key 1: Tensor(torch.Size([2, 5, 1]), dtype=torch.float32),\n",
      "        key 2: Tensor(torch.Size([2, 5, 5, 6]), dtype=torch.bool)},\n",
      "    batch_size=torch.Size([2, 5]),\n",
      "    device=cpu,\n",
      "    is_shared=True)\n",
      "permute(1, 0):  torch.Size([5, 2]) torch.Size([5, 2, 1])\n",
      "expand:  torch.Size([3, 2, 5]) torch.Size([3, 2, 5, 1])\n"
     ]
    }
   ],
   "source": [
    "# Other functionalities\n",
    "print(\"view(-1): \", tensordict.view(-1).batch_size, tensordict.view(-1).get(\"key 1\").shape)\n",
    "\n",
    "print(\"to device: \", tensordict.to(\"cpu\"))\n",
    "\n",
    "# print(\"pin_memory: \", tensordict.pin_memory())\n",
    "\n",
    "print(\"share memory: \", tensordict.share_memory_())\n",
    "\n",
    "print(\"permute(1, 0): \", \n",
    "      tensordict.permute(1, 0).batch_size, \n",
    "      tensordict.permute(1, 0).get(\"key 1\").shape)\n",
    "\n",
    "print(\"expand: \", \n",
    "      tensordict.expand(3, *tensordict.batch_size).batch_size, \n",
    "      tensordict.expand(3, *tensordict.batch_size).get(\"key 1\").shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bb26d8d-399f-498f-af7a-cf16597636d1",
   "metadata": {},
   "source": [
    "#### Nested tensordict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ac43369a-5bb7-45e8-afc0-8df76be1450e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorDict(\n",
       "    fields={\n",
       "        key 1: Tensor(torch.Size([5, 3]), dtype=torch.float32),\n",
       "        key 2: TensorDict(\n",
       "            fields={\n",
       "                sub-key 1: Tensor(torch.Size([5, 2, 1]), dtype=torch.float32)},\n",
       "            batch_size=torch.Size([5, 2]),\n",
       "            device=cpu,\n",
       "            is_shared=False)},\n",
       "    batch_size=torch.Size([5]),\n",
       "    device=cpu,\n",
       "    is_shared=False)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tensordict = TensorDict(source={\n",
    "    \"key 1\": torch.zeros(batch_size, 3),\n",
    "    \"key 2\": TensorDict(source={\n",
    "        \"sub-key 1\": torch.zeros(batch_size, 2, 1)\n",
    "    }, batch_size=[batch_size, 2])\n",
    "}, batch_size = [batch_size])\n",
    "tensordict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e94f6a8d-2429-45c8-9d50-abebef682836",
   "metadata": {},
   "source": [
    "### Replay buffers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f123c7f1-fa92-491d-97c8-c5d82f556003",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchrl.data import ReplayBuffer, PrioritizedReplayBuffer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ddbdd51d-a36f-4b07-b3c4-6699ca813835",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rb = ReplayBuffer(100, collate_fn=lambda x: x)\n",
    "rb.add(1)\n",
    "rb.sample(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d5254dd2-a2be-42cb-8304-a14324a51a2c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[2, 1, 2]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rb.extend([2, 3])\n",
    "rb.sample(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e607eb52-b55e-4416-93fc-a4aa400c47cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "rb = PrioritizedReplayBuffer(100, alpha=0.7, beta=1.1, collate_fn=lambda x: x)\n",
    "rb.add(1)\n",
    "rb.sample(1)\n",
    "rb.update_priority(1, 0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50f43a07-d7bb-4bfc-8c11-c5eb4ae0caf7",
   "metadata": {},
   "source": [
    "#### working with tensordicts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8ec9461d-afa4-4d35-ab04-6f66ac0e4036",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "collate_fn = torch.stack\n",
    "rb = ReplayBuffer(100, collate_fn=collate_fn)\n",
    "rb.add(TensorDict({\"a\": torch.randn(3)}, batch_size=[]))\n",
    "len(rb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "68d94666-3724-41a4-b888-e8d4bd75f853",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rb.extend(TensorDict({\"a\": torch.randn(2, 3)}, batch_size=[2]))\n",
    "len(rb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "6a808673-0682-48f0-bf27-749c2bff753c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LazyStackedTensorDict(\n",
       "    fields={\n",
       "        a: Tensor(torch.Size([10, 3]), dtype=torch.float32)},\n",
       "    batch_size=torch.Size([10]),\n",
       "    device=cpu,\n",
       "    is_shared=False)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rb.sample(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e79c254c-44d4-451b-ac76-733b38dd095e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorDict(\n",
       "    fields={\n",
       "        a: Tensor(torch.Size([2, 3]), dtype=torch.float32)},\n",
       "    batch_size=torch.Size([2]),\n",
       "    device=cpu,\n",
       "    is_shared=False)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rb.sample(2).contiguous()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "eb833a0f-01b9-4222-8239-cc8a7698b66d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorDict(\n",
       "    fields={\n",
       "        a: Tensor(torch.Size([2, 3]), dtype=torch.float32),\n",
       "        index: Tensor(torch.Size([2, 1]), dtype=torch.int32)},\n",
       "    batch_size=torch.Size([2]),\n",
       "    device=cpu,\n",
       "    is_shared=False)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "from torchrl.data import TensorDictPrioritizedReplayBuffer\n",
    "rb = TensorDictPrioritizedReplayBuffer(100, alpha=0.7, beta=1.1, priority_key=\"td_error\")\n",
    "rb.extend(TensorDict({\"a\": torch.randn(2, 3)}, batch_size=[2]))\n",
    "tensordict_sample = rb.sample(2).contiguous()\n",
    "tensordict_sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "1e029dc1-e5d0-4b7f-99fe-7f416725bdf2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1],\n",
       "        [0]], dtype=torch.int32)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tensordict_sample[\"index\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "70017242-5396-4319-905b-40e483b0f96f",
   "metadata": {},
   "outputs": [],
   "source": [
    "tensordict_sample[\"td_error\"] = torch.rand(2)\n",
    "rb.update_priority(tensordict_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "6e64849c-f1c7-42d1-8d31-6fc7bece6e30",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 0.28791671991348267\n",
      "1 0.06984968483448029\n",
      "2 0.0\n"
     ]
    }
   ],
   "source": [
    "for i, val in enumerate(rb._sum_tree):\n",
    "    print(i, val)\n",
    "    if i == len(rb):\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1a6d60d-3de1-43f9-a498-337abb98de1d",
   "metadata": {},
   "source": [
    "## Envs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "2bde26db-1880-4fcb-bd1d-f90420317b3d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n"
     ]
    }
   ],
   "source": [
    "from torchrl.envs.libs.gym import GymWrapper, GymEnv\n",
    "import gym\n",
    "\n",
    "gym_env = gym.make(\"Pendulum-v1\")\n",
    "env = GymWrapper(gym_env)\n",
    "env = GymEnv(\"Pendulum-v1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "92441513-0e7e-424c-bd92-348febfb6875",
   "metadata": {},
   "outputs": [],
   "source": [
    "tensordict = env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ba2264a4-0cad-4f68-93bb-841365e468f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorDict(\n",
       "    fields={\n",
       "        action: Tensor(torch.Size([1]), dtype=torch.float32),\n",
       "        done: Tensor(torch.Size([1]), dtype=torch.bool),\n",
       "        next_observation: Tensor(torch.Size([3]), dtype=torch.float32),\n",
       "        observation: Tensor(torch.Size([3]), dtype=torch.float32),\n",
       "        reward: Tensor(torch.Size([1]), dtype=torch.float32)},\n",
       "    batch_size=torch.Size([]),\n",
       "    device=cpu,\n",
       "    is_shared=False)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env.rand_step(tensordict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b01aebe-3517-41bf-9247-bc5f7bc44b7a",
   "metadata": {},
   "source": [
    "### changing environments config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "8296f4f9-9947-4053-a681-064eca21c2d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TensorDict(\n",
       "    fields={\n",
       "        done: Tensor(torch.Size([1]), dtype=torch.bool),\n",
       "        pixels: Tensor(torch.Size([500, 500, 3]), dtype=torch.uint8),\n",
       "        state: Tensor(torch.Size([3]), dtype=torch.float32)},\n",
       "    batch_size=torch.Size([]),\n",
       "    device=cpu,\n",
       "    is_shared=False)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env = GymEnv(\"Pendulum-v1\", frame_skip=3, from_pixels=True, pixels_only=False)\n",
    "env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "bb91ba2f-158e-4bcd-bc37-c8601da92384",
   "metadata": {},
   "outputs": [],
   "source": [
    "env.close()\n",
    "del env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "05d14d70-3307-40df-a190-3c48b767feca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n"
     ]
    }
   ],
   "source": [
    "from torchrl.envs import Compose, ObservationNorm, ToTensorImage, NoopResetEnv, TransformedEnv\n",
    "base_env = GymEnv(\"Pendulum-v1\", frame_skip=3, from_pixels=True, pixels_only=False)\n",
    "env = TransformedEnv(base_env, Compose(NoopResetEnv(3), ToTensorImage()))\n",
    "env.append_transform(ObservationNorm(keys_in=[\"next_pixels\"], loc=2, scale=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9509c09-b6a9-426a-a799-9766a195156b",
   "metadata": {},
   "source": [
    "### Transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "fb96959d-c587-4f72-af17-86171d1ad952",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchrl.envs import Compose, ObservationNorm, ToTensorImage, NoopResetEnv, TransformedEnv\n",
    "base_env = GymEnv(\"Pendulum-v1\", frame_skip=3, from_pixels=True, pixels_only=False)\n",
    "env = TransformedEnv(base_env, Compose(NoopResetEnv(3), ToTensorImage()))\n",
    "env.append_transform(ObservationNorm(keys_in=[\"next_pixels\"], loc=2, scale=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "f50a7a63-d156-4eac-94db-69395f799865",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorDict(\n",
       "    fields={\n",
       "        done: Tensor(torch.Size([1]), dtype=torch.bool),\n",
       "        pixels: Tensor(torch.Size([3, 500, 500]), dtype=torch.float32),\n",
       "        state: Tensor(torch.Size([3]), dtype=torch.float32)},\n",
       "    batch_size=torch.Size([]),\n",
       "    device=cpu,\n",
       "    is_shared=False)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "883a438f-1f78-4759-a47d-77025cc24920",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env:  TransformedEnv(env=GymEnv(env=Pendulum-v1, batch_size=torch.Size([]), device=cpu), transform=Compose(\n",
      "        NoopResetEnv(noops=3, random=True),\n",
      "        ToTensorImage(keys=['next_pixels']),\n",
      "        ObservationNorm(loc=2.0000, scale=1.0000, keys=['next_pixels'])))\n",
      "last transform parent:  TransformedEnv(env=GymEnv(env=Pendulum-v1, batch_size=torch.Size([]), device=cpu), transform=Compose(\n",
      "        NoopResetEnv(noops=3, random=True),\n",
      "        ToTensorImage(keys=['next_pixels'])))\n"
     ]
    }
   ],
   "source": [
    "print(\"env: \", env)\n",
    "print(\"last transform parent: \", env.transform[2].parent)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20f3160c-34e9-40a8-8ace-141a050111a2",
   "metadata": {},
   "source": [
    "### Vectorized environments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "64137b0a-0267-4e6f-b08b-4b947ad98823",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TensorDict(\n",
       "    fields={\n",
       "        done: Tensor(torch.Size([4, 1]), dtype=torch.bool),\n",
       "        pixels: Tensor(torch.Size([4, 3, 500, 500]), dtype=torch.float32),\n",
       "        state: Tensor(torch.Size([4, 3]), dtype=torch.float32)},\n",
       "    batch_size=torch.Size([4]),\n",
       "    device=cpu,\n",
       "    is_shared=False)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torchrl.envs import ParallelEnv\n",
    "base_env = ParallelEnv(4, lambda: GymEnv(\"Pendulum-v1\", frame_skip=3, from_pixels=True, pixels_only=False))\n",
    "env = TransformedEnv(base_env, Compose(NoopResetEnv(3), ToTensorImage()))  # applies transforms on batch of envs\n",
    "env.append_transform(ObservationNorm(keys_in=[\"next_pixels\"], loc=2, scale=1))\n",
    "env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "93fc6dd6-1ed5-412f-8f4b-4099b30b969d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NdBoundedTensorSpec(\n",
       "     shape=torch.Size([1]), space=ContinuousBox(minimum=tensor([-2.]), maximum=tensor([2.])), device=cpu, dtype=torch.float32, domain=continuous)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env.action_spec"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4b72279-24e4-486a-beca-e7c130164ed6",
   "metadata": {},
   "source": [
    "## Modules"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af6cb397-cbd6-4caf-bc1e-2f5388cd4c64",
   "metadata": {},
   "source": [
    "### Models\n",
    "#### MLP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "cb115026-b9ca-49c0-afd4-99795563e86b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MLP(\n",
      "  (0): LazyLinear(in_features=0, out_features=32, bias=True)\n",
      "  (1): ELU(alpha=1.0)\n",
      "  (2): Linear(in_features=32, out_features=64, bias=True)\n",
      "  (3): ELU(alpha=1.0)\n",
      "  (4): Linear(in_features=64, out_features=4, bias=True)\n",
      ")\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/vmoens/venv/rl/lib/python3.8/site-packages/torch/nn/modules/lazy.py:180: UserWarning: Lazy modules are a new feature under heavy development so changes to the API or functionality can happen at any moment.\n",
      "  warnings.warn('Lazy modules are a new feature under heavy development '\n"
     ]
    }
   ],
   "source": [
    "from torchrl.modules import MLP, ConvNet\n",
    "from torchrl.modules.models.utils import SquashDims\n",
    "from torch import nn\n",
    "net = MLP(num_cells=[32, 64], out_features=4, activation_class=nn.ELU)\n",
    "print(net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "f6007bbe-df30-4762-97f5-26676647a40e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10, 4])"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net(torch.randn(10, 3)).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a444a92-7b7d-42fc-a3d3-2bad7bc32880",
   "metadata": {},
   "source": [
    "#### CNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "9facfc1d-f207-43b2-8fd5-92fa6b32fb18",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ConvNet(\n",
      "  (0): LazyConv2d(0, 32, kernel_size=(8, 8), stride=(2, 2))\n",
      "  (1): ELU(alpha=1.0)\n",
      "  (2): Conv2d(32, 64, kernel_size=(4, 4), stride=(1, 1))\n",
      "  (3): ELU(alpha=1.0)\n",
      "  (4): SquashDims()\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "cnn = ConvNet(num_cells=[32, 64], kernel_sizes=[8, 4], strides=[2, 1], aggregator_class=SquashDims)\n",
    "print(cnn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "cd9d9eb1-ec44-4113-a3a1-5291f388e274",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10, 6400])"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cnn(torch.randn(10, 3, 32, 32)).shape  # last tensor is squashed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab751e1c-3fbc-4209-9339-a0dea73664e5",
   "metadata": {},
   "source": [
    "### TensorDictModules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "0144b74a-2230-4d78-9b07-ad29a4f39402",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorDict(\n",
      "    fields={\n",
      "        key 1: Tensor(torch.Size([10, 3]), dtype=torch.float32),\n",
      "        key 2: Tensor(torch.Size([10, 4]), dtype=torch.float32)},\n",
      "    batch_size=torch.Size([10]),\n",
      "    device=cpu,\n",
      "    is_shared=False)\n"
     ]
    }
   ],
   "source": [
    "from torchrl.modules import TensorDictModule\n",
    "tensordict = TensorDict({\"key 1\": torch.randn(10, 3)}, batch_size=[10])\n",
    "module = nn.Linear(3, 4)\n",
    "td_module = TensorDictModule(module, in_keys=[\"key 1\"], out_keys=[\"key 2\"])\n",
    "td_module(tensordict)\n",
    "print(tensordict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "819ab5df-d29c-44f4-8cc8-c15a2d553285",
   "metadata": {},
   "source": [
    "### Sequences of modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "e0e3e3ff-074a-4984-8e10-85fdd9d2328f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorDictSequential(\n",
      "    module=ModuleList(\n",
      "      (0): TensorDictModule(\n",
      "          module=Linear(in_features=5, out_features=3, bias=True), \n",
      "          device=cpu, \n",
      "          in_keys=['observation'], \n",
      "          out_keys=['hidden'])\n",
      "      (1): TensorDictModule(\n",
      "          module=Linear(in_features=3, out_features=4, bias=True), \n",
      "          device=cpu, \n",
      "          in_keys=['hidden'], \n",
      "          out_keys=['action'])\n",
      "      (2): TensorDictModule(\n",
      "          module=MLP(\n",
      "            (0): LazyLinear(in_features=0, out_features=4, bias=True)\n",
      "            (1): Tanh()\n",
      "            (2): Linear(in_features=4, out_features=5, bias=True)\n",
      "            (3): Tanh()\n",
      "            (4): Linear(in_features=5, out_features=1, bias=True)\n",
      "          ), \n",
      "          device=cpu, \n",
      "          in_keys=['hidden', 'action'], \n",
      "          out_keys=['value'])\n",
      "    ), \n",
      "    device=cpu, \n",
      "    in_keys=['observation'], \n",
      "    out_keys=['hidden', 'action', 'value'])\n"
     ]
    }
   ],
   "source": [
    "from torchrl.modules import TensorDictSequential\n",
    "backbone_module = nn.Linear(5, 3)\n",
    "backbone = TensorDictModule(backbone_module, in_keys=[\"observation\"], out_keys=[\"hidden\"])\n",
    "actor_module = nn.Linear(3, 4)\n",
    "actor = TensorDictModule(actor_module, in_keys=[\"hidden\"], out_keys=[\"action\"])\n",
    "value_module = MLP(out_features=1, num_cells=[4, 5])\n",
    "value = TensorDictModule(value_module, in_keys=[\"hidden\", \"action\"], out_keys=[\"value\"])\n",
    "\n",
    "sequence = TensorDictSequential(backbone, actor, value)\n",
    "print(sequence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "6e7980e8-7286-403e-82f7-8386021a6c85",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['observation'] ['hidden', 'action', 'value']\n"
     ]
    }
   ],
   "source": [
    "print(sequence.in_keys, sequence.out_keys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "b9be6bff-d543-45bc-975e-d820b9db6ce8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorDict(\n",
       "    fields={\n",
       "        action: Tensor(torch.Size([3, 4]), dtype=torch.float32),\n",
       "        hidden: Tensor(torch.Size([3, 3]), dtype=torch.float32),\n",
       "        observation: Tensor(torch.Size([3, 5]), dtype=torch.float32),\n",
       "        value: Tensor(torch.Size([3, 1]), dtype=torch.float32)},\n",
       "    batch_size=torch.Size([3]),\n",
       "    device=cpu,\n",
       "    is_shared=False)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tensordict = TensorDict(\n",
    "    {\"observation\": torch.randn(3, 5)}, [3],\n",
    ")\n",
    "backbone(tensordict)\n",
    "actor(tensordict)\n",
    "value(tensordict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "315d7d77-314e-4ffc-b9aa-e3651937dc98",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorDict(\n",
      "    fields={\n",
      "        action: Tensor(torch.Size([3, 4]), dtype=torch.float32),\n",
      "        hidden: Tensor(torch.Size([3, 3]), dtype=torch.float32),\n",
      "        observation: Tensor(torch.Size([3, 5]), dtype=torch.float32),\n",
      "        value: Tensor(torch.Size([3, 1]), dtype=torch.float32)},\n",
      "    batch_size=torch.Size([3]),\n",
      "    device=cpu,\n",
      "    is_shared=False)\n"
     ]
    }
   ],
   "source": [
    "tensordict = TensorDict(\n",
    "    {\"observation\": torch.randn(3, 5)}, [3],\n",
    ")\n",
    "sequence(tensordict)\n",
    "print(tensordict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a62cd71d-a33c-41dd-aa75-eb4cefef8c50",
   "metadata": {},
   "source": [
    "### Functional programming (ensembling / meta-RL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "f3496472-b697-4c78-9b77-972b74573884",
   "metadata": {},
   "outputs": [],
   "source": [
    "fsequence, (params, buffers) = sequence.make_functional_with_buffers()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "1577590f-5156-439f-a2f1-f8cba1fa3e78",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(list(fsequence.parameters()))  # functional modules have no parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "971618a2-9c4c-4af6-b170-082cdea4a756",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorDict(\n",
       "    fields={\n",
       "        action: Tensor(torch.Size([3, 4]), dtype=torch.float32),\n",
       "        hidden: Tensor(torch.Size([3, 3]), dtype=torch.float32),\n",
       "        observation: Tensor(torch.Size([3, 5]), dtype=torch.float32),\n",
       "        value: Tensor(torch.Size([3, 1]), dtype=torch.float32)},\n",
       "    batch_size=torch.Size([3]),\n",
       "    device=cpu,\n",
       "    is_shared=False)"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fsequence(tensordict, params=params, buffers=buffers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "ad98c6dc-918e-450a-9f3c-feb738e36d35",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorDict(\n",
      "    fields={\n",
      "        action: Tensor(torch.Size([4, 3, 4]), dtype=torch.float32),\n",
      "        hidden: Tensor(torch.Size([4, 3, 3]), dtype=torch.float32),\n",
      "        observation: Tensor(torch.Size([4, 3, 5]), dtype=torch.float32),\n",
      "        value: Tensor(torch.Size([4, 3, 1]), dtype=torch.float32)},\n",
      "    batch_size=torch.Size([4, 3]),\n",
      "    device=cpu,\n",
      "    is_shared=False)\n"
     ]
    }
   ],
   "source": [
    "params_expand = [p.expand(4, *p.shape) for p in params]\n",
    "buffers_expand = [b.expand(4, *b.shape) for b in buffers]\n",
    "tensordict_exp = fsequence(tensordict, params=params_expand, buffers=buffers, vmap=(0, 0, None))\n",
    "print(tensordict_exp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14084eb3-36e6-4729-8383-7ef4471fea5f",
   "metadata": {},
   "source": [
    "### Specialized classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "9c3f6d96-f213-4ef5-b700-133f40bf52f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.0137,  0.1524, -0.0641], grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "from torchrl.data import NdBoundedTensorSpec\n",
    "spec = NdBoundedTensorSpec(-torch.ones(3), torch.ones(3))\n",
    "base_module = nn.Linear(5, 3)\n",
    "module = TensorDictModule(module=base_module, spec=spec, in_keys=[\"obs\"], out_keys=[\"action\"], safe=True)\n",
    "tensordict = TensorDict({\"obs\": torch.randn(5)}, batch_size=[])\n",
    "module(tensordict)[\"action\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "441a1de4-e5e5-4ccf-a4a4-c7bb10e3ccc0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-1.,  1., -1.], grad_fn=<IndexPutBackward0>)"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tensordict = TensorDict({\"obs\": torch.randn(5)*100}, batch_size=[])\n",
    "module(tensordict)[\"action\"]  # safe=True projects the result within the set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "9ca25cc1-56bc-4e77-9feb-9298435042b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorDict(\n",
       "    fields={\n",
       "        action: Tensor(torch.Size([3]), dtype=torch.float32),\n",
       "        obs: Tensor(torch.Size([5]), dtype=torch.float32)},\n",
       "    batch_size=torch.Size([]),\n",
       "    device=cpu,\n",
       "    is_shared=False)"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torchrl.modules import Actor\n",
    "base_module = nn.Linear(5, 3)\n",
    "actor = Actor(base_module, in_keys=[\"obs\"])\n",
    "tensordict = TensorDict({\"obs\": torch.randn(5)}, batch_size=[])\n",
    "actor(tensordict)  # action is the default value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "0ba0507a-ff43-42d5-bd4f-c25fd006c00f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorDict(\n",
      "    fields={\n",
      "        action: Tensor(torch.Size([3, 2]), dtype=torch.float32),\n",
      "        input: Tensor(torch.Size([3, 5]), dtype=torch.float32),\n",
      "        loc: Tensor(torch.Size([3, 2]), dtype=torch.float32),\n",
      "        scale: Tensor(torch.Size([3, 2]), dtype=torch.float32)},\n",
      "    batch_size=torch.Size([3]),\n",
      "    device=cpu,\n",
      "    is_shared=False)\n"
     ]
    }
   ],
   "source": [
    "# Probabilistic modules\n",
    "from torchrl.modules import ProbabilisticTensorDictModule\n",
    "from torchrl.data import TensorDict\n",
    "from torchrl.modules import  TanhNormal, NormalParamWrapper\n",
    "td = TensorDict({\"input\": torch.randn(3, 5)}, [3,])\n",
    "net = NormalParamWrapper(nn.Linear(5, 4))  # splits the output in loc and scale\n",
    "module = TensorDictModule(net, in_keys=[\"input\"], out_keys=[\"loc\", \"scale\"])\n",
    "td_module = ProbabilisticTensorDictModule(\n",
    "   module=module,\n",
    "   dist_param_keys=[\"loc\", \"scale\"],\n",
    "   out_key_sample=[\"action\"],\n",
    "   distribution_class=TanhNormal,\n",
    "   return_log_prob=False,\n",
    ")\n",
    "td_module(td)\n",
    "print(td)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "a0a6dc50-a11c-408f-ae06-7c83795a8353",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorDict(\n",
      "    fields={\n",
      "        action: Tensor(torch.Size([3, 2]), dtype=torch.float32),\n",
      "        input: Tensor(torch.Size([3, 5]), dtype=torch.float32),\n",
      "        loc: Tensor(torch.Size([3, 2]), dtype=torch.float32),\n",
      "        sample_log_prob: Tensor(torch.Size([3, 1]), dtype=torch.float32),\n",
      "        scale: Tensor(torch.Size([3, 2]), dtype=torch.float32)},\n",
      "    batch_size=torch.Size([3]),\n",
      "    device=cpu,\n",
      "    is_shared=False)\n"
     ]
    }
   ],
   "source": [
    "# returning the log-probability\n",
    "td = TensorDict({\"input\": torch.randn(3, 5)}, [3,])\n",
    "td_module = ProbabilisticTensorDictModule(\n",
    "   module=module,\n",
    "   dist_param_keys=[\"loc\", \"scale\"],\n",
    "   out_key_sample=[\"action\"],\n",
    "   distribution_class=TanhNormal,\n",
    "   return_log_prob=True,\n",
    ")\n",
    "td_module(td)\n",
    "print(td)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "a84857c9-8a00-4526-92e4-8b6a05646bd5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "random: tensor([[ 0.8728, -0.1335],\n",
      "        [-0.9833,  0.3497],\n",
      "        [-0.6889, -0.6433]], grad_fn=<ClampBackward1>)\n",
      "mode: tensor([[-0.1131,  0.1761],\n",
      "        [-0.3425, -0.2665],\n",
      "        [ 0.2915,  0.6207]], grad_fn=<ClampBackward1>)\n",
      "mean: tensor([[-0.1131,  0.1441],\n",
      "        [-0.2375, -0.1242],\n",
      "        [ 0.1372,  0.3810]], grad_fn=<MeanBackward1>)\n"
     ]
    }
   ],
   "source": [
    "# Sampling vs mode / mean\n",
    "from torchrl.envs.utils import set_exploration_mode\n",
    "td = TensorDict({\"input\": torch.randn(3, 5)}, [3,])\n",
    "\n",
    "torch.manual_seed(0)\n",
    "with set_exploration_mode(\"random\"):\n",
    "    td_module(td)\n",
    "    print(\"random:\", td[\"action\"])\n",
    "    \n",
    "with set_exploration_mode(\"mode\"):\n",
    "    td_module(td)\n",
    "    print(\"mode:\", td[\"action\"])\n",
    "\n",
    "with set_exploration_mode(\"mean\"):\n",
    "    td_module(td)\n",
    "    print(\"mean:\", td[\"action\"])\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "141a232e-1472-4b7e-9d88-dfd0a19b8adf",
   "metadata": {},
   "source": [
    "## Using environments and modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "384a8372-3096-4897-b03d-af638b17e452",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n",
      "total steps: 99\n",
      "TensorDict(\n",
      "    fields={\n",
      "        action: Tensor(torch.Size([100, 1]), dtype=torch.float32),\n",
      "        done: Tensor(torch.Size([100, 1]), dtype=torch.bool),\n",
      "        next_observation: Tensor(torch.Size([100, 3]), dtype=torch.float32),\n",
      "        observation: Tensor(torch.Size([100, 3]), dtype=torch.float32),\n",
      "        reward: Tensor(torch.Size([100, 1]), dtype=torch.float32)},\n",
      "    batch_size=torch.Size([100]),\n",
      "    device=cpu,\n",
      "    is_shared=False)\n"
     ]
    }
   ],
   "source": [
    "from torchrl.envs.utils import step_mdp\n",
    "env = GymEnv(\"Pendulum-v1\")\n",
    "\n",
    "action_spec = env.action_spec\n",
    "actor_module = nn.Linear(3, 1)\n",
    "actor = TensorDictModule(actor_module, spec=action_spec, in_keys=[\"observation\"], out_keys=[\"action\"])\n",
    "\n",
    "torch.manual_seed(0)\n",
    "env.set_seed(0)\n",
    "\n",
    "max_steps = 100\n",
    "tensordict = env.reset()\n",
    "tensordicts = TensorDict({}, [max_steps])\n",
    "for i in range(max_steps):\n",
    "    actor(tensordict)\n",
    "    tensordicts[i] = env.step(tensordict)\n",
    "    tensordict = step_mdp(tensordict)  # roughly equivalent to obs = next_obs\n",
    "    if env.is_done:\n",
    "        break\n",
    "\n",
    "tensordicts_prealloc = tensordicts.clone()\n",
    "print(\"total steps:\", i)\n",
    "print(tensordicts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "71a2f7e7-815d-4e1c-bd8c-4b4942f3de7d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total steps: 99\n",
      "LazyStackedTensorDict(\n",
      "    fields={\n",
      "        action: Tensor(torch.Size([100, 1]), dtype=torch.float32),\n",
      "        done: Tensor(torch.Size([100, 1]), dtype=torch.bool),\n",
      "        next_observation: Tensor(torch.Size([100, 3]), dtype=torch.float32),\n",
      "        observation: Tensor(torch.Size([100, 3]), dtype=torch.float32),\n",
      "        reward: Tensor(torch.Size([100, 1]), dtype=torch.float32)},\n",
      "    batch_size=torch.Size([100]),\n",
      "    device=cpu,\n",
      "    is_shared=False)\n"
     ]
    }
   ],
   "source": [
    "# equivalent\n",
    "torch.manual_seed(0)\n",
    "env.set_seed(0)\n",
    "\n",
    "max_steps = 100\n",
    "tensordict = env.reset()\n",
    "tensordicts = []\n",
    "for i in range(max_steps):\n",
    "    actor(tensordict)\n",
    "    tensordicts.append(env.step(tensordict))\n",
    "    tensordict = step_mdp(tensordict)  # roughly equivalent to obs = next_obs\n",
    "    if env.is_done:\n",
    "        break\n",
    "tensordicts_stack = torch.stack(tensordicts, 0)\n",
    "print(\"total steps:\", i)\n",
    "print(tensordicts_stack)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "5380a357-dcb9-43a8-8a2e-f4be939db91f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(tensordicts_stack == tensordicts_prealloc).all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "ac59466c-1e39-4ecd-a840-72d5ec204b2d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorDict(\n",
       "    fields={\n",
       "        action: Tensor(torch.Size([100, 1]), dtype=torch.float32),\n",
       "        done: Tensor(torch.Size([100, 1]), dtype=torch.bool),\n",
       "        next_observation: Tensor(torch.Size([100, 3]), dtype=torch.float32),\n",
       "        observation: Tensor(torch.Size([100, 3]), dtype=torch.float32),\n",
       "        reward: Tensor(torch.Size([100, 1]), dtype=torch.float32)},\n",
       "    batch_size=torch.Size([100]),\n",
       "    device=cpu,\n",
       "    is_shared=False)"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# helper\n",
    "torch.manual_seed(0)\n",
    "env.set_seed(0)\n",
    "tensordict_rollout = env.rollout(policy=actor, max_steps=max_steps)\n",
    "tensordict_rollout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "7c7d8600-ecbb-4a55-b266-ed929f5d38c8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(tensordict_rollout == tensordicts_prealloc).all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9f8ef53-4c35-44fe-8763-792d9c237440",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "a0ae640d-777a-4aed-9c1d-0638d933afc9",
   "metadata": {},
   "source": [
    "## Collectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "02cfd1d3-150b-4430-8392-f8a629beb42d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n"
     ]
    }
   ],
   "source": [
    "from torchrl.envs import ParallelEnv, EnvCreator\n",
    "from torchrl.envs.libs.gym import GymEnv\n",
    "from torchrl.modules import TensorDictModule\n",
    "from torchrl.collectors import MultiSyncDataCollector, MultiaSyncDataCollector\n",
    "from torch import nn\n",
    "\n",
    "# EnvCreator makes sure that we can send a lambda function from process to process\n",
    "parallel_env = ParallelEnv(3, EnvCreator(lambda: GymEnv(\"Pendulum-v1\")))\n",
    "create_env_fn=[parallel_env, parallel_env]\n",
    "\n",
    "actor_module = nn.Linear(3, 1)\n",
    "actor = TensorDictModule(actor_module, in_keys=[\"observation\"], out_keys=[\"action\"])\n",
    "\n",
    "# Sync data collector\n",
    "devices = [\"cpu\", \"cpu\"]\n",
    "\n",
    "collector = MultiSyncDataCollector(\n",
    "    create_env_fn=create_env_fn,  # either a list of functions or a ParallelEnv\n",
    "    policy=actor,\n",
    "    total_frames=240,\n",
    "    max_frames_per_traj=-1,  # envs are terminating, we don't need to stop them early \n",
    "    frames_per_batch=60,  # we want 60 frames at a time (we have 3 envs per sub-collector)\n",
    "    passing_devices=devices,  # len must match len of env created\n",
    "    devices=devices,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "fe6091f2-2b33-4834-b437-fb8860b166f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorDict(\n",
      "    fields={\n",
      "        action: Tensor(torch.Size([6, 10, 1]), dtype=torch.float32),\n",
      "        done: Tensor(torch.Size([6, 10, 1]), dtype=torch.bool),\n",
      "        mask: Tensor(torch.Size([6, 10, 1]), dtype=torch.bool),\n",
      "        next_observation: Tensor(torch.Size([6, 10, 3]), dtype=torch.float32),\n",
      "        observation: Tensor(torch.Size([6, 10, 3]), dtype=torch.float32),\n",
      "        reward: Tensor(torch.Size([6, 10, 1]), dtype=torch.float32),\n",
      "        step_count: Tensor(torch.Size([6, 10, 1]), dtype=torch.int32),\n",
      "        traj_ids: Tensor(torch.Size([6, 10, 1]), dtype=torch.int64)},\n",
      "    batch_size=torch.Size([6, 10]),\n",
      "    device=cpu,\n",
      "    is_shared=False)\n",
      "3\n"
     ]
    }
   ],
   "source": [
    "for i, d in enumerate(collector):\n",
    "    if i == 0:\n",
    "        print(d)  # trajectories are split automatically in [6 workers x 10 steps]\n",
    "    collector.update_policy_weights_()  # make sure that our policies have the latest weights if working on multiple devices\n",
    "print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "b6a2e699-0d13-406e-84a4-62caf236f4ec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n",
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n",
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n",
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n",
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n",
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n",
      "TensorDict(\n",
      "    fields={\n",
      "        action: Tensor(torch.Size([3, 20, 1]), dtype=torch.float32),\n",
      "        done: Tensor(torch.Size([3, 20, 1]), dtype=torch.bool),\n",
      "        mask: Tensor(torch.Size([3, 20, 1]), dtype=torch.bool),\n",
      "        next_observation: Tensor(torch.Size([3, 20, 3]), dtype=torch.float32),\n",
      "        observation: Tensor(torch.Size([3, 20, 3]), dtype=torch.float32),\n",
      "        reward: Tensor(torch.Size([3, 20, 1]), dtype=torch.float32),\n",
      "        step_count: Tensor(torch.Size([3, 20, 1]), dtype=torch.int32),\n",
      "        traj_ids: Tensor(torch.Size([3, 20, 1]), dtype=torch.int64)},\n",
      "    batch_size=torch.Size([3, 20]),\n",
      "    device=cpu,\n",
      "    is_shared=False)\n",
      "3\n",
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n",
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n",
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n",
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n",
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n",
      "Discarding frameskip arg. This will be taken care of by TorchRL env wrapper.\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# async data collector: keeps working while you update your model\n",
    "collector = MultiaSyncDataCollector(\n",
    "    create_env_fn=create_env_fn,  # either a list of functions or a ParallelEnv\n",
    "    policy=actor,\n",
    "    total_frames=240,\n",
    "    max_frames_per_traj=-1,  # envs are terminating, we don't need to stop them early \n",
    "    frames_per_batch=60,  # we want 60 frames at a time (we have 3 envs per sub-collector)\n",
    "    passing_devices=devices,  # len must match len of env created\n",
    "    devices=devices,\n",
    ")\n",
    "\n",
    "for i, d in enumerate(collector):\n",
    "    if i == 0:\n",
    "        print(d)  # trajectories are split automatically in [6 workers x 10 steps]\n",
    "    collector.update_policy_weights_()  # make sure that our policies have the latest weights if working on multiple devices\n",
    "print(i)\n",
    "del collector"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cb3140a-d335-4a03-835a-0feba8b2581c",
   "metadata": {},
   "source": [
    "## Objectives"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "39794b46-f82d-4cd0-9121-d8fee34d352d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TorchRL delivers meta-RL compatible loss functions\n",
    "# Disclaimer: This APi may change in the future\n",
    "\n",
    "from torchrl.objectives import DDPGLoss\n",
    "from torchrl.data import TensorDict\n",
    "from torchrl.modules import TensorDictModule\n",
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "actor_module = nn.Linear(3, 1)\n",
    "actor = TensorDictModule(actor_module, in_keys=[\"observation\"], out_keys=[\"action\"])\n",
    "\n",
    "class ConcatModule(nn.Linear):\n",
    "    def forward(self, obs, action):\n",
    "        return super().forward(torch.cat([obs, action], -1))\n",
    "\n",
    "value_module = ConcatModule(4, 1)\n",
    "value = TensorDictModule(value_module, in_keys=[\"observation\", \"action\"], out_keys=[\"state_action_value\"])\n",
    "\n",
    "loss_fn = DDPGLoss(actor, value, gamma=0.99)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "78b5e1ea-eed0-48d7-b979-4e88efc4ff67",
   "metadata": {},
   "outputs": [],
   "source": [
    "tensordict = TensorDict({\n",
    "    \"observation\": torch.randn(10, 3), \n",
    "    \"next_observation\": torch.randn(10, 3),\n",
    "    \"reward\": torch.randn(10, 1),\n",
    "    \"action\": torch.randn(10, 1),\n",
    "    \"done\": torch.zeros(10, 1, dtype=torch.bool),\n",
    "}, batch_size=[10])\n",
    "loss_td = loss_fn(tensordict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "ca1ac32e-f948-432b-a40a-ff5927758377",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorDict(\n",
       "    fields={\n",
       "        loss_actor: Tensor(torch.Size([1]), dtype=torch.float32),\n",
       "        loss_value: Tensor(torch.Size([1]), dtype=torch.float32),\n",
       "        pred_value: Tensor(torch.Size([1]), dtype=torch.float32),\n",
       "        pred_value_max: Tensor(torch.Size([1]), dtype=torch.float32),\n",
       "        target_value: Tensor(torch.Size([1]), dtype=torch.float32),\n",
       "        target_value_max: Tensor(torch.Size([1]), dtype=torch.float32)},\n",
       "    batch_size=torch.Size([]),\n",
       "    device=cpu,\n",
       "    is_shared=False)"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss_td"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "eeacc666-42ca-4cee-9f23-c3082280dc47",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorDict(\n",
       "    fields={\n",
       "        action: Tensor(torch.Size([10, 1]), dtype=torch.float32),\n",
       "        done: Tensor(torch.Size([10, 1]), dtype=torch.bool),\n",
       "        next_observation: Tensor(torch.Size([10, 3]), dtype=torch.float32),\n",
       "        observation: Tensor(torch.Size([10, 3]), dtype=torch.float32),\n",
       "        reward: Tensor(torch.Size([10, 1]), dtype=torch.float32),\n",
       "        td_error: Tensor(torch.Size([10, 1]), dtype=torch.float32)},\n",
       "    batch_size=torch.Size([10]),\n",
       "    device=cpu,\n",
       "    is_shared=False)"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tensordict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3342145c-5a99-42b5-9564-f80f9cd14d41",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "500c3fbc-c6e6-448f-ba1a-8cb7916a12a0",
   "metadata": {},
   "source": [
    "## State of the library\n",
    "\n",
    "TorchRL is currently an **alpha-release**: there may be bugs and there is no guarantee about BC-breaking changes.\n",
    "We should be able to move to a beta-release by the end of the year. Our roadmap to get there comprises:\n",
    "- Distributed solutions\n",
    "- Offline RL\n",
    "- Greater support for meta-RL\n",
    "- Multi-task and hierarchical RL\n",
    "\n",
    "## Contributing:\n",
    "We are actively looking for contributors and early users. If you're working in RL (or just curious), try it! Give us feedback: what will make the success of TorchRL is how well it covers researchers needs. To do that, we need their input! Since the library is nascent, it is a great time for you to shape it the way you want!\n",
    "\n",
    "## Installing the library\n",
    "The library is on PyPI: \n",
    "```\n",
    "pip install torchrl\n",
    "```"
   ]
  }
 ],
 "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}