UROB/hw/hw4_RL/walker_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ff36817858940fd7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T18:58:06.054137907Z",
     "start_time": "2024-01-03T18:58:05.833563630Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6c99faf649cd52b",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Teaching a quadruped to walk  \n",
    "\n",
    "Time to try out the learning algorithms that you just implemented on a more difficult problem. The WalkerEnv implements a quadruped robot kind-of thing, see for yourself. The goal is to move in the $x$ direction as fast and as far as possible. \n",
    "\n",
    "Your goal is to implement a class ``WalkerPolicy`` with function ``determine_actions()`` just like the StochasticPolicy we used earlier to control the pendulum. Below is a template of this class, but feel free to alter it however you want. The only important thing is the ``determine_actions()`` function! \n",
    "\n",
    "After you implement it, copy ``WalkerPolicy`` into a separate file ``WalkerPolicy.py`` that you will upload to BRUTE together with the (optional) learned weights in a zip file. How the policy is implemented is up to you! You are constrained to only the libraries we used so far though, such as torch, numpy etc..\n",
    "\n",
    "You will get some free points just for uploading a working policy (irrelevant of the performance). Further 2 points will be awarded for successfully traversing a small distance in the x direction.\n",
    "\n",
    "Do mind, that the evaluation system might uses ``torch==1.12.1``."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d41290d3f9ccf033",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Hints\n",
    "\n",
    "There is no single easy way of doing this, but here are some suggestions on what you could try to improve your policy:\n",
    "\n",
    "1) This problem is much more difficult, than balancing a pendulum. It is a good idea to use a bit larger network than for the pendulum policy.\n",
    "\n",
    "2) You can also try to use a different optimizer, such as Adam and play with the hyperparameters.\n",
    "\n",
    "3) Using a neural network to compute the normal distribution scale $\\sigma$ can lead to too much randomness in the actions (i.e. exploration). You can use a fixed $\\sigma$ instead, or replace it with a learnable ```torch.Parameter``` initialized to some small constant. Make sure, you run it through an exponential, or softplus function to ensure $\\sigma$ is positive.\n",
    "\n",
    "4) The exploration can also be reduced by penalizing the variance of the action distribution in an additional loss term. \n",
    "\n",
    "5) If you see some undesirable behaviour, you can tweak the reward function to penalize it. Even though the $x$ distance is all we care about, adding extra terms to the reward can help guide the learning process (This is known as reward shaping). Simply define a reward function mapping the state $s_{t+1}$ and action $a_t$ to a scalar reward $r_t$ and put it in the config dictionary under the key ```'reward_fcn'```. See the ```WalkerEnv``` class for the implementation of the default reward.\n",
    "\n",
    "6) Using the normal distribution on a bounded action space can lead to certain problems caused by action clipping. This can be mitigated by using a different distribution, such as the Beta distribution. See the ```torch.distributions.beta``` module for more information. (Note that Beta distribution is defined on the interval [0,1] and works better with parameters $\\alpha,\\beta \\geq 1$.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d52d6512e1dc81e6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T18:58:06.260023207Z",
     "start_time": "2024-01-03T18:58:06.056887835Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from environment.WalkerEnv import WalkerEnv\n",
    "from solution import ppo_loss\n",
    "from torch import nn\n",
    "from torch.distributions.normal import Normal\n",
    "from solution import discount_cum_sum\n",
    "from solution import policy_gradient_loss_advantages, value_loss, policy_gradient_loss_discounted "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "bb0dc2f9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T18:58:06.434199241Z",
     "start_time": "2024-01-03T18:58:06.176469019Z"
    }
   },
   "outputs": [],
   "source": [
    "from environment.WalkerEnv import base_config\n",
    "def test_policy(pi, T=512, config=base_config, deterministic=True):\n",
    "    test_env = WalkerEnv(config)\n",
    "    mean_reward = 0\n",
    "    \n",
    "    s = test_env.vector_reset()\n",
    "    for i in range(T):\n",
    "        with torch.no_grad():\n",
    "            if deterministic:\n",
    "                actions = pi.determine_actions(torch.tensor(s).float())  # use deterministic actions based on the states\n",
    "            else:\n",
    "                actions = pi.sample_actions(torch.tensor(s).float())  # use random actions conditioned on the states\n",
    "        s, r = test_env.vector_step(actions.numpy())\n",
    "        mean_reward += sum(r)/(T*config['N'])\n",
    "        \n",
    "    test_env.close()\n",
    "    return mean_reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d1a6be7f9c93e125",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T18:58:06.434366298Z",
     "start_time": "2024-01-03T18:58:06.305533924Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1e318e36d11eb12d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T18:58:06.459079523Z",
     "start_time": "2024-01-03T18:58:06.346839924Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# def walker_reward(state, action):\n",
    "#     \"\"\"reward function for the walker environment, state is [29] vector, action is [8] vector\"\"\"\n",
    "#     pos = state[:15]  # first 15 elements of state vector are generalized coordinates [xyz, quat, joint_angles]\n",
    "#     vel = state[15:]  # last 14 elements of state vector are generalized velocities [xyz_vel, omega, joint_velocities]\n",
    "    \n",
    "#     joint_angles = state[16:24]  # Assuming joint angles are at indices 16 to 23\n",
    "#     # print(joint_angles)\n",
    "#     inactive_legs = [i for i, angle in enumerate(joint_angles) if angle < 0.2]\n",
    "    \n",
    "    \n",
    "    \n",
    "#     return pos[0] + vel[0]# - 0.1 * len(inactive_legs)  # return the x velocity as the reward by default\n",
    "\n",
    "\n",
    "# def sample_trajectories(env, pi, T):    \n",
    "#     \"\"\"given an environment env, a stochastic policy pi and number of timesteps T, interact with the environment for T steps \n",
    "#     using actions sampled from policy. Return torch tensors of collected states, actions and rewards\"\"\"\n",
    "#     states = np.zeros((T + 1, N, env.num_states), dtype=float)  # states from s(0) to s(T+1)\n",
    "#     actions = np.zeros((T, N, env.num_actions), dtype=float)  # actions from a(0) to a(T)\n",
    "#     rewards = np.zeros((T, N), dtype=float)  # rewards from r(0) to r(T)\n",
    "    \n",
    "#     s = env.vector_reset()\n",
    "#     states[0] = s\n",
    "#     for t in range(T):\n",
    "#         a = pi.sample_actions(torch.tensor(states[t]).float())  # policy needs float torch tensor (N, state_dim)\n",
    "#         s_next, r = env.vector_step(np.array(a))  # env needs numpy array of (Nx1)\n",
    "#         states[t + 1], actions[t], rewards[t] = s_next, a, r    \n",
    "        \n",
    "#     tensor_s = torch.tensor(states).float()  # (T+1, N, state_dim)  care for the extra timestep at the end!\n",
    "#     tensor_a = torch.tensor(actions).float()  # (T, N, 1)\n",
    "#     tensor_r = torch.tensor(rewards).float()  # (T, N)\n",
    "    \n",
    "#     return tensor_s, tensor_a, tensor_r \n",
    "\n",
    "# def compute_gae(tensor_r, values, gamma, lambda_):\n",
    "#     \"\"\"generalized advantage estimation (GAE) implementation\"\"\"\n",
    "#     delta_t = tensor_r + gamma * values[1:] - values[:-1]\n",
    "#     advantages = discount_cum_sum(delta_t, gamma * lambda_)\n",
    "#     value_targets = advantages + values[:-1]\n",
    "#     return value_targets, advantages\n",
    "\n",
    "# class WalkerPolicy(nn.Module):\n",
    "#     def __init__(self, state_dim=29, action_dim=8, fixed_sigma=None):\n",
    "#         super().__init__()\n",
    "#         # self.load_weights()  # load learned stored network weights after initialization\n",
    "\n",
    "#         if fixed_sigma is not None:\n",
    "#             self.sigma = torch.nn.Parameter(torch.tensor(fixed_sigma), requires_grad=False)\n",
    "#         else:\n",
    "#             self.sigma = torch.nn.Parameter(torch.tensor(0.01), requires_grad=True)  # Small constant as an exampl\n",
    "\n",
    "#         # shared layers between the action and value heads\n",
    "#         self.shared_layers = nn.Sequential(\n",
    "#             nn.Linear(state_dim, 128),\n",
    "#             nn.Tanh(),\n",
    "#             nn.Linear(128, 64),\n",
    "#             nn.ReLU(),\n",
    "#             nn.Linear(64, 32),\n",
    "#             nn.Tanh(),\n",
    "#         )\n",
    "        \n",
    "#         # action head\n",
    "#         self.action_layers = nn.Sequential(\n",
    "#             nn.Linear(32, 128),\n",
    "#             nn.Tanh(),\n",
    "#             nn.Linear(128,64),\n",
    "#             nn.ReLU(),\n",
    "#             nn.Linear(64, 2*action_dim),\n",
    "#         )\n",
    "\n",
    "#         # value head\n",
    "#         self.value_layers = nn.Sequential(\n",
    "#             nn.Linear(32, 128),\n",
    "#             nn.Tanh(),\n",
    "#             nn.Linear(128,64),      \n",
    "#             nn.ReLU(),\n",
    "#             nn.Linear(64, 1),\n",
    "#         )\n",
    "        \n",
    "#     def determine_actions(self, states):\n",
    "#         \"\"\"states is (N, state_dim) tensor, returns (N, action_dim) actions tensor. \n",
    "#         This function returns deterministic actions based on the input states. This would be used for control. \"\"\"\n",
    "#         params = self.action_layers(self.shared_layers(states))  # map states to distribution parameters\n",
    "#         mu, _ = torch.chunk(params, 2, -1)  # split the parameters into mean and std, return mean\n",
    "#         return mu\n",
    "    \n",
    "#     # TODO: implement a determine_actions() function mapping from (N, state_dim) states into (N, action_dim) actions\n",
    "\n",
    "#     def save_weights(self, path='walker_weights.pt'):\n",
    "#         # helper function to save your network weights\n",
    "#         torch.save(self.state_dict(), path)\n",
    "\n",
    "#     def load_weights(self, path='walker_weights.pt'):\n",
    "#         # helper function to load your network weights\n",
    "#         self.load_state_dict(torch.load(path))\n",
    "    \n",
    "#     def sample_actions(self, states):\n",
    "#         \"\"\"states is (T, N, state_dim) tensor, returns (T, N, action_dim) actions tensor. \n",
    "#         This function returns probabilistically sampled actions. This would be used for training the policy.\"\"\"\n",
    "#         params = self.action_layers(self.shared_layers(states))  # map states to distribution parameters\n",
    "#         mu, sigma = torch.chunk(params, 2, -1)  # split the parameters into mean and std\n",
    "#         # sigma = torch.nn.functional.softplus(self.sigma) + 0.001  # make sure std is positive\n",
    "#         sigma = torch.nn.functional.softplus(sigma)\n",
    "#         distribution = Normal(mu, sigma)  # create distribution of size (T, N, action_dim)\n",
    "#         actions = distribution.sample()  # sample actions\n",
    "#         return actions\n",
    "    \n",
    "#     def log_prob(self, actions, states):\n",
    "#         \"\"\"states is (T, N, state_dim) tensor. actions is (T, N, action_dim) tensor.\n",
    "#         This function returns the log-probabilities of the actions given the states. $\\log \\pi_\\theta(a_t | s_t)$\"\"\"\n",
    "#         params = self.action_layers(self.shared_layers(states))  # map states to distribution parameters\n",
    "#         mu, sigma = torch.chunk(params, 2, -1)  # split the parameters into mean and std\n",
    "#         # sigma = torch.nn.functional.softplus(self.sigma) + 0.001  # make sure std is positive\n",
    "#         sigma = torch.nn.functional.softplus(sigma)\n",
    "#         distribution = Normal(mu, sigma)  # create distribution of size (T, N, action_dim)\n",
    "#         logp = distribution.log_prob(actions)\n",
    "#         if len(logp.shape) == 3 and logp.shape[2] > 1:  # this allows generalization to multi-dim action spaces\n",
    "#             logp = logp.sum(dim=2, keepdim=True)  # sum over the action dimension\n",
    "#         return logp\n",
    "    \n",
    "#     def value_estimates(self, states):\n",
    "#         \"\"\"states is (T, N, state_dim) tensor, returns (T, N) values tensor. Useful for value estimation during training.\"\"\"\n",
    "#         return self.value_layers(self.shared_layers(states)).squeeze()\n",
    "    \n",
    "# # training parameters\n",
    "# N = 32\n",
    "# T = 128\n",
    "# config = {'N': N, 'vis': 1,'reward_fcn':walker_reward}\n",
    "# epochs = 500\n",
    "# lr = 0.003\n",
    "# gamma=0.99\n",
    "# epsilon = 0.2\n",
    "\n",
    "# sgd_iters = 5\n",
    "\n",
    "# # policy, environment and optimizer\n",
    "# pi = WalkerPolicy()#fixed_sigma=0.05)\n",
    "# train_env = WalkerEnv(config)\n",
    "\n",
    "# # optim = torch.optim.SGD(pi.parameters(), lr=lr)\n",
    "# optim = torch.optim.Adam(pi.parameters(),lr=lr)\n",
    "\n",
    "# mean_rewards, p_losses, v_losses = np.zeros(epochs), np.zeros(epochs), np.zeros(epochs)  # for logging mean rewards over epochs\n",
    "# for epoch in range(epochs):\n",
    "#     tensor_s, tensor_a, tensor_r = sample_trajectories(train_env, pi, T)  # collect trajectories using current policy\n",
    "        \n",
    "#     with torch.no_grad():  # compute the old probabilities\n",
    "#         logp_old = pi.log_prob(tensor_a, tensor_s[:T]).squeeze(2)  # compute log(pi(a_t | s_t))\n",
    "    \n",
    "#     for i in range(sgd_iters):  # we can even do multiple gradient steps\n",
    "#         values = pi.value_estimates(tensor_s)  # estimate value function for all states \n",
    "#         logp = pi.log_prob(tensor_a, tensor_s[:T]).squeeze(2)  # compute log(pi(a_t | s_t))\n",
    "\n",
    "#         with torch.no_grad():  # no need for gradients when computing the advantages and value targets\n",
    "#             # value_targets, advantage_estimates = compute_advantage_estimates(tensor_r, values, gamma, bootstrap=True)\n",
    "#             value_targets, advantage_estimates = compute_gae(tensor_r, values, gamma, lambda_=0.97)\n",
    "#             advantage_estimates = (advantage_estimates - advantage_estimates.mean()) / advantage_estimates.std()  # normalize advantages\n",
    "            \n",
    "#         L_v = value_loss(values[:T], value_targets)  # add the value loss\n",
    "        \n",
    "#         p_ratios = torch.exp(logp - logp_old)  # compute the ratios r_\\theta(a_t | s_t)\n",
    "#         L_ppo = ppo_loss(p_ratios, advantage_estimates, epsilon=epsilon)  # compute the policy gradient loss\n",
    "#         total_loss = L_v + L_ppo\n",
    "        \n",
    "#         optim.zero_grad()\n",
    "#         total_loss.backward()  # backprop and gradient step\n",
    "#         optim.step()\n",
    "    \n",
    "#     if epoch % 10 == 0:\n",
    "#         print('Epoch %d, mean reward: %.3f, value loss: %.3f' % (epoch, tensor_r.mean(), L_v.item()))\n",
    "#     mean_rewards[epoch] = tensor_r.mean()\n",
    "#     v_losses[epoch] = L_v.item()\n",
    "#     p_losses[epoch] = L_ppo.item()\n",
    "    \n",
    "# train_env.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0808fe84",
   "metadata": {},
   "source": [
    "### 5.30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "deeefe6c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T19:00:43.704300404Z",
     "start_time": "2024-01-03T18:58:06.444883087Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Environment ready\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "(python:31030): Gtk-WARNING **: 19:58:08.397: gtk_disable_setlocale() must be called before gtk_init()\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, mean reward: -0.032, value loss: 15.924\n",
      "Epoch 10, mean reward: 0.034, value loss: 30.399\n",
      "Epoch 20, mean reward: 0.058, value loss: 62.763\n",
      "Epoch 30, mean reward: 0.044, value loss: 47.541\n",
      "Epoch 40, mean reward: 0.115, value loss: 22.554\n",
      "Epoch 50, mean reward: 0.070, value loss: 18.800\n",
      "Epoch 60, mean reward: 0.042, value loss: 13.170\n",
      "Epoch 70, mean reward: 0.042, value loss: 12.338\n",
      "Epoch 80, mean reward: 0.094, value loss: 17.915\n",
      "Epoch 90, mean reward: 0.088, value loss: 25.348\n",
      "Epoch 100, mean reward: 0.044, value loss: 28.715\n",
      "Epoch 110, mean reward: 0.150, value loss: 28.737\n",
      "Epoch 120, mean reward: 0.079, value loss: 50.370\n",
      "Epoch 130, mean reward: -0.029, value loss: 39.567\n",
      "Epoch 140, mean reward: -0.009, value loss: 52.795\n",
      "Epoch 150, mean reward: -0.046, value loss: 42.176\n",
      "Epoch 160, mean reward: -0.120, value loss: 35.987\n",
      "Epoch 170, mean reward: 0.106, value loss: 19.057\n",
      "Epoch 180, mean reward: -0.032, value loss: 42.672\n",
      "Epoch 190, mean reward: 0.045, value loss: 61.032\n"
     ]
    }
   ],
   "source": [
    "from solution import discount_cum_sum\n",
    "from solution import policy_gradient_loss_advantages, value_loss\n",
    "\n",
    "def compute_advantage_estimates(tensor_r, values, gamma, bootstrap=False):\n",
    "    \"\"\"given reward tensor (T, N), value estimates tensor (T+1, N) and gamma scalar\"\"\"\n",
    "    if bootstrap:  # use last value estimates as a return estimate\n",
    "        terminal_value_estimates = values[-1].unsqueeze(0)  # values of the last states (1, N)\n",
    "        rs_v = torch.cat((tensor_r, terminal_value_estimates), dim=0)\n",
    "        value_targets = discount_cum_sum(rs_v, gamma)[:-1]\n",
    "    else:\n",
    "        value_targets = discount_cum_sum(tensor_r, gamma)\n",
    "    advantages = value_targets - values[:-1]\n",
    "    return value_targets, advantages\n",
    "\n",
    "def compute_gae(tensor_r, values, gamma, lambda_):\n",
    "    \"\"\"generalized advantage estimation (GAE) implementation\"\"\"\n",
    "    delta_t = tensor_r + gamma * values[1:] - values[:-1]\n",
    "    advantages = discount_cum_sum(delta_t, gamma * lambda_)\n",
    "    value_targets = advantages + values[:-1]\n",
    "    return value_targets, advantages\n",
    "\n",
    "def sample_trajectories(env, pi, T):    \n",
    "    \"\"\"given an environment env, a stochastic policy pi and number of timesteps T, interact with the environment for T steps \n",
    "    using actions sampled from policy. Return torch tensors of collected states, actions and rewards\"\"\"\n",
    "    states = np.zeros((T + 1, N, env.num_states), dtype=float)  # states from s(0) to s(T+1)\n",
    "    actions = np.zeros((T, N, env.num_actions), dtype=float)  # actions from a(0) to a(T)\n",
    "    rewards = np.zeros((T, N), dtype=float)  # rewards from r(0) to r(T)\n",
    "    \n",
    "    s = env.vector_reset()\n",
    "    states[0] = s\n",
    "    for t in range(T):\n",
    "        a = pi.sample_actions(torch.tensor(states[t]).float())  # policy needs float torch tensor (N, state_dim)\n",
    "        s_next, r = env.vector_step(np.array(a))  # env needs numpy array of (Nx1)\n",
    "        states[t + 1], actions[t], rewards[t] = s_next, a, r    \n",
    "        \n",
    "    tensor_s = torch.tensor(states).float()  # (T+1, N, state_dim)  care for the extra timestep at the end!\n",
    "    tensor_a = torch.tensor(actions).float()  # (T, N, 1)\n",
    "    tensor_r = torch.tensor(rewards).float()  # (T, N)\n",
    "    \n",
    "    return tensor_s, tensor_a, tensor_r \n",
    "\n",
    "class WalkerPolicy(nn.Module):\n",
    "    def __init__(self, state_dim=29, action_dim=8):\n",
    "        super(WalkerPolicy, self).__init__()\n",
    "        \n",
    "        # shared layers between the action and value heads\n",
    "        self.shared_layers = nn.Sequential(\n",
    "            nn.Linear(state_dim, 32),\n",
    "            nn.Tanh(),\n",
    "        )\n",
    "        \n",
    "        # action head\n",
    "        self.action_layers = nn.Sequential(\n",
    "            nn.Linear(32, 64),\n",
    "            nn.Tanh(),\n",
    "            nn.Linear(64, 2*action_dim),\n",
    "        )\n",
    "\n",
    "        # value head\n",
    "        self.value_layers = nn.Sequential(\n",
    "            nn.Linear(32, 64),\n",
    "            nn.Tanh(),\n",
    "            nn.Linear(64, 1),\n",
    "        )\n",
    "        # self.load_weights()  # load learned stored network weights after initialization\n",
    "        \n",
    "    def determine_actions(self, states):\n",
    "        \"\"\"states is (N, state_dim) tensor, returns (N, action_dim) actions tensor. \n",
    "        This function returns deterministic actions based on the input states. This would be used for control. \"\"\"\n",
    "        params = self.action_layers(self.shared_layers(states))  # map states to distribution parameters\n",
    "        mu, _ = torch.chunk(params, 2, -1)  # split the parameters into mean and std, return mean\n",
    "        return mu\n",
    "    \n",
    "    def save_weights(self, path='walker_weights.pt'):\n",
    "        # helper function to save your network weights\n",
    "        torch.save(self.state_dict(), path)\n",
    "\n",
    "    def load_weights(self, path='walker_weights.pt'):\n",
    "        # helper function to load your network weights\n",
    "        self.load_state_dict(torch.load(path))\n",
    "\n",
    "    def sample_actions(self, states):\n",
    "        \"\"\"states is (T, N, state_dim) tensor, returns (T, N, action_dim) actions tensor. \n",
    "        This function returns probabilistically sampled actions. This would be used for training the policy.\"\"\"\n",
    "        params = self.action_layers(self.shared_layers(states))  # map states to distribution parameters\n",
    "        mu, sigma = torch.chunk(params, 2, -1)  # split the parameters into mean and std\n",
    "        sigma = torch.nn.functional.softplus(sigma)  # make sure std is positive\n",
    "        distribution = Normal(mu, sigma)  # create distribution of size (T, N, action_dim)\n",
    "        actions = distribution.sample()  # sample actions\n",
    "        return actions\n",
    "    \n",
    "    def log_prob(self, actions, states):\n",
    "        \"\"\"states is (T, N, state_dim) tensor. actions is (T, N, action_dim) tensor.\n",
    "        This function returns the log-probabilities of the actions given the states. $\\log \\pi_\\theta(a_t | s_t)$\"\"\"\n",
    "        params = self.action_layers(self.shared_layers(states))  # map states to distribution parameters\n",
    "        mu, sigma = torch.chunk(params, 2, -1)  # split the parameters into mean and std\n",
    "        sigma = torch.nn.functional.softplus(sigma)  # make sure std is positive\n",
    "        distribution = Normal(mu, sigma)  # create distribution of size (T, N, action_dim)\n",
    "        logp = distribution.log_prob(actions)\n",
    "        if len(logp.shape) == 3 and logp.shape[2] > 1:  # this allows generalization to multi-dim action spaces\n",
    "            logp = logp.sum(dim=2, keepdim=True)  # sum over the action dimension\n",
    "        return logp\n",
    "    \n",
    "    def value_estimates(self, states):\n",
    "        \"\"\"states is (T, N, state_dim) tensor, returns (T, N) values tensor. Useful for value estimation during training.\"\"\"\n",
    "        return self.value_layers(self.shared_layers(states)).squeeze()\n",
    "\n",
    "# training parameters\n",
    "N = 32\n",
    "T = 128\n",
    "config = {'N': N, 'vis': 1}\n",
    "epochs = 200\n",
    "lr = 0.01\n",
    "gamma=0.95\n",
    "\n",
    "# policy, environment and optimizer\n",
    "pi = WalkerPolicy()\n",
    "train_env = WalkerEnv(config)\n",
    "# optim = torch.optim.SGD(pi.parameters(), lr=lr)\n",
    "optim = torch.optim.Adam(pi.parameters(), lr=lr)\n",
    "\n",
    "mean_rewards, p_losses, v_losses = np.zeros(epochs), np.zeros(epochs), np.zeros(epochs)  # logging\n",
    "for epoch in range(epochs):\n",
    "    tensor_s, tensor_a, tensor_r = sample_trajectories(train_env, pi, T)  # collect trajectories using current policy\n",
    "    \n",
    "    values = pi.value_estimates(tensor_s)  # estimate value function for all states \n",
    "    logp = pi.log_prob(tensor_a, tensor_s[:T]).squeeze(2)  # compute log(pi(a_t | s_t))\n",
    "\n",
    "    with torch.no_grad():  # no need for gradients when computing the advantages and value targets\n",
    "        value_targets, advantage_estimates = compute_advantage_estimates(tensor_r, values, gamma, bootstrap=True)\n",
    "        # value_targets, advantage_estimates = compute_gae(tensor_r, values, gamma, lambda_=0.99)\n",
    "\n",
    "    L_pg = policy_gradient_loss_advantages(logp, advantage_estimates)  # compute the policy gradient loss\n",
    "    L_v = value_loss(values[:T], value_targets)  # add the value loss\n",
    "    total_loss = L_pg + L_v\n",
    "    \n",
    "    optim.zero_grad()\n",
    "    total_loss.backward()  # backprop and gradient step\n",
    "    optim.step()\n",
    "    \n",
    "    if epoch % 10 == 0:\n",
    "        print('Epoch %d, mean reward: %.3f, value loss: %.3f' % (epoch, tensor_r.mean(), L_v.item()))\n",
    "    mean_rewards[epoch] = tensor_r.mean()\n",
    "    v_losses[epoch] = L_v.item()\n",
    "    p_losses[epoch] = L_pg.item()\n",
    "\n",
    "train_env.close()\n",
    "    \n",
    "# plot_training(mean_rewards, p_losses, v_losses)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c58f8a0bc117e27e",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### 5.08"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d98a9cfd693691f8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T19:00:43.833029104Z",
     "start_time": "2024-01-03T19:00:43.706909067Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# import torch\n",
    "# from torch import nn\n",
    "# from torch.distributions.normal import Normal\n",
    "# from solution import policy_gradient_loss_simple\n",
    "# \n",
    "# \n",
    "# def walker_reward(state, action):\n",
    "#     \"\"\"reward function for the walker environment, state is [29] vector, action is [8] vector\"\"\"\n",
    "#     pos = state[:15]  # first 15 elements of state vector are generalized coordinates [xyz, quat, joint_angles]\n",
    "#     vel = state[15:]  # last 14 elements of state vector are generalized velocities [xyz_vel, omega, joint_velocities]\n",
    "#     \n",
    "#     return pos[0] + vel[0] # return the x velocity as the reward by default\n",
    "# \n",
    "# class WalkerPolicy(nn.Module):\n",
    "#     def __init__(self, state_dim=29, action_dim=8):\n",
    "#         super().__init__()\n",
    "#         # self.load_weights()  # load learned stored network weights after initialization\n",
    "# \n",
    "#         self.network = nn.Sequential(\n",
    "#             nn.Linear(state_dim, 128),\n",
    "#             nn.Tanh(),\n",
    "#             nn.Linear(128, 64),\n",
    "#             nn.Tanh(),\n",
    "#             nn.Linear(64, 2*action_dim),\n",
    "#         )\n",
    "# \n",
    "#     def determine_actions(self, states):\n",
    "#         \"\"\"states is (N, state_dim) tensor, returns (N, action_dim) actions tensor. \n",
    "#         This function returns deterministic actions based on the input states. This would be used for control. \"\"\"\n",
    "#         params = self.network(states)  # map states to distribution parameters\n",
    "#         mu, _ = torch.chunk(params, 2, -1)  # split the parameters into mean and std, return mean\n",
    "#         return mu\n",
    "#     \n",
    "#     def save_weights(self, path='walker_weights.pt'):\n",
    "#         # helper function to save your network weights\n",
    "#         torch.save(self.state_dict(), path)\n",
    "# \n",
    "#     def load_weights(self, path='walker_weights.pt'):\n",
    "#         # helper function to load your network weights\n",
    "#         self.load_state_dict(torch.load(path))\n",
    "#         \n",
    "#     def sample_actions(self, states):\n",
    "#         \"\"\"states is (T, N, state_dim) tensor, returns (T, N, action_dim) actions tensor. \n",
    "#         This function returns probabilistically sampled actions. This would be used for training the policy.\"\"\"\n",
    "#         params = self.network(states)  # map states to distribution parameters\n",
    "#         mu, sigma = torch.chunk(params, 2, -1)  # split the parameters into mean and std\n",
    "#         sigma = torch.nn.functional.softplus(sigma)  # make sure std is positive\n",
    "#         distribution = Normal(mu, sigma)  # create distribution of size (T, N, action_dim)\n",
    "#         actions = distribution.sample()  # sample actions\n",
    "#         return actions\n",
    "# \n",
    "#     def log_prob(self, actions, states):\n",
    "#         \"\"\"states is (T, N, state_dim) tensor. actions is (T, N, action_dim) tensor.\n",
    "#         This function returns the log-probabilities of the actions given the states. $\\log \\pi_\\theta(a_t | s_t)$\"\"\"\n",
    "#         params = self.network(states)  # map states to distribution parameters\n",
    "#         mu, sigma = torch.chunk(params, 2, -1)  # split the parameters into mean and std\n",
    "#         sigma = torch.nn.functional.softplus(sigma)  # make sure std is positive\n",
    "#         distribution = Normal(mu, sigma)  # create distribution of size (T, N, action_dim)\n",
    "#         logp = distribution.log_prob(actions)\n",
    "#         if len(logp.shape) == 3 and logp.shape[2] > 1:  # this allows generalization to multi-dim action spaces\n",
    "#             logp = logp.sum(dim=2, keepdim=True)  # sum over the action dimension\n",
    "#         return logp\n",
    "#     \n",
    "# def sample_trajectories(env, pi, T):    \n",
    "#     \"\"\"given an environment env, a stochastic policy pi and number of timesteps T, interact with the environment for T steps \n",
    "#     using actions sampled from policy. Return torch tensors of collected states, actions and rewards\"\"\"\n",
    "#     states = np.zeros((T + 1, N, env.num_states), dtype=float)  # states from s(0) to s(T+1)\n",
    "#     actions = np.zeros((T, N, env.num_actions), dtype=float)  # actions from a(0) to a(T)\n",
    "#     rewards = np.zeros((T, N), dtype=float)  # rewards from r(0) to r(T)\n",
    "#     \n",
    "#     s = env.vector_reset()\n",
    "#     states[0] = s\n",
    "#     for t in range(T):\n",
    "#         a = pi.sample_actions(torch.tensor(states[t]).float())  # policy needs float torch tensor (N, state_dim)\n",
    "#         s_next, r = env.vector_step(np.array(a))  # env needs numpy array of (Nx1)\n",
    "#         states[t + 1], actions[t], rewards[t] = s_next, a, r    \n",
    "#         \n",
    "#     tensor_s = torch.tensor(states).float()  # (T+1, N, state_dim)  care for the extra timestep at the end!\n",
    "#     tensor_a = torch.tensor(actions).float()  # (T, N, 1)\n",
    "#     tensor_r = torch.tensor(rewards).float()  # (T, N)\n",
    "#     \n",
    "#     return tensor_s, tensor_a, tensor_r\n",
    "# \n",
    "# # training parameters\n",
    "# N = 32\n",
    "# T = 128\n",
    "# config = {'N': N, 'vis': 1, \"reward_fcn\":walker_reward}\n",
    "# epochs = 500\n",
    "# lr = 0.01\n",
    "# gamma=0.99\n",
    "# \n",
    "# # policy, environment and optimizer\n",
    "# pi = WalkerPolicy()\n",
    "# train_env = WalkerEnv(config)\n",
    "# # optim = torch.optim.SGD(pi.parameters(), lr=lr)\n",
    "# optim = torch.optim.Adam(pi.parameters(), lr=lr)\n",
    "# \n",
    "# mean_rewards, p_losses = np.zeros(epochs), np.zeros(epochs)\n",
    "# for epoch in range(epochs):\n",
    "#     tensor_s, tensor_a, tensor_r = sample_trajectories(train_env, pi, T)  # collect trajectories using current policy\n",
    "# \n",
    "#     logp = pi.log_prob(tensor_a, tensor_s[:T]).squeeze(2)  # compute log(pi(a_t | s_t))\n",
    "#     loss = policy_gradient_loss_discounted(logp, tensor_r, gamma)  # compute the policy gradient loss\n",
    "#     \n",
    "#     optim.zero_grad()\n",
    "#     loss.backward()  # backprop and gradient step\n",
    "#     optim.step()\n",
    "#     \n",
    "#     if epoch % 10 == 0:\n",
    "#         print('Epoch %d, mean reward: %.3f' % (epoch, tensor_r.mean()))\n",
    "#     mean_rewards[epoch] = tensor_r.mean()\n",
    "#     p_losses[epoch] = loss.item()\n",
    "#     \n",
    "# train_env.close()\n",
    "# \n",
    "# pi.save_weights()\n",
    "# # plot_training(mean_rewards, p_losses)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36072da0",
   "metadata": {},
   "source": [
    "### 3.42 m:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b884ffe7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T19:00:43.925588246Z",
     "start_time": "2024-01-03T19:00:43.841426215Z"
    }
   },
   "outputs": [],
   "source": [
    "# import torch\n",
    "# from torch import nn\n",
    "# from torch.distributions.normal import Normal\n",
    "# from solution import policy_gradient_loss_simple\n",
    "\n",
    "# class WalkerPolicy(nn.Module):\n",
    "#     def __init__(self, state_dim=29, action_dim=8):\n",
    "#         super().__init__()\n",
    "#         # self.load_weights()  # load learned stored network weights after initialization\n",
    "\n",
    "#         self.network = nn.Sequential(\n",
    "#             nn.Linear(state_dim, 128),\n",
    "#             nn.Tanh(),\n",
    "#             nn.Linear(128, 64),\n",
    "#             nn.Tanh(),\n",
    "#             nn.Linear(64, 2*action_dim),\n",
    "#         )\n",
    "\n",
    "#     def determine_actions(self, states):\n",
    "#         \"\"\"states is (N, state_dim) tensor, returns (N, action_dim) actions tensor. \n",
    "#         This function returns deterministic actions based on the input states. This would be used for control. \"\"\"\n",
    "#         params = self.network(states)  # map states to distribution parameters\n",
    "#         mu, _ = torch.chunk(params, 2, -1)  # split the parameters into mean and std, return mean\n",
    "#         return mu\n",
    "    \n",
    "#     def save_weights(self, path='walker_weights.pt'):\n",
    "#         # helper function to save your network weights\n",
    "#         torch.save(self.state_dict(), path)\n",
    "\n",
    "#     def load_weights(self, path='walker_weights.pt'):\n",
    "#         # helper function to load your network weights\n",
    "#         self.load_state_dict(torch.load(path))\n",
    "        \n",
    "#     def sample_actions(self, states):\n",
    "#         \"\"\"states is (T, N, state_dim) tensor, returns (T, N, action_dim) actions tensor. \n",
    "#         This function returns probabilistically sampled actions. This would be used for training the policy.\"\"\"\n",
    "#         params = self.network(states)  # map states to distribution parameters\n",
    "#         mu, sigma = torch.chunk(params, 2, -1)  # split the parameters into mean and std\n",
    "#         sigma = torch.nn.functional.softplus(sigma)  # make sure std is positive\n",
    "#         distribution = Normal(mu, sigma)  # create distribution of size (T, N, action_dim)\n",
    "#         actions = distribution.sample()  # sample actions\n",
    "#         return actions\n",
    "\n",
    "#     def log_prob(self, actions, states):\n",
    "#         \"\"\"states is (T, N, state_dim) tensor. actions is (T, N, action_dim) tensor.\n",
    "#         This function returns the log-probabilities of the actions given the states. $\\log \\pi_\\theta(a_t | s_t)$\"\"\"\n",
    "#         params = self.network(states)  # map states to distribution parameters\n",
    "#         mu, sigma = torch.chunk(params, 2, -1)  # split the parameters into mean and std\n",
    "#         sigma = torch.nn.functional.softplus(sigma)  # make sure std is positive\n",
    "#         distribution = Normal(mu, sigma)  # create distribution of size (T, N, action_dim)\n",
    "#         logp = distribution.log_prob(actions)\n",
    "#         if len(logp.shape) == 3 and logp.shape[2] > 1:  # this allows generalization to multi-dim action spaces\n",
    "#             logp = logp.sum(dim=2, keepdim=True)  # sum over the action dimension\n",
    "#         return logp\n",
    "    \n",
    "# def sample_trajectories(env, pi, T):    \n",
    "#     \"\"\"given an environment env, a stochastic policy pi and number of timesteps T, interact with the environment for T steps \n",
    "#     using actions sampled from policy. Return torch tensors of collected states, actions and rewards\"\"\"\n",
    "#     states = np.zeros((T + 1, N, env.num_states), dtype=float)  # states from s(0) to s(T+1)\n",
    "#     actions = np.zeros((T, N, env.num_actions), dtype=float)  # actions from a(0) to a(T)\n",
    "#     rewards = np.zeros((T, N), dtype=float)  # rewards from r(0) to r(T)\n",
    "    \n",
    "#     s = env.vector_reset()\n",
    "#     states[0] = s\n",
    "#     for t in range(T):\n",
    "#         a = pi.sample_actions(torch.tensor(states[t]).float())  # policy needs float torch tensor (N, state_dim)\n",
    "#         s_next, r = env.vector_step(np.array(a))  # env needs numpy array of (Nx1)\n",
    "#         states[t + 1], actions[t], rewards[t] = s_next, a, r    \n",
    "        \n",
    "#     tensor_s = torch.tensor(states).float()  # (T+1, N, state_dim)  care for the extra timestep at the end!\n",
    "#     tensor_a = torch.tensor(actions).float()  # (T, N, 1)\n",
    "#     tensor_r = torch.tensor(rewards).float()  # (T, N)\n",
    "    \n",
    "#     return tensor_s, tensor_a, tensor_r\n",
    "\n",
    "# # training parameters\n",
    "# N = 32\n",
    "# T = 128\n",
    "# config = {'N': N, 'vis': 1}\n",
    "# epochs = 500\n",
    "# lr = 0.01\n",
    "\n",
    "# # policy, environment and optimizer\n",
    "# pi = WalkerPolicy()\n",
    "# train_env = WalkerEnv(config)\n",
    "# optim = torch.optim.SGD(pi.parameters(), lr=lr)\n",
    "\n",
    "# mean_rewards, p_losses = np.zeros(epochs), np.zeros(epochs)  # for logging mean rewards over epochs\n",
    "# for epoch in range(epochs):\n",
    "#     tensor_s, tensor_a, tensor_r = sample_trajectories(train_env, pi, T)  # collect trajectories using current policy\n",
    "    \n",
    "#     logp = pi.log_prob(tensor_a, tensor_s[:T]).squeeze(2)  # compute log(pi(a_t | s_t))\n",
    "#     loss = policy_gradient_loss_simple(logp, tensor_r)  # compute the policy gradient loss\n",
    "    \n",
    "#     optim.zero_grad()\n",
    "#     loss.backward()  # backprop and gradient step\n",
    "#     optim.step()\n",
    "    \n",
    "#     if epoch % 10 == 0:\n",
    "#         print('Epoch %d, mean reward: %.3f' % (epoch, tensor_r.mean()))\n",
    "#     mean_rewards[epoch] = tensor_r.mean()\n",
    "#     p_losses[epoch] = loss.item()\n",
    "    \n",
    "# train_env.close()\n",
    "\n",
    "# # pi.save_weights()\n",
    "# # plot_training(mean_rewards, p_losses)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a372f7d",
   "metadata": {},
   "source": [
    "### Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e64a34ffdb26d39b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T19:00:44.017254773Z",
     "start_time": "2024-01-03T19:00:43.926176998Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# config = {'N': 1, 'vis': 1}\n",
    "# env = WalkerEnv(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "f1159688c1501d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T19:00:44.107849028Z",
     "start_time": "2024-01-03T19:00:44.017644812Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# obs = env.vector_reset()\n",
    "# env.render()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8017d2863ead4fbb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T19:00:44.199294126Z",
     "start_time": "2024-01-03T19:00:44.109103721Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# obs[0, 0]  # this is the x coordinate of the robot, we want to maximize this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a9eb6e91ef52c5de",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T19:00:44.301341461Z",
     "start_time": "2024-01-03T19:00:44.199739539Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# T = 512\n",
    "# for i in range(512):\n",
    "    # tensor_s, tensor_a, tensor_r = sample_trajectories(train_env, pi, T)  # collect trajectories using\n",
    "    # values = pi.value_estimates(tensor_s)  # estimate value function for all states\n",
    "    # logp = pi.log_prob(tensor_a, tensor_s[:T]).squeeze(2)  # compute log(pi(a_t | s_t))\n",
    "    \n",
    "#     a = np.random.randn(1, 8)\n",
    "#     obs, reward = env.vector_step(a)\n",
    "# env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "b6120ca3cad7b4b6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T19:00:44.352542256Z",
     "start_time": "2024-01-03T19:00:44.303403881Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ba0bdbb9ec0d76d7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T19:00:44.397562035Z",
     "start_time": "2024-01-03T19:00:44.345438911Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "7885edb6aa1a9ea8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T19:00:53.530161876Z",
     "start_time": "2024-01-03T19:00:44.397735800Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Environment ready\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_31030/2741340417.py:10: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at ../torch/csrc/utils/tensor_new.cpp:261.)\n",
      "  actions = pi.determine_actions(torch.tensor(s).float())  # use deterministic actions based on the states\n",
      "\n",
      "(python:31030): Gtk-WARNING **: 20:00:44.656: gtk_disable_setlocale() must be called before gtk_init()\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.15915431915709632"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_policy(pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "ed6b551c19e2ac09",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-03T19:00:53.531089413Z",
     "start_time": "2024-01-03T19:00:53.511456034Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}