tutorial_Creating_ChasingVecEnv.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "name": "tutorial_Creating_ChasingVecEnv.ipynb",
   "provenance": [],
   "authorship_tag": "ABX9TyO8laSVRpW52hqbf2o1VXta",
   "include_colab_link": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "view-in-github",
    "colab_type": "text"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/AI4Finance-Foundation/ElegantRL/blob/master/tutorial_Creating_ChasingVecEnv.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "8g5P1KyKLer3",
    "outputId": "3e5b00f0-17b4-4fa3-e4c5-6b7bc25e3c3b"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Env num: 4096\n",
      "Steps: tensor([1., 1., 1.,  ..., 1., 1., 1.], device='cuda:0')\n",
      "Steps: tensor([2., 2., 2.,  ..., 2., 2., 2.], device='cuda:0')\n",
      "Steps: tensor([3., 3., 3.,  ..., 3., 3., 3.], device='cuda:0')\n",
      "Steps: tensor([4., 4., 4.,  ..., 4., 4., 4.], device='cuda:0')\n",
      "Steps: tensor([5., 5., 5.,  ..., 5., 5., 5.], device='cuda:0')\n",
      "Steps: tensor([6., 6., 6.,  ..., 6., 6., 6.], device='cuda:0')\n",
      "Steps: tensor([7., 7., 7.,  ..., 7., 7., 7.], device='cuda:0')\n",
      "Steps: tensor([8., 8., 8.,  ..., 8., 8., 8.], device='cuda:0')\n",
      "Steps: tensor([9., 9., 9.,  ..., 9., 9., 9.], device='cuda:0')\n",
      "Steps: tensor([10., 10., 10.,  ..., 10., 10., 10.], device='cuda:0')\n",
      "Steps: tensor([11., 11., 11.,  ..., 11., 11., 11.], device='cuda:0')\n",
      "Steps: tensor([12., 12., 12.,  ..., 12., 12., 12.], device='cuda:0')\n",
      "Steps: tensor([13., 13., 13.,  ..., 13., 13., 13.], device='cuda:0')\n",
      "Steps: tensor([14., 14., 14.,  ..., 14., 14., 14.], device='cuda:0')\n",
      "Steps: tensor([15., 15., 15.,  ..., 15., 15., 15.], device='cuda:0')\n",
      "Steps: tensor([16., 16., 16.,  ..., 16., 16., 16.], device='cuda:0')\n",
      "Steps: tensor([17., 17., 17.,  ..., 17., 17., 17.], device='cuda:0')\n",
      "Steps: tensor([18., 18., 18.,  ..., 18., 18., 18.], device='cuda:0')\n",
      "Steps: tensor([19., 19., 19.,  ..., 19., 19., 19.], device='cuda:0')\n",
      "Steps: tensor([20., 20., 20.,  ..., 20., 20., 20.], device='cuda:0')\n",
      "Steps: tensor([21., 21., 21.,  ..., 21., 21., 21.], device='cuda:0')\n",
      "Steps: tensor([22., 22., 22.,  ..., 22., 22., 22.], device='cuda:0')\n",
      "Steps: tensor([23., 23., 23.,  ..., 23., 23., 23.], device='cuda:0')\n",
      "Steps: tensor([24., 24., 24.,  ..., 24., 24., 24.], device='cuda:0')\n",
      "Steps: tensor([25., 25., 25.,  ..., 25., 25., 25.], device='cuda:0')\n",
      "Steps: tensor([26., 26., 26.,  ..., 26., 26., 26.], device='cuda:0')\n",
      "Steps: tensor([27., 27., 27.,  ..., 27., 27., 27.], device='cuda:0')\n",
      "Steps: tensor([28., 28., 28.,  ..., 28., 28., 28.], device='cuda:0')\n",
      "Steps: tensor([29., 29., 29.,  ..., 29., 29., 29.], device='cuda:0')\n",
      "Steps: tensor([30., 30., 30.,  ..., 30., 30., 30.], device='cuda:0')\n",
      "Steps: tensor([31., 31., 31.,  ..., 31., 31., 31.], device='cuda:0')\n",
      "Steps: tensor([32., 32., 32.,  ..., 32., 32., 32.], device='cuda:0')\n",
      "Steps: tensor([33., 33., 33.,  ..., 33., 33., 33.], device='cuda:0')\n",
      "Steps: tensor([34., 34., 34.,  ..., 34., 34., 34.], device='cuda:0')\n",
      "Steps: tensor([35., 35., 35.,  ..., 35., 35., 35.], device='cuda:0')\n",
      "Steps: tensor([36., 36., 36.,  ..., 36., 36., 36.], device='cuda:0')\n",
      "Steps: tensor([37., 37., 37.,  ..., 37., 37., 37.], device='cuda:0')\n",
      "Steps: tensor([38., 38., 38.,  ..., 38., 38., 38.], device='cuda:0')\n",
      "Steps: tensor([39., 39., 39.,  ..., 39., 39., 39.], device='cuda:0')\n",
      "Steps: tensor([40., 40., 40.,  ..., 40., 40., 40.], device='cuda:0')\n",
      "Steps: tensor([41., 41., 41.,  ..., 41., 41., 41.], device='cuda:0')\n",
      "Steps: tensor([42., 42., 42.,  ..., 42., 42., 42.], device='cuda:0')\n",
      "Steps: tensor([43., 43., 43.,  ..., 43., 43., 43.], device='cuda:0')\n",
      "Steps: tensor([44., 44., 44.,  ..., 44., 44., 44.], device='cuda:0')\n",
      "Steps: tensor([45., 45., 45.,  ..., 45., 45., 45.], device='cuda:0')\n",
      "Steps: tensor([46., 46., 46.,  ..., 46., 46., 46.], device='cuda:0')\n",
      "Steps: tensor([47., 47., 47.,  ..., 47., 47., 47.], device='cuda:0')\n",
      "Steps: tensor([48., 48., 48.,  ..., 48., 48., 48.], device='cuda:0')\n",
      "Steps: tensor([49., 49., 49.,  ..., 49., 49., 49.], device='cuda:0')\n",
      "Steps: tensor([50., 50., 50.,  ..., 50., 50., 50.], device='cuda:0')\n",
      "Steps: tensor([51., 51., 51.,  ..., 51., 51., 51.], device='cuda:0')\n",
      "Steps: tensor([52., 52., 52.,  ..., 52., 52., 52.], device='cuda:0')\n",
      "Steps: tensor([53., 53., 53.,  ..., 53., 53., 53.], device='cuda:0')\n",
      "Steps: tensor([54., 54., 54.,  ..., 54., 54., 54.], device='cuda:0')\n",
      "Steps: tensor([55., 55., 55.,  ..., 55., 55., 55.], device='cuda:0')\n",
      "Steps: tensor([56., 56., 56.,  ..., 56., 56., 56.], device='cuda:0')\n",
      "Steps: tensor([57., 57., 57.,  ..., 57., 57., 57.], device='cuda:0')\n",
      "Steps: tensor([58., 58., 58.,  ..., 58., 58., 58.], device='cuda:0')\n",
      "Steps: tensor([59., 59., 59.,  ..., 59., 59., 59.], device='cuda:0')\n",
      "Steps: tensor([60., 60., 60.,  ..., 60., 60., 60.], device='cuda:0')\n",
      "Steps: tensor([61., 61., 61.,  ..., 61., 61., 61.], device='cuda:0')\n",
      "Steps: tensor([62., 62., 62.,  ..., 62., 62., 62.], device='cuda:0')\n",
      "Steps: tensor([63., 63., 63.,  ..., 63., 63., 63.], device='cuda:0')\n",
      "Steps: tensor([64., 64., 64.,  ..., 64., 64., 64.], device='cuda:0')\n",
      "Steps: tensor([65., 65., 65.,  ..., 65., 65., 65.], device='cuda:0')\n",
      "Steps: tensor([66., 66., 66.,  ..., 66., 66., 66.], device='cuda:0')\n",
      "Steps: tensor([67., 67., 67.,  ..., 67., 67., 67.], device='cuda:0')\n",
      "Steps: tensor([68., 68., 68.,  ..., 68., 68., 68.], device='cuda:0')\n",
      "Steps: tensor([69., 69., 69.,  ..., 69., 69., 69.], device='cuda:0')\n",
      "Steps: tensor([70., 70., 70.,  ..., 70., 70., 70.], device='cuda:0')\n",
      "Steps: tensor([71., 71., 71.,  ..., 71., 71., 71.], device='cuda:0')\n",
      "Steps: tensor([72., 72., 72.,  ..., 72., 72., 72.], device='cuda:0')\n",
      "Steps: tensor([73., 73., 73.,  ..., 73., 73., 73.], device='cuda:0')\n",
      "Steps: tensor([74., 74., 74.,  ..., 74., 74., 74.], device='cuda:0')\n",
      "Steps: tensor([75., 75., 75.,  ..., 75., 75., 75.], device='cuda:0')\n",
      "Steps: tensor([76., 76., 76.,  ..., 76., 76., 76.], device='cuda:0')\n",
      "Steps: tensor([77., 77., 77.,  ..., 77., 77., 77.], device='cuda:0')\n",
      "Steps: tensor([78., 78., 78.,  ..., 78., 78., 78.], device='cuda:0')\n",
      "Steps: tensor([79., 79., 79.,  ..., 79., 79., 79.], device='cuda:0')\n",
      "Steps: tensor([80., 80., 80.,  ..., 80., 80., 80.], device='cuda:0')\n",
      "Steps: tensor([81., 81., 81.,  ..., 81., 81., 81.], device='cuda:0')\n",
      "Steps: tensor([82., 82., 82.,  ..., 82., 82., 82.], device='cuda:0')\n",
      "Steps: tensor([83., 83., 83.,  ..., 83., 83., 83.], device='cuda:0')\n",
      "Steps: tensor([84., 84., 84.,  ..., 84., 84., 84.], device='cuda:0')\n",
      "Steps: tensor([85., 85., 85.,  ..., 85., 85., 85.], device='cuda:0')\n",
      "Steps: tensor([86., 86., 86.,  ..., 86., 86., 86.], device='cuda:0')\n",
      "Steps: tensor([87., 87., 87.,  ..., 87., 87., 87.], device='cuda:0')\n",
      "Steps: tensor([88., 88., 88.,  ..., 88., 88., 88.], device='cuda:0')\n",
      "Steps: tensor([89., 89., 89.,  ..., 89., 89., 89.], device='cuda:0')\n",
      "Steps: tensor([90., 90., 90.,  ..., 90., 90., 90.], device='cuda:0')\n",
      "Steps: tensor([91., 91., 91.,  ..., 91., 91., 91.], device='cuda:0')\n",
      "Steps: tensor([92., 92., 92.,  ..., 92., 92., 92.], device='cuda:0')\n",
      "Steps: tensor([93., 93., 93.,  ..., 93., 93., 93.], device='cuda:0')\n",
      "Steps: tensor([94., 94., 94.,  ..., 94., 94., 94.], device='cuda:0')\n",
      "Steps: tensor([95., 95., 95.,  ..., 95., 95., 95.], device='cuda:0')\n",
      "Steps: tensor([96., 96., 96.,  ..., 96., 96., 96.], device='cuda:0')\n",
      "Steps: tensor([97., 97., 97.,  ..., 97., 97., 97.], device='cuda:0')\n",
      "Steps: tensor([98., 98., 98.,  ..., 98., 98., 98.], device='cuda:0')\n",
      "Steps: tensor([99., 99., 99.,  ..., 99., 99., 99.], device='cuda:0')\n",
      "Steps: tensor([100., 100., 100.,  ..., 100., 100., 100.], device='cuda:0')\n",
      "Steps: tensor([101., 101., 101.,  ..., 101., 101., 101.], device='cuda:0')\n",
      "Steps: tensor([102., 102., 102.,  ..., 102., 102., 102.], device='cuda:0')\n",
      "Steps: tensor([103., 103., 103.,  ..., 103., 103., 103.], device='cuda:0')\n",
      "Steps: tensor([104., 104., 104.,  ..., 104., 104., 104.], device='cuda:0')\n",
      "Steps: tensor([105., 105., 105.,  ..., 105., 105., 105.], device='cuda:0')\n",
      "Steps: tensor([106., 106., 106.,  ..., 106., 106., 106.], device='cuda:0')\n",
      "Steps: tensor([107., 107., 107.,  ..., 107., 107., 107.], device='cuda:0')\n",
      "Steps: tensor([108., 108., 108.,  ..., 108., 108., 108.], device='cuda:0')\n",
      "Steps: tensor([109., 109., 109.,  ..., 109., 109., 109.], device='cuda:0')\n",
      "Steps: tensor([110., 110., 110.,  ..., 110., 110., 110.], device='cuda:0')\n",
      "Steps: tensor([111., 111., 111.,  ..., 111., 111., 111.], device='cuda:0')\n",
      "Steps: tensor([112., 112., 112.,  ..., 112., 112., 112.], device='cuda:0')\n",
      "Steps: tensor([113., 113., 113.,  ..., 113., 113., 113.], device='cuda:0')\n",
      "Steps: tensor([114., 114., 114.,  ..., 114., 114., 114.], device='cuda:0')\n",
      "Steps: tensor([115., 115., 115.,  ..., 115., 115., 115.], device='cuda:0')\n",
      "Steps: tensor([116., 116., 116.,  ..., 116., 116., 116.], device='cuda:0')\n",
      "Steps: tensor([117., 117., 117.,  ..., 117., 117., 117.], device='cuda:0')\n",
      "Steps: tensor([118., 118., 118.,  ..., 118., 118., 118.], device='cuda:0')\n",
      "Steps: tensor([119., 119., 119.,  ..., 119., 119., 119.], device='cuda:0')\n",
      "Steps: tensor([120., 120., 120.,  ..., 120., 120., 120.], device='cuda:0')\n",
      "Steps: tensor([121., 121., 121.,  ..., 121., 121., 121.], device='cuda:0')\n",
      "Steps: tensor([122., 122., 122.,  ..., 122., 122., 122.], device='cuda:0')\n",
      "Steps: tensor([123., 123., 123.,  ..., 123., 123., 123.], device='cuda:0')\n",
      "Steps: tensor([124., 124., 124.,  ..., 124., 124., 124.], device='cuda:0')\n",
      "Steps: tensor([125., 125., 125.,  ..., 125., 125., 125.], device='cuda:0')\n",
      "Steps: tensor([126., 126., 126.,  ..., 126., 126., 126.], device='cuda:0')\n",
      "Steps: tensor([127., 127., 127.,  ..., 127., 127., 127.], device='cuda:0')\n",
      "Steps: tensor([128., 128., 128.,  ..., 128., 128., 128.], device='cuda:0')\n",
      "Steps: tensor([129., 129., 129.,  ..., 129., 129., 129.], device='cuda:0')\n",
      "Steps: tensor([130., 130., 130.,  ..., 130., 130., 130.], device='cuda:0')\n",
      "Steps: tensor([131., 131., 131.,  ..., 131., 131., 131.], device='cuda:0')\n",
      "Steps: tensor([132., 132., 132.,  ..., 132., 132., 132.], device='cuda:0')\n",
      "Steps: tensor([133., 133., 133.,  ..., 133., 133., 133.], device='cuda:0')\n",
      "Steps: tensor([134., 134., 134.,  ..., 134., 134., 134.], device='cuda:0')\n",
      "Steps: tensor([135., 135., 135.,  ..., 135., 135., 135.], device='cuda:0')\n",
      "Steps: tensor([136., 136., 136.,  ..., 136., 136., 136.], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import numpy.random as rd\n",
    "\n",
    "\n",
    "TargetReturnDict = {\n",
    "    2: 5.5,\n",
    "    3: 3.5,\n",
    "    4: 2.5,\n",
    "    8: -1.5,  # -1.37\n",
    "}\n",
    "\n",
    "\n",
    "class ChasingVecEnv:\n",
    "    def __init__(self, dim=2, env_num=32, device_id=0):\n",
    "        self.dim = dim\n",
    "        self.init_distance = 8.0\n",
    "\n",
    "        # reset\n",
    "        self.p0s = None  # position\n",
    "        self.v0s = None  # velocity\n",
    "        self.p1s = None\n",
    "        self.v1s = None\n",
    "\n",
    "        self.distances = None  # a tensor of distance between point0 and point1\n",
    "        self.cur_steps = None  # a tensor of current step number\n",
    "        # env.step() is a function, so I can't name it `steps`\n",
    "\n",
    "        \"\"\"env info\"\"\"\n",
    "        self.env_name = \"ChasingVecEnv\"\n",
    "        self.state_dim = self.dim * 4\n",
    "        self.action_dim = self.dim\n",
    "        self.max_step = 2**10\n",
    "        self.if_discrete = False\n",
    "        self.target_return = TargetReturnDict[dim]\n",
    "\n",
    "        self.env_num = env_num\n",
    "        self.device = torch.device(\"cpu\" if device_id == -1 else f\"cuda:{device_id}\")\n",
    "\n",
    "    def reset(self):\n",
    "        self.p0s = torch.zeros(\n",
    "            (self.env_num, self.dim), dtype=torch.float32, device=self.device\n",
    "        )\n",
    "        self.v0s = torch.zeros(\n",
    "            (self.env_num, self.dim), dtype=torch.float32, device=self.device\n",
    "        )\n",
    "        self.p1s = torch.zeros(\n",
    "            (self.env_num, self.dim), dtype=torch.float32, device=self.device\n",
    "        )\n",
    "        self.v1s = torch.zeros(\n",
    "            (self.env_num, self.dim), dtype=torch.float32, device=self.device\n",
    "        )\n",
    "\n",
    "        self.cur_steps = torch.zeros(\n",
    "            self.env_num, dtype=torch.float32, device=self.device\n",
    "        )\n",
    "\n",
    "        for env_i in range(self.env_num):\n",
    "            self.reset_env_i(env_i)\n",
    "\n",
    "        self.distances = ((self.p0s - self.p1s) ** 2).sum(dim=1) ** 0.5\n",
    "\n",
    "        return self.get_state()\n",
    "\n",
    "    def reset_env_i(self, i):\n",
    "        self.p0s[i] = torch.normal(0, 1, size=(self.dim,))\n",
    "        self.v0s[i] = torch.zeros((self.dim,))\n",
    "        self.p1s[i] = torch.normal(-self.init_distance, 1, size=(self.dim,))\n",
    "        self.v1s[i] = torch.zeros((self.dim,))\n",
    "\n",
    "        self.cur_steps[i] = 0\n",
    "\n",
    "    def step(self, actions):\n",
    "        \"\"\"\n",
    "        :param actions: [tensor] actions.shape == (env_num, action_dim)\n",
    "        :return: next_states [tensor] next_states.shape == (env_num, state_dim)\n",
    "        :return: rewards [tensor] rewards == (env_num, )\n",
    "        :return: dones [tensor] dones == (env_num, ), done = 1. if done else 0.\n",
    "        :return: None [None or dict]\n",
    "        \"\"\"\n",
    "        # assert actions.get_device() == self.device.index\n",
    "        actions_l2 = (actions**2).sum(dim=1, keepdim=True) ** 0.5\n",
    "        actions_l2 = actions_l2.clamp_min(1.0)\n",
    "        actions = actions / actions_l2\n",
    "\n",
    "        self.v1s *= 0.75\n",
    "        self.v1s += actions\n",
    "        self.p1s += self.v1s * 0.01\n",
    "\n",
    "        self.v0s *= 0.50\n",
    "        self.v0s += torch.rand(\n",
    "            size=(self.env_num, self.dim), dtype=torch.float32, device=self.device\n",
    "        )\n",
    "        self.p0s += self.v0s * 0.01\n",
    "\n",
    "        \"\"\"reward\"\"\"\n",
    "        distances = ((self.p0s - self.p1s) ** 2).sum(dim=1) ** 0.5\n",
    "        rewards = self.distances - distances - actions_l2.squeeze(1) * 0.02\n",
    "        self.distances = distances\n",
    "\n",
    "        \"\"\"done\"\"\"\n",
    "        self.cur_steps += 1  # array\n",
    "        dones = (distances < self.dim) | (self.cur_steps == self.max_step)\n",
    "        for env_i in range(self.env_num):\n",
    "            if dones[env_i]:\n",
    "                self.reset_env_i(env_i)\n",
    "        dones = dones.type(torch.float32)\n",
    "\n",
    "        \"\"\"next_state\"\"\"\n",
    "        next_states = self.get_state()\n",
    "\n",
    "        # assert next_states.get_device() == self.device.index\n",
    "        # assert rewards.get_device() == self.device.index\n",
    "        # assert dones.get_device() == self.device.index\n",
    "        return next_states, rewards, dones, None\n",
    "\n",
    "    def get_state(self):\n",
    "        return torch.cat((self.p0s, self.v0s, self.p1s, self.v1s), dim=1)\n",
    "\n",
    "    @staticmethod\n",
    "    def get_action(states):\n",
    "        states_reshape = states.reshape((states.shape[0], 4, -1))\n",
    "        p0s = states_reshape[:, 0]\n",
    "        p1s = states_reshape[:, 2]\n",
    "        actions = p0s - p1s\n",
    "        return actions\n",
    "\n",
    "\n",
    "def check_chasing_vec_env():\n",
    "    env = ChasingVecEnv(dim=2, env_num=4096, device_id=0)\n",
    "    print(\"Env num:\", env.env_num)\n",
    "\n",
    "    reward_sums = [\n",
    "        0.0,\n",
    "    ] * env.env_num  # episode returns\n",
    "    reward_sums_list = [\n",
    "        [],\n",
    "    ] * env.env_num\n",
    "\n",
    "    states = env.reset()\n",
    "    for _ in range(env.max_step * 4):\n",
    "        actions = env.get_action(states)\n",
    "        states, rewards, dones, _ = env.step(actions)\n",
    "        print(\"Steps:\", env.cur_steps)\n",
    "        for env_i in range(env.env_num):\n",
    "            reward_sums[env_i] += rewards[env_i].item()\n",
    "\n",
    "            if dones[env_i]:\n",
    "                print(\n",
    "                    f\"{env.distances[env_i].item():8.4f}    {actions[env_i].detach().cpu().numpy().round(2)}\"\n",
    "                )\n",
    "                reward_sums_list[env_i].append(reward_sums[env_i])\n",
    "                reward_sums[env_i] = 0.0\n",
    "\n",
    "    reward_sums_list = np.array(reward_sums_list)\n",
    "    print(\"shape:\", reward_sums_list.shape)\n",
    "    print(\"mean: \", np.mean(reward_sums_list, axis=1))\n",
    "    print(\"std:  \", np.std(reward_sums_list, axis=1))\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    check_chasing_vec_env()"
   ]
  }
 ]
}