sac_kernel.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "928119c5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "PERIOD = 7\n",
    "\n",
    "# Env\n",
    "import gym, json, time\n",
    "import argparse\n",
    "from gym import spaces\n",
    "from epipolicy.core.epidemic import construct_epidemic\n",
    "from epipolicy.obj.act import construct_act\n",
    "import numpy as np\n",
    "\n",
    "class EpiEnv(gym.Env):\n",
    "    \"\"\"Custom Environment that follows gym interface\"\"\"\n",
    "    metadata = {'render.modes': ['human']}\n",
    "\n",
    "    def __init__(self, session):\n",
    "        super(EpiEnv, self).__init__()\n",
    "        self.epi = construct_epidemic(session)\n",
    "        total_population = np.sum(self.epi.static.default_state.obs.current_comp)\n",
    "        obs_count = self.epi.static.compartment_count * self.epi.static.locale_count * self.epi.static.group_count\n",
    "        action_count = 0\n",
    "        action_param_count =  0\n",
    "        for itv in self.epi.static.interventions:\n",
    "            if not itv.is_cost:\n",
    "                action_count += 1\n",
    "                action_param_count += len(itv.cp_list)\n",
    "        self.act_domain = np.zeros((action_param_count, 2), dtype=np.float32)\n",
    "        index = 0\n",
    "        for itv in self.epi.static.interventions:\n",
    "            if not itv.is_cost:\n",
    "                for cp in itv.cp_list:\n",
    "                    self.act_domain[index, 0] = cp.min_value\n",
    "                    self.act_domain[index, 1] = cp.max_value\n",
    "                    index += 1\n",
    "        # Define action and observation space\n",
    "        # They must be gym.spaces objects\n",
    "        # Example when using discrete actions:\n",
    "        self.action_space = spaces.Box(low=0, high=1, shape=(action_count,), dtype=np.float32)\n",
    "        # Example for using image as input:\n",
    "        self.observation_space = spaces.Box(low=0, high=total_population, shape=(obs_count,), dtype=np.float32)\n",
    "\n",
    "    def step(self, action):\n",
    "        expanded_action = np.zeros(len(self.act_domain), dtype=np.float32)\n",
    "        index = 0\n",
    "        for i in range(len(self.act_domain)):\n",
    "            if self.act_domain[i, 0] == self.act_domain[i, 1]:\n",
    "                expanded_action[i] = self.act_domain[i, 0]\n",
    "            else:\n",
    "                expanded_action[i] = min(max(0, action[index]), 1)\n",
    "                index += 1\n",
    "\n",
    "        epi_action = []\n",
    "        index = 0\n",
    "        for itv_id, itv in enumerate(self.epi.static.interventions):\n",
    "            if not itv.is_cost:\n",
    "                epi_action.append(construct_act(itv_id, expanded_action[index:index+len(itv.cp_list)]))\n",
    "                index += len(itv.cp_list)\n",
    "\n",
    "        total_r = 0\n",
    "        for i in range(PERIOD):\n",
    "            state, r, done = self.epi.step(epi_action)\n",
    "            total_r += r\n",
    "            if done:\n",
    "                break\n",
    "        return state.obs.current_comp.flatten(), total_r, done, dict()\n",
    "\n",
    "    def reset(self):\n",
    "        state = self.epi.reset()\n",
    "        return state.obs.current_comp.flatten()  # reward, done, info can't be included\n",
    "    def render(self, mode='human'):\n",
    "        pass\n",
    "    def close(self):\n",
    "        pass\n",
    "\n",
    "class RewardScale(gym.RewardWrapper):\n",
    "    def __init__(self, env, scale):\n",
    "        super().__init__(env)\n",
    "        self.scale = scale\n",
    "    \n",
    "    def reward(self, rew):\n",
    "        # modify rew\n",
    "        return rew * self.scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "57301ef1",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "epi_ids = [\"SIR_A\", \"SIR_B\", \"SIRV_A\", \"SIRV_B\", \"COVID_A\", \"COVID_B\", \"COVID_C\"]\n",
    "\n",
    "def make_env(gym_id, seed, idx):\n",
    "    def thunk():\n",
    "        if gym_id in epi_ids:\n",
    "            fp = open('jsons/{}.json'.format(gym_id), 'r')\n",
    "            session = json.load(fp)\n",
    "            env = EpiEnv(session)\n",
    "        else:\n",
    "            env = gym.make(gym_id)\n",
    "        env = gym.wrappers.RecordEpisodeStatistics(env)\n",
    "        env = gym.wrappers.ClipAction(env)\n",
    "        env = gym.wrappers.NormalizeObservation(env)\n",
    "        env = gym.wrappers.TransformObservation(env, lambda obs: np.clip(obs, -10, 10))\n",
    "        env = gym.wrappers.NormalizeReward(env)\n",
    "        env = gym.wrappers.TransformReward(env, lambda reward: np.clip(reward, -10, 10))\n",
    "        env = RewardScale(env, 10)\n",
    "        # Our env is deterministic\n",
    "        # env.seed(seed)\n",
    "        env.action_space.seed(seed)\n",
    "        env.observation_space.seed(seed)\n",
    "        return env\n",
    "\n",
    "    return thunk\n",
    "\n",
    "def make_primal_env(gym_id):\n",
    "    def thunk():\n",
    "        if gym_id in epi_ids:\n",
    "            fp = open('jsons/{}.json'.format(gym_id), 'r')\n",
    "            session = json.load(fp)\n",
    "            env = EpiEnv(session)\n",
    "        else:\n",
    "            env = gym.make(gym_id)\n",
    "        return env\n",
    "    return thunk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d68b1dc7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def parse_args(main_args = None):\n",
    "    # fmt: off\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument(\"--exp-name\", type=str, default=\"SAC\",\n",
    "        help=\"the name of this experiment\")\n",
    "    parser.add_argument(\"--gym-id\", type=str, default=\"HalfCheetahBulletEnv-v0\",\n",
    "        help=\"the id of the gym environment\")\n",
    "#     parser.add_argument(\"--learning-rate\", type=float, default=3e-4,\n",
    "#         help=\"the learning rate of the optimizer\")\n",
    "    parser.add_argument(\"--seed\", type=int, default=1,\n",
    "        help=\"seed of the experiment\")\n",
    "    parser.add_argument(\"--total-timesteps\", type=int, default=700000,\n",
    "        help=\"total timesteps of the experiments\")\n",
    "#     parser.add_argument(\"--torch-deterministic\", type=lambda x: bool(strtobool(x)), default=True, nargs=\"?\", const=True,\n",
    "#         help=\"if toggled, `torch.backends.cudnn.deterministic=False`\")\n",
    "#     parser.add_argument(\"--cuda\", type=lambda x: bool(strtobool(x)), default=True, nargs=\"?\", const=True,\n",
    "#         help=\"if toggled, cuda will be enabled by default\")\n",
    "#     parser.add_argument(\"--track\", type=lambda x: bool(strtobool(x)), default=False, nargs=\"?\", const=True,\n",
    "#         help=\"if toggled, this experiment will be tracked with Weights and Biases\")\n",
    "#     parser.add_argument(\"--wandb-project-name\", type=str, default=\"ppo-implementation-details\",\n",
    "#         help=\"the wandb's project name\")\n",
    "#     parser.add_argument(\"--wandb-entity\", type=str, default=None,\n",
    "#         help=\"the entity (team) of wandb's project\")\n",
    "#     parser.add_argument(\"--capture-video\", type=lambda x: bool(strtobool(x)), default=False, nargs=\"?\", const=True,\n",
    "#         help=\"weather to capture videos of the agent performances (check out `videos` folder)\")\n",
    "\n",
    "#     parser.add_argument(\"--policy_plot_interval\", type=int, default=1,\n",
    "#         help=\"seed of the experiment\")\n",
    "\n",
    "    # Algorithm specific arguments\n",
    "    parser.add_argument(\"--learning-starts\", type=int, default=5000,\n",
    "        help=\"learning starts\")\n",
    "    parser.add_argument(\"--target-entropy-scale\", type=int, default=1,\n",
    "        help=\"scale of target entropy with the dimension of action\")\n",
    "    parser.add_argument(\"--train-freq\", type=int, default=146,\n",
    "        help=\"train freq\")\n",
    "    parser.add_argument(\"--gradient-steps\", type=int, default=1,\n",
    "        help=\"gradient steps\")\n",
    "#     parser.add_argument(\"--num-envs\", type=int, default=1,\n",
    "#         help=\"the number of parallel game environments\")\n",
    "#     parser.add_argument(\"--num-steps\", type=int, default=2048,\n",
    "#         help=\"the number of steps to run in each environment per policy rollout\")\n",
    "#     parser.add_argument(\"--anneal-lr\", type=lambda x: bool(strtobool(x)), default=True, nargs=\"?\", const=True,\n",
    "#         help=\"Toggle learning rate annealing for policy and value networks\")\n",
    "#     parser.add_argument(\"--gae\", type=lambda x: bool(strtobool(x)), default=True, nargs=\"?\", const=True,\n",
    "#         help=\"Use GAE for advantage computation\")\n",
    "#     parser.add_argument(\"--gamma\", type=float, default=0.99,\n",
    "#         help=\"the discount factor gamma\")\n",
    "#     parser.add_argument(\"--gae-lambda\", type=float, default=0.95,\n",
    "#         help=\"the lambda for the general advantage estimation\")\n",
    "#     parser.add_argument(\"--num-minibatches\", type=int, default=32,\n",
    "#         help=\"the number of mini-batches\")\n",
    "#     parser.add_argument(\"--update-epochs\", type=int, default=10,\n",
    "#         help=\"the K epochs to update the policy\")\n",
    "#     parser.add_argument(\"--norm-adv\", type=lambda x: bool(strtobool(x)), default=True, nargs=\"?\", const=True,\n",
    "#         help=\"Toggles advantages normalization\")\n",
    "#     parser.add_argument(\"--clip-coef\", type=float, default=0.2,\n",
    "#         help=\"the surrogate clipping coefficient\")\n",
    "#     parser.add_argument(\"--clip-vloss\", type=lambda x: bool(strtobool(x)), default=True, nargs=\"?\", const=True,\n",
    "#         help=\"Toggles whether or not to use a clipped loss for the value function, as per the paper.\")\n",
    "#     parser.add_argument(\"--ent-coef\", type=float, default=0.0,\n",
    "#         help=\"coefficient of the entropy\")\n",
    "#     parser.add_argument(\"--vf-coef\", type=float, default=0.5,\n",
    "#         help=\"coefficient of the value function\")\n",
    "#     parser.add_argument(\"--max-grad-norm\", type=float, default=0.5,\n",
    "#         help=\"the maximum norm for the gradient clipping\")\n",
    "#     parser.add_argument(\"--target-kl\", type=float, default=None,\n",
    "#         help=\"the target KL divergence threshold\")\n",
    "    if main_args is not None:\n",
    "        args = parser.parse_args(main_args.split())\n",
    "    else:\n",
    "        args = parser.parse_args()\n",
    "#     args.num_steps //= PERIOD\n",
    "    args.total_timesteps //= PERIOD\n",
    "#     args.batch_size = int(args.num_envs * args.num_steps)\n",
    "#     args.minibatch_size = int(args.batch_size // args.num_minibatches)\n",
    "    # fmt: on\n",
    "    return args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ce82c583",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running with SIR_A__SAC__1__1660127166\n",
      "Using cpu device\n",
      "Wrapping the env with a `Monitor` wrapper\n",
      "Wrapping the env in a DummyVecEnv.\n",
      "Logging to runs/SIR_A__SAC__1__1660127166_1\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -91.9    |\n",
      "| time/              |          |\n",
      "|    episodes        | 4        |\n",
      "|    fps             | 39       |\n",
      "|    time_elapsed    | 5        |\n",
      "|    total_timesteps | 212      |\n",
      "---------------------------------\n",
      "At global step 292, total_rewards=-94383246.70791376\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -66.8    |\n",
      "| time/              |          |\n",
      "|    episodes        | 8        |\n",
      "|    fps             | 33       |\n",
      "|    time_elapsed    | 13       |\n",
      "|    total_timesteps | 451      |\n",
      "---------------------------------\n",
      "At global step 584, total_rewards=-94667611.9288868\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -59.4    |\n",
      "| time/              |          |\n",
      "|    episodes        | 12       |\n",
      "|    fps             | 31       |\n",
      "|    time_elapsed    | 21       |\n",
      "|    total_timesteps | 690      |\n",
      "---------------------------------\n",
      "At global step 876, total_rewards=-94818796.22871138\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -54.5    |\n",
      "| time/              |          |\n",
      "|    episodes        | 16       |\n",
      "|    fps             | 30       |\n",
      "|    time_elapsed    | 30       |\n",
      "|    total_timesteps | 929      |\n",
      "---------------------------------\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -52.4    |\n",
      "| time/              |          |\n",
      "|    episodes        | 20       |\n",
      "|    fps             | 19       |\n",
      "|    time_elapsed    | 58       |\n",
      "|    total_timesteps | 1141     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -1.77    |\n",
      "|    critic_loss     | 0.535    |\n",
      "|    ent_coef        | 0.992    |\n",
      "|    ent_coef_loss   | -0.027   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 28       |\n",
      "---------------------------------\n",
      "At global step 1168, total_rewards=-103099628.15075529\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -51.6    |\n",
      "| time/              |          |\n",
      "|    episodes        | 24       |\n",
      "|    fps             | 12       |\n",
      "|    time_elapsed    | 107      |\n",
      "|    total_timesteps | 1380     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -1.95    |\n",
      "|    critic_loss     | 0.158    |\n",
      "|    ent_coef        | 0.978    |\n",
      "|    ent_coef_loss   | -0.0737  |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 75       |\n",
      "---------------------------------\n",
      "At global step 1460, total_rewards=-95516212.78113538\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -48.6    |\n",
      "| time/              |          |\n",
      "|    episodes        | 28       |\n",
      "|    fps             | 10       |\n",
      "|    time_elapsed    | 156      |\n",
      "|    total_timesteps | 1619     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -1.78    |\n",
      "|    critic_loss     | 1.02     |\n",
      "|    ent_coef        | 0.964    |\n",
      "|    ent_coef_loss   | -0.12    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 123      |\n",
      "---------------------------------\n",
      "At global step 1752, total_rewards=-94160372.9295787\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -46      |\n",
      "| time/              |          |\n",
      "|    episodes        | 32       |\n",
      "|    fps             | 9        |\n",
      "|    time_elapsed    | 205      |\n",
      "|    total_timesteps | 1858     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -1.98    |\n",
      "|    critic_loss     | 0.112    |\n",
      "|    ent_coef        | 0.95     |\n",
      "|    ent_coef_loss   | -0.168   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 171      |\n",
      "---------------------------------\n",
      "At global step 2044, total_rewards=-92593702.00879583\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -43.9    |\n",
      "| time/              |          |\n",
      "|    episodes        | 36       |\n",
      "|    fps             | 8        |\n",
      "|    time_elapsed    | 252      |\n",
      "|    total_timesteps | 2097     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -1.99    |\n",
      "|    critic_loss     | 38.3     |\n",
      "|    ent_coef        | 0.937    |\n",
      "|    ent_coef_loss   | -0.217   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 219      |\n",
      "---------------------------------\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -42.3    |\n",
      "| time/              |          |\n",
      "|    episodes        | 40       |\n",
      "|    fps             | 7        |\n",
      "|    time_elapsed    | 293      |\n",
      "|    total_timesteps | 2309     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -1.94    |\n",
      "|    critic_loss     | 0.329    |\n",
      "|    ent_coef        | 0.925    |\n",
      "|    ent_coef_loss   | -0.251   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 261      |\n",
      "---------------------------------\n",
      "At global step 2336, total_rewards=-95167496.78739323\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -41.8    |\n",
      "| time/              |          |\n",
      "|    episodes        | 44       |\n",
      "|    fps             | 7        |\n",
      "|    time_elapsed    | 342      |\n",
      "|    total_timesteps | 2548     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.15    |\n",
      "|    critic_loss     | 0.114    |\n",
      "|    ent_coef        | 0.912    |\n",
      "|    ent_coef_loss   | -0.306   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 309      |\n",
      "---------------------------------\n",
      "At global step 2628, total_rewards=-101196841.62345222\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -41      |\n",
      "| time/              |          |\n",
      "|    episodes        | 48       |\n",
      "|    fps             | 7        |\n",
      "|    time_elapsed    | 390      |\n",
      "|    total_timesteps | 2787     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.05    |\n",
      "|    critic_loss     | 0.154    |\n",
      "|    ent_coef        | 0.899    |\n",
      "|    ent_coef_loss   | -0.348   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 357      |\n",
      "---------------------------------\n",
      "At global step 2920, total_rewards=-97418878.18352558\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -40.7    |\n",
      "| time/              |          |\n",
      "|    episodes        | 52       |\n",
      "|    fps             | 6        |\n",
      "|    time_elapsed    | 438      |\n",
      "|    total_timesteps | 3026     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.05    |\n",
      "|    critic_loss     | 38.9     |\n",
      "|    ent_coef        | 0.886    |\n",
      "|    ent_coef_loss   | -0.398   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 405      |\n",
      "---------------------------------\n",
      "At global step 3212, total_rewards=-95667774.89713675\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -40      |\n",
      "| time/              |          |\n",
      "|    episodes        | 56       |\n",
      "|    fps             | 6        |\n",
      "|    time_elapsed    | 485      |\n",
      "|    total_timesteps | 3265     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.25    |\n",
      "|    critic_loss     | 1.15     |\n",
      "|    ent_coef        | 0.874    |\n",
      "|    ent_coef_loss   | -0.441   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 452      |\n",
      "---------------------------------\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -39.3    |\n",
      "| time/              |          |\n",
      "|    episodes        | 60       |\n",
      "|    fps             | 6        |\n",
      "|    time_elapsed    | 527      |\n",
      "|    total_timesteps | 3477     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.33    |\n",
      "|    critic_loss     | 0.127    |\n",
      "|    ent_coef        | 0.863    |\n",
      "|    ent_coef_loss   | -0.49    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 495      |\n",
      "---------------------------------\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "At global step 3504, total_rewards=-85926251.13566485\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -38.7    |\n",
      "| time/              |          |\n",
      "|    episodes        | 64       |\n",
      "|    fps             | 6        |\n",
      "|    time_elapsed    | 577      |\n",
      "|    total_timesteps | 3716     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.43    |\n",
      "|    critic_loss     | 0.124    |\n",
      "|    ent_coef        | 0.85     |\n",
      "|    ent_coef_loss   | -0.531   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 543      |\n",
      "---------------------------------\n",
      "At global step 3796, total_rewards=-85101817.36182101\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -38.2    |\n",
      "| time/              |          |\n",
      "|    episodes        | 68       |\n",
      "|    fps             | 6        |\n",
      "|    time_elapsed    | 626      |\n",
      "|    total_timesteps | 3955     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.55    |\n",
      "|    critic_loss     | 0.836    |\n",
      "|    ent_coef        | 0.838    |\n",
      "|    ent_coef_loss   | -0.583   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 590      |\n",
      "---------------------------------\n",
      "At global step 4088, total_rewards=-94060781.01388818\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -37.7    |\n",
      "| time/              |          |\n",
      "|    episodes        | 72       |\n",
      "|    fps             | 6        |\n",
      "|    time_elapsed    | 675      |\n",
      "|    total_timesteps | 4194     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.46    |\n",
      "|    critic_loss     | 40.7     |\n",
      "|    ent_coef        | 0.827    |\n",
      "|    ent_coef_loss   | -0.624   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 638      |\n",
      "---------------------------------\n",
      "At global step 4380, total_rewards=-85016480.10450166\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -37.4    |\n",
      "| time/              |          |\n",
      "|    episodes        | 76       |\n",
      "|    fps             | 6        |\n",
      "|    time_elapsed    | 724      |\n",
      "|    total_timesteps | 4433     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.54    |\n",
      "|    critic_loss     | 0.165    |\n",
      "|    ent_coef        | 0.815    |\n",
      "|    ent_coef_loss   | -0.661   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 686      |\n",
      "---------------------------------\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -37.2    |\n",
      "| time/              |          |\n",
      "|    episodes        | 80       |\n",
      "|    fps             | 6        |\n",
      "|    time_elapsed    | 764      |\n",
      "|    total_timesteps | 4645     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.45    |\n",
      "|    critic_loss     | 39.9     |\n",
      "|    ent_coef        | 0.805    |\n",
      "|    ent_coef_loss   | -0.719   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 728      |\n",
      "---------------------------------\n",
      "At global step 4672, total_rewards=-89488171.03607967\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -36.9    |\n",
      "| time/              |          |\n",
      "|    episodes        | 84       |\n",
      "|    fps             | 6        |\n",
      "|    time_elapsed    | 813      |\n",
      "|    total_timesteps | 4884     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.47    |\n",
      "|    critic_loss     | 0.157    |\n",
      "|    ent_coef        | 0.793    |\n",
      "|    ent_coef_loss   | -0.76    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 776      |\n",
      "---------------------------------\n",
      "At global step 4964, total_rewards=-87390897.1712424\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -36.8    |\n",
      "| time/              |          |\n",
      "|    episodes        | 88       |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 862      |\n",
      "|    total_timesteps | 5123     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.67    |\n",
      "|    critic_loss     | 0.214    |\n",
      "|    ent_coef        | 0.782    |\n",
      "|    ent_coef_loss   | -0.789   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 824      |\n",
      "---------------------------------\n",
      "At global step 5256, total_rewards=-82984933.14528038\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -36.5    |\n",
      "| time/              |          |\n",
      "|    episodes        | 92       |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 910      |\n",
      "|    total_timesteps | 5362     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.8     |\n",
      "|    critic_loss     | 0.174    |\n",
      "|    ent_coef        | 0.771    |\n",
      "|    ent_coef_loss   | -0.836   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 872      |\n",
      "---------------------------------\n",
      "At global step 5548, total_rewards=-85986441.68091908\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -36.2    |\n",
      "| time/              |          |\n",
      "|    episodes        | 96       |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 959      |\n",
      "|    total_timesteps | 5601     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.79    |\n",
      "|    critic_loss     | 0.0988   |\n",
      "|    ent_coef        | 0.76     |\n",
      "|    ent_coef_loss   | -0.905   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 920      |\n",
      "---------------------------------\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -36.1    |\n",
      "| time/              |          |\n",
      "|    episodes        | 100      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1000     |\n",
      "|    total_timesteps | 5813     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.92    |\n",
      "|    critic_loss     | 0.344    |\n",
      "|    ent_coef        | 0.75     |\n",
      "|    ent_coef_loss   | -0.93    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 962      |\n",
      "---------------------------------\n",
      "At global step 5840, total_rewards=-85319245.7735851\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -33.8    |\n",
      "| time/              |          |\n",
      "|    episodes        | 104      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1049     |\n",
      "|    total_timesteps | 6052     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.85    |\n",
      "|    critic_loss     | 0.201    |\n",
      "|    ent_coef        | 0.74     |\n",
      "|    ent_coef_loss   | -0.982   |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1010     |\n",
      "---------------------------------\n",
      "At global step 6132, total_rewards=-82011858.01360163\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -33.4    |\n",
      "| time/              |          |\n",
      "|    episodes        | 108      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1100     |\n",
      "|    total_timesteps | 6291     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.95    |\n",
      "|    critic_loss     | 0.238    |\n",
      "|    ent_coef        | 0.729    |\n",
      "|    ent_coef_loss   | -1.01    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1058     |\n",
      "---------------------------------\n",
      "At global step 6424, total_rewards=-80359713.7055182\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -32.8    |\n",
      "| time/              |          |\n",
      "|    episodes        | 112      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1149     |\n",
      "|    total_timesteps | 6530     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.99    |\n",
      "|    critic_loss     | 0.176    |\n",
      "|    ent_coef        | 0.719    |\n",
      "|    ent_coef_loss   | -1.05    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1105     |\n",
      "---------------------------------\n",
      "At global step 6716, total_rewards=-79413086.95335521\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -32.4    |\n",
      "| time/              |          |\n",
      "|    episodes        | 116      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1201     |\n",
      "|    total_timesteps | 6769     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.83    |\n",
      "|    critic_loss     | 0.753    |\n",
      "|    ent_coef        | 0.709    |\n",
      "|    ent_coef_loss   | -1.11    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1153     |\n",
      "---------------------------------\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -31.9    |\n",
      "| time/              |          |\n",
      "|    episodes        | 120      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1241     |\n",
      "|    total_timesteps | 6981     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.04    |\n",
      "|    critic_loss     | 0.18     |\n",
      "|    ent_coef        | 0.7      |\n",
      "|    ent_coef_loss   | -1.17    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1196     |\n",
      "---------------------------------\n",
      "At global step 7008, total_rewards=-81651465.17165944\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -31.2    |\n",
      "| time/              |          |\n",
      "|    episodes        | 124      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1291     |\n",
      "|    total_timesteps | 7220     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.07    |\n",
      "|    critic_loss     | 0.13     |\n",
      "|    ent_coef        | 0.69     |\n",
      "|    ent_coef_loss   | -1.19    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1243     |\n",
      "---------------------------------\n",
      "At global step 7300, total_rewards=-80677988.45741206\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -31.1    |\n",
      "| time/              |          |\n",
      "|    episodes        | 128      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1341     |\n",
      "|    total_timesteps | 7459     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.12    |\n",
      "|    critic_loss     | 0.167    |\n",
      "|    ent_coef        | 0.68     |\n",
      "|    ent_coef_loss   | -1.24    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1291     |\n",
      "---------------------------------\n",
      "At global step 7592, total_rewards=-78433922.34606275\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -31.2    |\n",
      "| time/              |          |\n",
      "|    episodes        | 132      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1390     |\n",
      "|    total_timesteps | 7698     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.04    |\n",
      "|    critic_loss     | 0.137    |\n",
      "|    ent_coef        | 0.671    |\n",
      "|    ent_coef_loss   | -1.29    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1339     |\n",
      "---------------------------------\n",
      "At global step 7884, total_rewards=-78975957.46991266\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -31.3    |\n",
      "| time/              |          |\n",
      "|    episodes        | 136      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1438     |\n",
      "|    total_timesteps | 7937     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.17    |\n",
      "|    critic_loss     | 0.0939   |\n",
      "|    ent_coef        | 0.661    |\n",
      "|    ent_coef_loss   | -1.34    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1387     |\n",
      "---------------------------------\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -31.4    |\n",
      "| time/              |          |\n",
      "|    episodes        | 140      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1479     |\n",
      "|    total_timesteps | 8149     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.99    |\n",
      "|    critic_loss     | 0.143    |\n",
      "|    ent_coef        | 0.653    |\n",
      "|    ent_coef_loss   | -1.39    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1429     |\n",
      "---------------------------------\n",
      "At global step 8176, total_rewards=-78523636.34999043\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -31.2    |\n",
      "| time/              |          |\n",
      "|    episodes        | 144      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1530     |\n",
      "|    total_timesteps | 8388     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.99    |\n",
      "|    critic_loss     | 0.124    |\n",
      "|    ent_coef        | 0.644    |\n",
      "|    ent_coef_loss   | -1.42    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1477     |\n",
      "---------------------------------\n",
      "At global step 8468, total_rewards=-75233949.19534516\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -31      |\n",
      "| time/              |          |\n",
      "|    episodes        | 148      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1579     |\n",
      "|    total_timesteps | 8627     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.25    |\n",
      "|    critic_loss     | 0.151    |\n",
      "|    ent_coef        | 0.635    |\n",
      "|    ent_coef_loss   | -1.44    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1525     |\n",
      "---------------------------------\n",
      "At global step 8760, total_rewards=-77246499.65993024\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.7    |\n",
      "| time/              |          |\n",
      "|    episodes        | 152      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1628     |\n",
      "|    total_timesteps | 8866     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.96    |\n",
      "|    critic_loss     | 0.443    |\n",
      "|    ent_coef        | 0.626    |\n",
      "|    ent_coef_loss   | -1.52    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1573     |\n",
      "---------------------------------\n",
      "At global step 9052, total_rewards=-75097728.07388023\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.7    |\n",
      "| time/              |          |\n",
      "|    episodes        | 156      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1679     |\n",
      "|    total_timesteps | 9105     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.27    |\n",
      "|    critic_loss     | 0.143    |\n",
      "|    ent_coef        | 0.617    |\n",
      "|    ent_coef_loss   | -1.54    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1620     |\n",
      "---------------------------------\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.8    |\n",
      "| time/              |          |\n",
      "|    episodes        | 160      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1719     |\n",
      "|    total_timesteps | 9317     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.14    |\n",
      "|    critic_loss     | 0.142    |\n",
      "|    ent_coef        | 0.609    |\n",
      "|    ent_coef_loss   | -1.61    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1663     |\n",
      "---------------------------------\n",
      "At global step 9344, total_rewards=-75443866.65373527\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.8    |\n",
      "| time/              |          |\n",
      "|    episodes        | 164      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1769     |\n",
      "|    total_timesteps | 9556     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.92    |\n",
      "|    critic_loss     | 0.6      |\n",
      "|    ent_coef        | 0.601    |\n",
      "|    ent_coef_loss   | -1.64    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1711     |\n",
      "---------------------------------\n",
      "At global step 9636, total_rewards=-74909199.70519502\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.8    |\n",
      "| time/              |          |\n",
      "|    episodes        | 168      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1819     |\n",
      "|    total_timesteps | 9795     |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.18    |\n",
      "|    critic_loss     | 0.206    |\n",
      "|    ent_coef        | 0.592    |\n",
      "|    ent_coef_loss   | -1.67    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1758     |\n",
      "---------------------------------\n",
      "At global step 9928, total_rewards=-73161022.6022377\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.9    |\n",
      "| time/              |          |\n",
      "|    episodes        | 172      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1870     |\n",
      "|    total_timesteps | 10034    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.15    |\n",
      "|    critic_loss     | 0.119    |\n",
      "|    ent_coef        | 0.584    |\n",
      "|    ent_coef_loss   | -1.73    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1806     |\n",
      "---------------------------------\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "At global step 10220, total_rewards=-73310531.44310786\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.8    |\n",
      "| time/              |          |\n",
      "|    episodes        | 176      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1919     |\n",
      "|    total_timesteps | 10273    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.26    |\n",
      "|    critic_loss     | 0.083    |\n",
      "|    ent_coef        | 0.576    |\n",
      "|    ent_coef_loss   | -1.75    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1854     |\n",
      "---------------------------------\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.6    |\n",
      "| time/              |          |\n",
      "|    episodes        | 180      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 1960     |\n",
      "|    total_timesteps | 10485    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.27    |\n",
      "|    critic_loss     | 0.0617   |\n",
      "|    ent_coef        | 0.569    |\n",
      "|    ent_coef_loss   | -1.83    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1896     |\n",
      "---------------------------------\n",
      "At global step 10512, total_rewards=-73153161.79674536\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.6    |\n",
      "| time/              |          |\n",
      "|    episodes        | 184      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 2011     |\n",
      "|    total_timesteps | 10724    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.26    |\n",
      "|    critic_loss     | 0.069    |\n",
      "|    ent_coef        | 0.561    |\n",
      "|    ent_coef_loss   | -1.84    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1944     |\n",
      "---------------------------------\n",
      "At global step 10804, total_rewards=-72215839.97194454\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.5    |\n",
      "| time/              |          |\n",
      "|    episodes        | 188      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 2060     |\n",
      "|    total_timesteps | 10963    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.07    |\n",
      "|    critic_loss     | 0.111    |\n",
      "|    ent_coef        | 0.553    |\n",
      "|    ent_coef_loss   | -1.85    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 1992     |\n",
      "---------------------------------\n",
      "At global step 11096, total_rewards=-71361139.83003181\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.5    |\n",
      "| time/              |          |\n",
      "|    episodes        | 192      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 2111     |\n",
      "|    total_timesteps | 11202    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.17    |\n",
      "|    critic_loss     | 0.114    |\n",
      "|    ent_coef        | 0.545    |\n",
      "|    ent_coef_loss   | -1.92    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 2040     |\n",
      "---------------------------------\n",
      "At global step 11388, total_rewards=-68767759.26510009\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.5    |\n",
      "| time/              |          |\n",
      "|    episodes        | 196      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 2161     |\n",
      "|    total_timesteps | 11441    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.16    |\n",
      "|    critic_loss     | 0.0751   |\n",
      "|    ent_coef        | 0.537    |\n",
      "|    ent_coef_loss   | -1.94    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 2088     |\n",
      "---------------------------------\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.3    |\n",
      "| time/              |          |\n",
      "|    episodes        | 200      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 2203     |\n",
      "|    total_timesteps | 11653    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.12    |\n",
      "|    critic_loss     | 0.075    |\n",
      "|    ent_coef        | 0.531    |\n",
      "|    ent_coef_loss   | -2.03    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 2130     |\n",
      "---------------------------------\n",
      "At global step 11680, total_rewards=-70922773.90878592\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.1    |\n",
      "| time/              |          |\n",
      "|    episodes        | 204      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 2252     |\n",
      "|    total_timesteps | 11892    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.98    |\n",
      "|    critic_loss     | 0.119    |\n",
      "|    ent_coef        | 0.523    |\n",
      "|    ent_coef_loss   | -2.07    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 2178     |\n",
      "---------------------------------\n",
      "At global step 11972, total_rewards=-70779446.74270637\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.1    |\n",
      "| time/              |          |\n",
      "|    episodes        | 208      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 2302     |\n",
      "|    total_timesteps | 12131    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.1     |\n",
      "|    critic_loss     | 0.0749   |\n",
      "|    ent_coef        | 0.516    |\n",
      "|    ent_coef_loss   | -2.08    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 2226     |\n",
      "---------------------------------\n",
      "At global step 12264, total_rewards=-70028455.94211055\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30.2    |\n",
      "| time/              |          |\n",
      "|    episodes        | 212      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 2351     |\n",
      "|    total_timesteps | 12370    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.06    |\n",
      "|    critic_loss     | 0.167    |\n",
      "|    ent_coef        | 0.509    |\n",
      "|    ent_coef_loss   | -2.06    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 2273     |\n",
      "---------------------------------\n",
      "At global step 12556, total_rewards=-68491049.87548564\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30      |\n",
      "| time/              |          |\n",
      "|    episodes        | 216      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 2400     |\n",
      "|    total_timesteps | 12609    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.97    |\n",
      "|    critic_loss     | 0.0885   |\n",
      "|    ent_coef        | 0.502    |\n",
      "|    ent_coef_loss   | -2.22    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 2321     |\n",
      "---------------------------------\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -30      |\n",
      "| time/              |          |\n",
      "|    episodes        | 220      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 2441     |\n",
      "|    total_timesteps | 12821    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -3.13    |\n",
      "|    critic_loss     | 0.0777   |\n",
      "|    ent_coef        | 0.495    |\n",
      "|    ent_coef_loss   | -2.21    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 2364     |\n",
      "---------------------------------\n",
      "At global step 12848, total_rewards=-66963317.80816582\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -29.9    |\n",
      "| time/              |          |\n",
      "|    episodes        | 224      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 2490     |\n",
      "|    total_timesteps | 13060    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.94    |\n",
      "|    critic_loss     | 0.0527   |\n",
      "|    ent_coef        | 0.489    |\n",
      "|    ent_coef_loss   | -2.23    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 2411     |\n",
      "---------------------------------\n",
      "At global step 13140, total_rewards=-68838358.50210287\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -29.9    |\n",
      "| time/              |          |\n",
      "|    episodes        | 228      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 2540     |\n",
      "|    total_timesteps | 13299    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.87    |\n",
      "|    critic_loss     | 0.274    |\n",
      "|    ent_coef        | 0.482    |\n",
      "|    ent_coef_loss   | -2.23    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 2459     |\n",
      "---------------------------------\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "At global step 13432, total_rewards=-67085955.40263419\n",
      "---------------------------------\n",
      "| rollout/           |          |\n",
      "|    ep_len_mean     | 53       |\n",
      "|    ep_rew_mean     | -29.9    |\n",
      "| time/              |          |\n",
      "|    episodes        | 232      |\n",
      "|    fps             | 5        |\n",
      "|    time_elapsed    | 2589     |\n",
      "|    total_timesteps | 13538    |\n",
      "| train/             |          |\n",
      "|    actor_loss      | -2.8     |\n",
      "|    critic_loss     | 0.101    |\n",
      "|    ent_coef        | 0.475    |\n",
      "|    ent_coef_loss   | -2.22    |\n",
      "|    learning_rate   | 0.0003   |\n",
      "|    n_updates       | 2507     |\n",
      "---------------------------------\n",
      "At global step 13724, total_rewards=-68219832.84810112\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Input \u001b[0;32mIn [8]\u001b[0m, in \u001b[0;36m<cell line: 68>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     65\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mRunning with \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mrun_name\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     66\u001b[0m model \u001b[38;5;241m=\u001b[39m SAC(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMlpPolicy\u001b[39m\u001b[38;5;124m\"\u001b[39m, env, verbose\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m, tensorboard_log\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mruns/\u001b[39m\u001b[38;5;124m\"\u001b[39m, learning_starts\u001b[38;5;241m=\u001b[39margs\u001b[38;5;241m.\u001b[39mlearning_starts, target_entropy\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m-\u001b[39margs\u001b[38;5;241m.\u001b[39mtarget_entropy_scale \u001b[38;5;241m*\u001b[39m env\u001b[38;5;241m.\u001b[39maction_space\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m0\u001b[39m],\n\u001b[1;32m     67\u001b[0m            train_freq\u001b[38;5;241m=\u001b[39margs\u001b[38;5;241m.\u001b[39mtrain_freq, gradient_steps\u001b[38;5;241m=\u001b[39margs\u001b[38;5;241m.\u001b[39mgradient_steps)\n\u001b[0;32m---> 68\u001b[0m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlearn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtotal_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtotal_timesteps\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlog_interval\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcallback\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mSummaryWriterCallback\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtb_log_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrun_name\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.8/site-packages/stable_baselines3/sac/sac.py:292\u001b[0m, in \u001b[0;36mSAC.learn\u001b[0;34m(self, total_timesteps, callback, log_interval, eval_env, eval_freq, n_eval_episodes, tb_log_name, eval_log_path, reset_num_timesteps)\u001b[0m\n\u001b[1;32m    279\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mlearn\u001b[39m(\n\u001b[1;32m    280\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m    281\u001b[0m     total_timesteps: \u001b[38;5;28mint\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    289\u001b[0m     reset_num_timesteps: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m    290\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m OffPolicyAlgorithm:\n\u001b[0;32m--> 292\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mSAC\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlearn\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    293\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtotal_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtotal_timesteps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    294\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcallback\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcallback\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    295\u001b[0m \u001b[43m        \u001b[49m\u001b[43mlog_interval\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlog_interval\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    296\u001b[0m \u001b[43m        \u001b[49m\u001b[43meval_env\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43meval_env\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    297\u001b[0m \u001b[43m        \u001b[49m\u001b[43meval_freq\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43meval_freq\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    298\u001b[0m \u001b[43m        \u001b[49m\u001b[43mn_eval_episodes\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mn_eval_episodes\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    299\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtb_log_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtb_log_name\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    300\u001b[0m \u001b[43m        \u001b[49m\u001b[43meval_log_path\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43meval_log_path\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    301\u001b[0m \u001b[43m        \u001b[49m\u001b[43mreset_num_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreset_num_timesteps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    302\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.8/site-packages/stable_baselines3/common/off_policy_algorithm.py:366\u001b[0m, in \u001b[0;36mOffPolicyAlgorithm.learn\u001b[0;34m(self, total_timesteps, callback, log_interval, eval_env, eval_freq, n_eval_episodes, tb_log_name, eval_log_path, reset_num_timesteps)\u001b[0m\n\u001b[1;32m    364\u001b[0m         \u001b[38;5;66;03m# Special case when the user passes `gradient_steps=0`\u001b[39;00m\n\u001b[1;32m    365\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m gradient_steps \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[0;32m--> 366\u001b[0m             \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbatch_size\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient_steps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mgradient_steps\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    368\u001b[0m callback\u001b[38;5;241m.\u001b[39mon_training_end()\n\u001b[1;32m    370\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n",
      "File \u001b[0;32m~/.local/lib/python3.8/site-packages/stable_baselines3/sac/sac.py:206\u001b[0m, in \u001b[0;36mSAC.train\u001b[0;34m(self, gradient_steps, batch_size)\u001b[0m\n\u001b[1;32m    203\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mactor\u001b[38;5;241m.\u001b[39mreset_noise()\n\u001b[1;32m    205\u001b[0m \u001b[38;5;66;03m# Action by the current actor for the sampled state\u001b[39;00m\n\u001b[0;32m--> 206\u001b[0m actions_pi, log_prob \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mactor\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43maction_log_prob\u001b[49m\u001b[43m(\u001b[49m\u001b[43mreplay_data\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mobservations\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    207\u001b[0m log_prob \u001b[38;5;241m=\u001b[39m log_prob\u001b[38;5;241m.\u001b[39mreshape(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m    209\u001b[0m ent_coef_loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/.local/lib/python3.8/site-packages/stable_baselines3/sac/policies.py:180\u001b[0m, in \u001b[0;36mActor.action_log_prob\u001b[0;34m(self, obs)\u001b[0m\n\u001b[1;32m    179\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21maction_log_prob\u001b[39m(\u001b[38;5;28mself\u001b[39m, obs: th\u001b[38;5;241m.\u001b[39mTensor) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tuple[th\u001b[38;5;241m.\u001b[39mTensor, th\u001b[38;5;241m.\u001b[39mTensor]:\n\u001b[0;32m--> 180\u001b[0m     mean_actions, log_std, kwargs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_action_dist_params\u001b[49m\u001b[43m(\u001b[49m\u001b[43mobs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    181\u001b[0m     \u001b[38;5;66;03m# return action and associated log prob\u001b[39;00m\n\u001b[1;32m    182\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maction_dist\u001b[38;5;241m.\u001b[39mlog_prob_from_params(mean_actions, log_std, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[0;32m~/.local/lib/python3.8/site-packages/stable_baselines3/sac/policies.py:163\u001b[0m, in \u001b[0;36mActor.get_action_dist_params\u001b[0;34m(self, obs)\u001b[0m\n\u001b[1;32m    155\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    156\u001b[0m \u001b[38;5;124;03mGet the parameters for the action distribution.\u001b[39;00m\n\u001b[1;32m    157\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    160\u001b[0m \u001b[38;5;124;03m    Mean, standard deviation and optional keyword arguments.\u001b[39;00m\n\u001b[1;32m    161\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    162\u001b[0m features \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mextract_features(obs)\n\u001b[0;32m--> 163\u001b[0m latent_pi \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlatent_pi\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfeatures\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    164\u001b[0m mean_actions \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmu(latent_pi)\n\u001b[1;32m    166\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39muse_sde:\n",
      "File \u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1106\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1107\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1108\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1109\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1110\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1111\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1112\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/container.py:141\u001b[0m, in \u001b[0;36mSequential.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    139\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m):\n\u001b[1;32m    140\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m:\n\u001b[0;32m--> 141\u001b[0m         \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mmodule\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    142\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28minput\u001b[39m\n",
      "File \u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1106\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1107\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1108\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1109\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1110\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1111\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1112\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/activation.py:98\u001b[0m, in \u001b[0;36mReLU.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m     97\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Tensor) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tensor:\n\u001b[0;32m---> 98\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrelu\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minplace\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minplace\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/functional.py:1442\u001b[0m, in \u001b[0;36mrelu\u001b[0;34m(input, inplace)\u001b[0m\n\u001b[1;32m   1440\u001b[0m     result \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mrelu_(\u001b[38;5;28minput\u001b[39m)\n\u001b[1;32m   1441\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1442\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrelu\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1443\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "args = parse_args(\"--gym-id SIR_A --seed 1 --learning-starts 1000 --train-freq 5\")\n",
    "run_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\n",
    "env = make_env(args.gym_id, args.seed, 0)()\n",
    "test_env = make_primal_env(args.gym_id)()\n",
    "\n",
    "import torch\n",
    "from stable_baselines3 import SAC\n",
    "from stable_baselines3.common.callbacks import BaseCallback\n",
    "from stable_baselines3.common.logger import TensorBoardOutputFormat\n",
    "\n",
    "class SummaryWriterCallback(BaseCallback):\n",
    "    def _on_training_start(self):\n",
    "        self._log_freq = 292  # log every 1000 calls\n",
    "\n",
    "        output_formats = self.logger.output_formats\n",
    "        # Save reference to tensorboard formatter object\n",
    "        # note: the failure case (not formatter found) is not handled here, should be done with try/except.\n",
    "        self.tb_formatter = next(formatter for formatter in output_formats if isinstance(formatter, TensorBoardOutputFormat))\n",
    "\n",
    "    def _on_step(self) -> bool:\n",
    "        if self.n_calls % self._log_freq == 0:\n",
    "            env_obs = torch.Tensor(self.training_env.reset())\n",
    "            test_obs = torch.Tensor(test_env.reset())\n",
    "            done = False\n",
    "            timestep = 0\n",
    "            total_r = 0\n",
    "            itv_line = []\n",
    "            while not done:\n",
    "                with torch.no_grad():\n",
    "                    action_mean, _ = self.model.predict(env_obs, deterministic=True)\n",
    "                    action_mean = np.array(action_mean)\n",
    "                    test_action_mean = np.clip(np.mean(action_mean, axis=0), 0, 1)\n",
    "                    \n",
    "                test_obs, r, done, _ = test_env.step(test_action_mean)\n",
    "                test_obs = torch.Tensor(test_obs)\n",
    "                itv_index = 0\n",
    "                itv_array = []\n",
    "                for itv in test_env.epi.static.interventions:\n",
    "                    if not itv.is_cost:\n",
    "                        v = float(test_action_mean[itv_index])\n",
    "                        self.tb_formatter.writer.add_scalar('charts/policy_{}/{}'.format(self.num_timesteps, itv.name), v, timestep)\n",
    "                        itv_array.append(v)\n",
    "                        itv_index += 1\n",
    "                itv_line.append(itv_array)\n",
    "                        \n",
    "                env_obs, _, _, _ = self.training_env.step(action_mean)\n",
    "                env_obs = torch.Tensor(env_obs)\n",
    "                \n",
    "                total_r += r\n",
    "                timestep += PERIOD\n",
    "            line = '|'.join([str(self.num_timesteps), str(total_r), str(itv_line)]) + '\\n'\n",
    "            self.tb_formatter.writer.add_scalar(\"charts/learning_curve\", total_r, self.num_timesteps)\n",
    "            print(\"At global step {}, total_rewards={}\".format(self.num_timesteps, total_r))\n",
    "            self.tb_formatter.writer.flush()\n",
    "            csv_file = open('runs/{}_1/records.csv'.format(run_name), 'a')\n",
    "            csv_file.write(line)\n",
    "            csv_file.close()\n",
    "\n",
    "# learning_starts = [1000, 5000, 10000]\n",
    "# target_entropies = [-env.action_space.shape[0], -2*env.action_space.shape[0], -4*env.action_space.shape[0]]\n",
    "# choice = 0\n",
    "# learning_start = learning_starts[choice // 3]\n",
    "# target_entropy = target_entropies[choice % 3]\n",
    "\n",
    "print(f\"Running with {run_name}\")\n",
    "model = SAC(\"MlpPolicy\", env, verbose=1, tensorboard_log=\"runs/\", learning_starts=args.learning_starts, target_entropy=-args.target_entropy_scale * env.action_space.shape[0],\n",
    "           train_freq=args.train_freq, gradient_steps=args.gradient_steps)\n",
    "model.learn(total_timesteps=args.total_timesteps, log_interval=4, callback=SummaryWriterCallback(), tb_log_name=run_name)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e4a732b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9195e3a5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95be68ce",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}