3D quad race INDI inner loop.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import stable_baselines3\n",
    "import sys\n",
    "import numpy as np\n",
    "\n",
    "print(\"python version:\", sys.version)\n",
    "print(\"stable_baselines3 version:\", stable_baselines3.__version__)\n",
    "print(\"torch version:\", torch.__version__)\n",
    "print(\"cuda available:\", torch.cuda.is_available())\n",
    "print(\"cuda version:\", torch.version.cuda)\n",
    "print(\"cudnn version:\", torch.backends.cudnn.version())\n",
    "\n",
    "# set device\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(\"device:\", device)\n",
    "\n",
    "# set torch default device\n",
    "torch.set_default_device(device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Equation of Motion 3D Quadcopter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reload sympy\n",
    "from sympy import *\n",
    "\n",
    "# Equations of motion 3D quadcopter + INDI inner loop\n",
    "\n",
    "state = symbols('x y z v_x v_y v_z phi theta psi p q r T_norm')\n",
    "x,y,z,vx,vy,vz,phi,theta,psi,p,q,r,T_norm = state\n",
    "# commands normalized to [-1,1]\n",
    "control = symbols('p_cmd,q_cmd,r_cmd,T_cmd')\n",
    "p_cmd,q_cmd,r_cmd,T_cmd = control\n",
    "\n",
    "# Physics\n",
    "g = 9.81\n",
    "\n",
    "# INDI params\n",
    "taup = 0.03\n",
    "tauq = 0.03\n",
    "taur = 0.03\n",
    "tauT = 0.03\n",
    "\n",
    "p_min =-3.0 #-6.0\n",
    "p_max = 3.0 #6.0\n",
    "q_min =-3.0 #-6.0\n",
    "q_max = 3.0 #6.0\n",
    "r_min =-2.0\n",
    "r_max = 2.0\n",
    "T_min = 0.0\n",
    "T_max = 16.0\n",
    "\n",
    "# Drag model\n",
    "k_x = 0.33915248\n",
    "k_y = 0.4314916\n",
    "\n",
    "# Rotation matrix \n",
    "Rx = Matrix([[1, 0, 0], [0, cos(phi), -sin(phi)], [0, sin(phi), cos(phi)]])\n",
    "Ry = Matrix([[cos(theta), 0, sin(theta)], [0, 1, 0], [-sin(theta), 0, cos(theta)]])\n",
    "Rz = Matrix([[cos(psi), -sin(psi), 0], [sin(psi), cos(psi), 0], [0, 0, 1]])\n",
    "R = Rz*Ry*Rx\n",
    "\n",
    "# Body velocity\n",
    "vbx, vby, vbz = R.T@Matrix([vx,vy,vz])\n",
    "\n",
    "# Drag\n",
    "Dx = -k_x*vbx\n",
    "Dy = -k_y*vby\n",
    "\n",
    "# Dynamics\n",
    "d_x = vx\n",
    "d_y = vy\n",
    "d_z = vz\n",
    "\n",
    "T = (T_norm+1)/2*(T_max-T_min)+T_min\n",
    "d_vx, d_vy, d_vz = Matrix([0,0,g]) + R@Matrix([Dx, Dy,-T])\n",
    "\n",
    "d_phi   = p + q*sin(phi)*tan(theta) + r*cos(phi)*tan(theta)\n",
    "d_theta = q*cos(phi) - r*sin(phi)\n",
    "d_psi   = q*sin(phi)/cos(theta) + r*cos(phi)/cos(theta)\n",
    "\n",
    "d_p      = ((p_cmd+1)/2*(p_max-p_min)+p_min - p)/taup\n",
    "d_q      = ((q_cmd+1)/2*(q_max-q_min)+q_min - q)/tauq\n",
    "d_r      = ((r_cmd+1)/2*(r_max-r_min)+r_min - r)/taur\n",
    "d_T_norm = (T_cmd - T_norm)/tauT\n",
    "\n",
    "# State space model\n",
    "f = [d_x, d_y, d_z, d_vx, d_vy, d_vz, d_phi, d_theta, d_psi, d_p, d_q, d_r, d_T_norm]\n",
    "\n",
    "# lambdify\n",
    "f_func = lambdify((Array(state), Array(control)), Array(f), 'numpy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "axs, ays, azs = R*Matrix([0,0,x])\n",
    "print(ccode(axs))\n",
    "print(ccode(ays))\n",
    "print(ccode(azs))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define the 3D Quad Gates Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Efficient vectorized version of the environment\n",
    "from gym import spaces\n",
    "from stable_baselines3.common.vec_env import VecEnv\n",
    "\n",
    "class Quadcopter3DGates(VecEnv):\n",
    "    def __init__(self,\n",
    "                 num_envs,\n",
    "                 gates_pos,\n",
    "                 gate_yaw,\n",
    "                 start_pos,\n",
    "                 gates_ahead=0,\n",
    "                 pause_if_collision=False,\n",
    "                 ):\n",
    "        \n",
    "        # Define the race track\n",
    "        self.start_pos = start_pos.astype(np.float32)\n",
    "        self.gate_pos = gates_pos.astype(np.float32)\n",
    "        self.gate_yaw = gate_yaw.astype(np.float32)\n",
    "        self.num_gates = gates_pos.shape[0]\n",
    "        self.gates_ahead = gates_ahead\n",
    "        \n",
    "        # Pause if collision\n",
    "        self.pause_if_collision = pause_if_collision\n",
    "\n",
    "        # Calculate relative gates\n",
    "        # pos,yaw of gate i in reference frame of gate i-1 (assumes a looped track)\n",
    "        self.gate_pos_rel = np.zeros((self.num_gates,3), dtype=np.float32)\n",
    "        self.gate_yaw_rel = np.zeros(self.num_gates, dtype=np.float32)\n",
    "        for i in range(0,self.num_gates):\n",
    "            self.gate_pos_rel[i] = self.gate_pos[i] - self.gate_pos[i-1]\n",
    "            # Rotation matrix\n",
    "            R = np.array([\n",
    "                [np.cos(self.gate_yaw[i-1]), np.sin(self.gate_yaw[i-1])],\n",
    "                [-np.sin(self.gate_yaw[i-1]), np.cos(self.gate_yaw[i-1])]\n",
    "            ])\n",
    "            self.gate_pos_rel[i,0:2] = R@self.gate_pos_rel[i,0:2]\n",
    "            self.gate_yaw_rel[i] = self.gate_yaw[i] - self.gate_yaw[i-1]\n",
    "\n",
    "        # Define the target gate for each environment\n",
    "        self.target_gates = np.zeros(num_envs, dtype=int)\n",
    "\n",
    "        # action space: [p_cmd,q_cmd,r_cmd,T_cmd] normalized to [-1,1]\n",
    "        action_space = spaces.Box(low=-1, high=1, shape=(4,))\n",
    "\n",
    "        # observation space: pos[G], vel[G], att[eulerB->G], rates[B], rpms, future_gates[G], future_gate_dirs[G]\n",
    "        # [G] = reference frame aligned with target gate\n",
    "        # [B] = body frame\n",
    "        self.state_len = 13+4*self.gates_ahead\n",
    "        observation_space = spaces.Box(\n",
    "            low  = np.array([-np.inf]*self.state_len),\n",
    "            high = np.array([ np.inf]*self.state_len)\n",
    "        )\n",
    "\n",
    "        # Initialize the VecEnv\n",
    "        VecEnv.__init__(self,num_envs, observation_space, action_space)\n",
    "\n",
    "        # world state: pos[W], vel[W], att[eulerB->W], rates[B], rpms\n",
    "        self.world_states = np.zeros((num_envs,13), dtype=np.float32)\n",
    "        # observation state\n",
    "        self.states = np.zeros((num_envs,self.state_len), dtype=np.float32)\n",
    "\n",
    "        # Define any other environment-specific parameters\n",
    "        self.max_steps = 1200      # Maximum number of steps in an episode\n",
    "        self.dt = np.float32(0.01) # Time step duration\n",
    "\n",
    "        self.step_counts = np.zeros(num_envs, dtype=int)\n",
    "        self.actions = np.zeros((num_envs,4), dtype=np.float32)\n",
    "        self.dones = np.zeros(num_envs, dtype=bool)\n",
    "        self.final_gate_passed = np.zeros(num_envs, dtype=bool)\n",
    "\n",
    "        self.update_states = self.update_states_gate\n",
    "\n",
    "        self.pause = False\n",
    "\n",
    "    def update_states_world(self):\n",
    "        self.states = self.world_states\n",
    "\n",
    "    def update_states_gate(self):\n",
    "        # Transform pos and vel in gate frame\n",
    "        gate_pos = self.gate_pos[self.target_gates%self.num_gates]\n",
    "        gate_yaw = self.gate_yaw[self.target_gates%self.num_gates]\n",
    "\n",
    "        # Rotation matrix from world frame to gate frame\n",
    "        R = np.array([\n",
    "            [np.cos(gate_yaw), np.sin(gate_yaw)],\n",
    "            [-np.sin(gate_yaw), np.cos(gate_yaw)]\n",
    "        ]).transpose((2,1,0))\n",
    "\n",
    "        # new state array to prevent the weird bug related to indexing ([:] syntax)\n",
    "        new_states = np.zeros_like(self.states)\n",
    "\n",
    "        # Update positions\n",
    "        pos_W = self.world_states[:,0:3]\n",
    "        pos_G = (pos_W[:,np.newaxis,0:2] - gate_pos[:,np.newaxis,0:2]) @ R\n",
    "        new_states[:,0:2] = pos_G[:,0,:]\n",
    "        new_states[:,2] = pos_W[:,2] - gate_pos[:,2]\n",
    "\n",
    "        # Update velocities\n",
    "        vel_W = self.world_states[:,3:6]\n",
    "        vel_G = (vel_W[:,np.newaxis,0:2]) @ R\n",
    "        new_states[:,3:5] = vel_G[:,0,:]\n",
    "        new_states[:,5] = vel_W[:,2]\n",
    "\n",
    "        # Update attitude\n",
    "        new_states[:,6:8] = self.world_states[:,6:8]\n",
    "        yaw = self.world_states[:,8] - gate_yaw\n",
    "        yaw %= 2*np.pi\n",
    "        yaw[yaw > np.pi] -= 2*np.pi\n",
    "        yaw[yaw < -np.pi] += 2*np.pi\n",
    "        new_states[:,8] = yaw\n",
    "\n",
    "        # Update rates\n",
    "        new_states[:,9:12] = self.world_states[:,9:12]\n",
    "\n",
    "        # Update Thrust\n",
    "        new_states[:,12] = self.world_states[:,12]\n",
    "\n",
    "        # Update future gates relative to current gate ([0,0,0,0] for out of bounds)\n",
    "        for i in range(self.gates_ahead):\n",
    "            indices = (self.target_gates+i+1)\n",
    "            # loop when out of bounds\n",
    "            indices = indices % self.num_gates\n",
    "            valid = indices < self.num_gates\n",
    "            new_states[valid,13+4*i:13+4*i+3] = self.gate_pos_rel[indices[valid]]\n",
    "            new_states[valid,13+4*i+3] = self.gate_yaw_rel[indices[valid]]\n",
    "\n",
    "        self.states = new_states\n",
    "\n",
    "    def reset_(self, dones):\n",
    "        num_reset = dones.sum()\n",
    "\n",
    "        x0 = np.random.uniform(-0.5,0.5, size=(num_reset,)) + self.start_pos[0]\n",
    "        y0 = np.random.uniform(-0.5,0.5, size=(num_reset,)) + self.start_pos[1]\n",
    "        z0 = np.random.uniform(-0.5,0.5, size=(num_reset,)) + self.start_pos[2]\n",
    "        \n",
    "        vx0 = np.random.uniform(-0.5,0.5, size=(num_reset,))\n",
    "        vy0 = np.random.uniform(-0.5,0.5, size=(num_reset,))\n",
    "        vz0 = np.random.uniform(-0.5,0.5, size=(num_reset,))\n",
    "        \n",
    "        phi0   = np.random.uniform(-np.pi/9,np.pi/9, size=(num_reset,))\n",
    "        theta0 = np.random.uniform(-np.pi/9,np.pi/9, size=(num_reset,))\n",
    "        psi0   = np.random.uniform(-np.pi,np.pi, size=(num_reset,))\n",
    "        \n",
    "        p0 = np.random.uniform(-0.1,0.1, size=(num_reset,))\n",
    "        q0 = np.random.uniform(-0.1,0.1, size=(num_reset,))\n",
    "        r0 = np.random.uniform(-0.1,0.1, size=(num_reset,))\n",
    "        \n",
    "        T0 = np.random.uniform(-.1,.1, size=(num_reset,))\n",
    "\n",
    "        self.world_states[dones] = np.stack([x0, y0, z0, vx0, vy0, vz0, phi0, theta0, psi0, p0, q0, r0, T0], axis=1)\n",
    "\n",
    "        self.step_counts[dones] = np.zeros(num_reset)\n",
    "        \n",
    "        self.target_gates[dones] = np.zeros(num_reset, dtype=int)\n",
    "\n",
    "        # update states\n",
    "        self.update_states()\n",
    "        return self.states\n",
    "    \n",
    "    def reset(self):\n",
    "        return self.reset_(np.ones(self.num_envs, dtype=bool))\n",
    "\n",
    "    def step_async(self, actions):\n",
    "        self.actions = actions\n",
    "    \n",
    "    def step_wait(self):\n",
    "        new_states = self.world_states + self.dt*f_func(self.world_states.T, self.actions.T).T\n",
    "        self.step_counts += 1\n",
    "\n",
    "        pos_old = self.world_states[:,0:3]\n",
    "        pos_new = new_states[:,0:3]\n",
    "        pos_gate = self.gate_pos[self.target_gates%self.num_gates]\n",
    "        yaw_gate = self.gate_yaw[self.target_gates%self.num_gates]\n",
    "\n",
    "        # Rewards from [...]\n",
    "        d2g_old = np.linalg.norm(pos_old - pos_gate, axis=1)\n",
    "        d2g_new = np.linalg.norm(pos_new - pos_gate, axis=1)\n",
    "        rat_penalty = 0*0.01*np.linalg.norm(new_states[:,9:12], axis=1)\n",
    "        rewards = d2g_old - d2g_new - rat_penalty\n",
    "        \n",
    "        normal = np.array([np.cos(yaw_gate), np.sin(yaw_gate)]).T\n",
    "        # dot product of normal and position vector over axis 1\n",
    "        pos_old_projected = (pos_old[:,0]-pos_gate[:,0])*normal[:,0] + (pos_old[:,1]-pos_gate[:,1])*normal[:,1]\n",
    "        pos_new_projected = (pos_new[:,0]-pos_gate[:,0])*normal[:,0] + (pos_new[:,1]-pos_gate[:,1])*normal[:,1]\n",
    "        passed_gate_plane = (pos_old_projected < 0) & (pos_new_projected > 0)\n",
    "        gate_passed = passed_gate_plane & np.all(np.abs(pos_new - pos_gate)<0.5, axis=1)\n",
    "        gate_collision = passed_gate_plane & np.any(np.abs(pos_new - pos_gate)>0.5, axis=1)\n",
    "        \n",
    "        # gate reward (+ dist penalty)\n",
    "        rewards[gate_passed] = 10 - 10*d2g_new[gate_passed]\n",
    "        \n",
    "        # Check for gate collision\n",
    "        rewards[gate_collision] = -10\n",
    "\n",
    "        # Check ground collision (z > 0)\n",
    "        ground_collision = new_states[:,2] > 0\n",
    "        rewards[ground_collision] = -10\n",
    "        \n",
    "        # Check out of bounds\n",
    "        # outside grid abs(x,y)>10\n",
    "        # prevent numerical issues: abs(p,q,r) < 1000\n",
    "        out_of_bounds = np.any(np.abs(new_states[:,0:2]) > 10, axis=1) | np.any(np.abs(new_states[:,9:12]) > 1000, axis=1)\n",
    "        rewards[out_of_bounds] = -10\n",
    "        \n",
    "        # Check number of steps\n",
    "        max_steps_reached = self.step_counts >= self.max_steps\n",
    "\n",
    "        # Update target gate\n",
    "        self.target_gates[gate_passed] += 1\n",
    "        self.target_gates[gate_passed] %= self.num_gates\n",
    "        \n",
    "        # Check if final gate has been passed\n",
    "        # self.final_gate_passed = self.target_gates >= self.num_gates\n",
    "\n",
    "        # give reward for passing final gate\n",
    "        rewards[self.final_gate_passed] = 10\n",
    "        \n",
    "        # Check if the episode is done\n",
    "        dones = max_steps_reached | ground_collision | gate_collision | out_of_bounds #| self.final_gate_passed\n",
    "        self.dones = dones\n",
    "        \n",
    "        # Pause if collision\n",
    "        if self.pause:\n",
    "            dones = dones & ~dones\n",
    "            self.dones = dones\n",
    "        elif self.pause_if_collision:\n",
    "            # dones = max_steps_reached | final_gate_passed | out_of_bounds\n",
    "            update = ~dones #~(gate_collision | ground_collision)\n",
    "            # Update world states\n",
    "            self.world_states[update] = new_states[update]\n",
    "            self.update_states()\n",
    "            # Reset env if done (and update states)\n",
    "            # self.reset_(dones)\n",
    "        else:\n",
    "            # Update world states\n",
    "            self.world_states = new_states\n",
    "            # reset env if done (and update states)\n",
    "            self.reset_(dones)\n",
    "\n",
    "\n",
    "        # Write info dicts\n",
    "        infos = [{}] * self.num_envs\n",
    "        for i in range(self.num_envs):\n",
    "            if dones[i]:\n",
    "                infos[i][\"terminal_observation\"] = self.states[i]\n",
    "            if max_steps_reached[i]:\n",
    "                infos[i][\"TimeLimit.truncated\"] = True\n",
    "        return self.states, rewards, dones, infos\n",
    "    \n",
    "    def close(self):\n",
    "        pass\n",
    "\n",
    "    def seed(self, seed=None):\n",
    "        pass\n",
    "\n",
    "    def get_attr(self, attr_name, indices=None):\n",
    "        pass\n",
    "\n",
    "    def set_attr(self, attr_name, value, indices=None):\n",
    "        pass\n",
    "\n",
    "    def env_method(self, method_name, *method_args, indices=None, **method_kwargs):\n",
    "        pass\n",
    "\n",
    "    def env_is_wrapped(self, wrapper_class, indices=None):\n",
    "        return [False]*self.num_envs\n",
    "\n",
    "    def render(self, mode='human'):\n",
    "        # Outputs a dict containing all information for rendering\n",
    "        state_dict = dict(zip(['x','y','z','vx','vy','vz','phi','theta','psi','p','q','r','T'], self.world_states.T))\n",
    "        # Rescale actions to [0,1] for rendering\n",
    "        action_dict = dict(zip(['u1','u2','u3','u4'], (np.array(self.actions.T)+1)/2))\n",
    "        return {**state_dict, **action_dict}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define Race Track"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib\n",
    "from quadcopter_animation import animation\n",
    "importlib.reload(animation)\n",
    "\n",
    "# Define the race track\n",
    "# gate_pos = np.array([\n",
    "#     [-1.5,-2,-1.5],\n",
    "#     [1.5,2,-1.5],\n",
    "#     [1.5,-2,-2.5],\n",
    "#     [-1.5,2,-1.5]\n",
    "# ]*2)\n",
    "\n",
    "gate_pos = np.array([\n",
    "    [ 2,-1.5,-1.5],\n",
    "    [ 2, 1.5,-1.5],\n",
    "    [-2, 1.5,-1.5],\n",
    "    [-2,-1.5,-1.5]\n",
    "]*2)\n",
    "\n",
    "# gate_yaw = np.array([\n",
    "#     0,\n",
    "#     0,\n",
    "#     np.pi,\n",
    "#     np.pi\n",
    "# ]*2)\n",
    "gate_yaw = np.array([\n",
    "    np.pi/4,\n",
    "    3*np.pi/4,\n",
    "    5*np.pi/4,\n",
    "    7*np.pi/4\n",
    "]*2)\n",
    "\n",
    "\n",
    "# start_pos = gate_pos[0] - np.array([2,0,0])\n",
    "start_pos = gate_pos[3] #- np.array([2,0,0])\n",
    "\n",
    "num = 10\n",
    "env = Quadcopter3DGates(num_envs=num, gates_pos=gate_pos, gate_yaw=gate_yaw, start_pos=start_pos, pause_if_collision=False)\n",
    "\n",
    "# Run a random agent\n",
    "env.reset()\n",
    "\n",
    "done = False\n",
    "def run():\n",
    "    global done\n",
    "    action = np.random.uniform(-1,1, size=(num,4))\n",
    "    # action[:,3] = 1\n",
    "    state, reward, done, _ = env.step(action)\n",
    "    # print(state[0][-4:])\n",
    "    # print(env.disturbances[0])\n",
    "    if reward[0] > 1:\n",
    "        print(\"reward:\", reward)\n",
    "    return env.render()\n",
    "\n",
    "animation.view(run, gate_pos=gate_pos, gate_yaw=gate_yaw) #, record_steps=1000, show_window=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train PPO Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from stable_baselines3 import PPO\n",
    "from datetime import datetime\n",
    "from stable_baselines3.common.vec_env import VecMonitor\n",
    "import importlib\n",
    "from quadcopter_animation import animation\n",
    "\n",
    "models_dir = 'models/PPO_INDI'\n",
    "log_dir = 'logs/PPO_INDI'\n",
    "video_log_dir = 'videos/PPO_INDI'\n",
    "\n",
    "if not os.path.exists(models_dir):\n",
    "    os.makedirs(models_dir)\n",
    "if not os.path.exists(log_dir):\n",
    "    os.makedirs(log_dir)\n",
    "if not os.path.exists(video_log_dir):\n",
    "    os.makedirs(video_log_dir)\n",
    "\n",
    "# Date and time string for unique folder names\n",
    "datetime_str = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "\n",
    "# Create the environment\n",
    "env = Quadcopter3DGates(num_envs=100, gates_pos=gate_pos, gate_yaw=gate_yaw, start_pos=start_pos, gates_ahead=1)\n",
    "test_env = Quadcopter3DGates(num_envs=10, gates_pos=gate_pos, gate_yaw=gate_yaw, start_pos=start_pos, gates_ahead=1, pause_if_collision=True)\n",
    "\n",
    "# Wrap the environment in a Monitor wrapper\n",
    "env = VecMonitor(env)\n",
    "\n",
    "# MODEL DEFINITION\n",
    "policy_kwargs = dict(activation_fn=torch.nn.ReLU, net_arch=[dict(pi=[120,120,120], vf=[120,120,120])], log_std_init = 0)\n",
    "model = PPO(\n",
    "    \"MlpPolicy\",\n",
    "    env,\n",
    "    policy_kwargs=policy_kwargs,\n",
    "    verbose=0,\n",
    "    tensorboard_log=log_dir,\n",
    "    n_steps=1000,\n",
    "    batch_size=5000,\n",
    "    n_epochs=10,\n",
    "    gamma=0.999\n",
    ")\n",
    "\n",
    "print(model.policy)\n",
    "print(model.num_timesteps)\n",
    "\n",
    "def animate_policy(model, env, deterministic=False, log_times=False, **kwargs):\n",
    "    env.reset()\n",
    "    def run():\n",
    "        actions, _ = model.predict(env.states, deterministic=deterministic)\n",
    "        states, rewards, dones, infos = env.step(actions)\n",
    "        if log_times:\n",
    "            if rewards[0] == 10:\n",
    "                print(env.step_counts[0]*env.dt)\n",
    "        return env.render()\n",
    "    animation.view(run, gate_pos=env.gate_pos, gate_yaw=env.gate_yaw, **kwargs)\n",
    "\n",
    "# animate untrained policy (use this to set the recording camera position)\n",
    "animate_policy(model, test_env)\n",
    "\n",
    "def train(model, test_env, log_name, n=10000000000):\n",
    "    # save every 10 policy rollouts\n",
    "    TIMESTEPS = model.n_steps*env.num_envs*10\n",
    "    for i in range(0,n):\n",
    "        model.learn(total_timesteps=TIMESTEPS, reset_num_timesteps=False, tb_log_name=log_name)\n",
    "        time_steps = model.num_timesteps\n",
    "        # save model\n",
    "        model.save(models_dir + '/' + log_name + '/' + str(time_steps))\n",
    "        # save policy animation\n",
    "        animate_policy(\n",
    "            model,\n",
    "            test_env,\n",
    "            record_steps=1000,\n",
    "            record_file=video_log_dir + '/' + log_name + '/' + str(time_steps) + '.mp4',\n",
    "            show_window=False\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run training loop\n",
    "# train(model, test_env, 'INDI_gate_penalty_no_rate_penalty')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simulate PPO model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_env = Quadcopter3DGates(num_envs=1, gates_pos=gate_pos, gate_yaw=gate_yaw, start_pos=start_pos, gates_ahead=1, pause_if_collision=True)\n",
    "\n",
    "# path = 'models/PPO_INDI/INDI_TEST2_no_rate_penalty/375000000'\n",
    "path = 'models/PPO_INDI/INDI_TEST2_with_rate_penalty/375000000'\n",
    "path = 'models/PPO_INDI/INDI_TEST3_pq_max=3_rate_penalty/65000000'\n",
    "path = 'models/PPO_INDI/INDI_TEST3_pq_max=3_no_rate_penalty/117000000'\n",
    "path = 'models/PPO_INDI/INDI_gate_penalty_no_rate_penalty/305000000'\n",
    "\n",
    "\n",
    "model = PPO.load(path)\n",
    "animate_policy(model, test_env, deterministic=False, log_times=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_env.reset()\n",
    "action_list = []\n",
    "\n",
    "def run():\n",
    "    global action_list\n",
    "    actions, _ = model.predict(test_env.states, deterministic=False)\n",
    "    action_list.append(actions[0])\n",
    "    states, rewards, dones, infos = test_env.step(actions)\n",
    "    return test_env.render()\n",
    "\n",
    "animation.view(run, gate_pos=test_env.gate_pos, gate_yaw=test_env.gate_yaw)\n",
    "action_list = np.array(action_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.plot(action_list[:,0])\n",
    "plt.plot(action_list[:,1])\n",
    "plt.plot(action_list[:,2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LOGGING\n",
    "test_env.max_steps = 10000\n",
    "test_env.reset()\n",
    "\n",
    "time = []\n",
    "state_traj = []\n",
    "action_traj = []\n",
    "\n",
    "def run():\n",
    "    actions, _ = model.predict(test_env.states, deterministic=False)\n",
    "    states, rewards, dones, infos = test_env.step(actions)\n",
    "    state_traj.append(test_env.world_states[0].copy())\n",
    "    action_traj.append(actions[0].copy())\n",
    "    time.append(test_env.step_counts[0]*test_env.dt)\n",
    "        \n",
    "    return test_env.render()\n",
    "\n",
    "animation.view(run, gate_pos=env.gate_pos, gate_yaw=env.gate_yaw)\n",
    "\n",
    "# make log dict\n",
    "log_dict = {}\n",
    "state_traj = np.array(state_traj)\n",
    "action_traj = np.array(action_traj)\n",
    "\n",
    "log_dict['t'] = time\n",
    "log_dict['x'] = state_traj[:,0]\n",
    "log_dict['y'] = state_traj[:,1]\n",
    "log_dict['z'] = state_traj[:,2]\n",
    "log_dict['vx'] = state_traj[:,3]\n",
    "log_dict['vy'] = state_traj[:,4]\n",
    "log_dict['vz'] = state_traj[:,5]\n",
    "log_dict['V'] = np.sqrt(state_traj[:,3]**2 + state_traj[:,4]**2 + state_traj[:,5]**2)\n",
    "log_dict['phi'] = state_traj[:,6]\n",
    "log_dict['theta'] = state_traj[:,7]\n",
    "log_dict['psi'] = state_traj[:,8]\n",
    "log_dict['u1'] = (action_traj[:,0]+1)/2\n",
    "log_dict['u2'] = (action_traj[:,1]+1)/2\n",
    "log_dict['u3'] = (action_traj[:,2]+1)/2\n",
    "log_dict['u4'] = (action_traj[:,3]+1)/2\n",
    "log_dict['u'] = np.stack([log_dict['u1'], log_dict['u2'], log_dict['u3'], log_dict['u4']], axis=1)\n",
    "\n",
    "# save log dict to flight_data/simulation_logs\n",
    "# make folder\n",
    "name = 'INDI_NET_sim_with_gate_penalty'\n",
    "folder = 'flight_data/simulation_logs'\n",
    "if not os.path.exists(folder):\n",
    "    os.makedirs(folder)\n",
    "# save dict as npz\n",
    "np.savez(folder + '/' + name, **log_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate C code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "# first test only moment randomization\n",
    "# path = 'models/PPO_INDI/INDI_TEST2_no_rate_penalty/375000000'\n",
    "# path = 'models/PPO_INDI/INDI_TEST2_with_rate_penalty/375000000'\n",
    "model = PPO.load(path)\n",
    "\n",
    "# get network\n",
    "network = list(model.policy.mlp_extractor.policy_net) + [model.policy.action_net]\n",
    "network = nn.Sequential(*network)\n",
    "print('NETWORK:')\n",
    "print(network)\n",
    "\n",
    "print(model.policy.action_dist)\n",
    "print(model.policy.log_std)\n",
    "print(model.policy.log_std.exp())\n",
    "network_std = model.policy.log_std.exp().cpu().detach().numpy()\n",
    "print(network_std)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_env = Quadcopter3DGates(num_envs=1, gates_pos=gate_pos, gate_yaw=gate_yaw, start_pos=start_pos, gates_ahead=1, pause_if_collision=True)\n",
    "\n",
    "animate_policy(model, test_env, deterministic=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import os\n",
    "import subprocess\n",
    "\n",
    "# remove the c_code folder and all of its contents\n",
    "subprocess.call('rm -rf c_code', shell=True)\n",
    "# create a new c_code folder\n",
    "subprocess.call('mkdir c_code', shell=True)\n",
    "\n",
    "# Create the \"c_code\" folder if it doesn't exist\n",
    "output_folder = \"c_code\"\n",
    "os.makedirs(output_folder, exist_ok=True)\n",
    "\n",
    "# Generate the C file and the header file inside the \"c_code\" folder\n",
    "source_file_path = os.path.join(output_folder, \"neural_network.c\")\n",
    "header_file_path = os.path.join(output_folder, \"neural_network.h\")\n",
    "\n",
    "# np.float32 to str\n",
    "float_to_str = lambda x: str(float(x))\n",
    "\n",
    "# Generate the C file\n",
    "with open(source_file_path, \"w\") as file:\n",
    "    file.write('#include \"neural_network.h\"\\n')\n",
    "    file.write(\"#include <stdio.h>\\n\")\n",
    "    file.write(\"#include <math.h>\\n\\n\")\n",
    "\n",
    "    # Define weights and biases as global constant float arrays\n",
    "    i = 1\n",
    "    for layer in network:\n",
    "        if isinstance(layer, nn.Linear):\n",
    "            weights_layer = layer.weight.data.cpu().numpy()\n",
    "            biases_layer = layer.bias.data.cpu().numpy()\n",
    "\n",
    "            file.write(f\"const float weights_fc{i}[] = {{\\n\")\n",
    "            file.write(\",\\n\".join([\", \".join(map(float_to_str, row)) for row in weights_layer]))\n",
    "            file.write(\"\\n};\\n\\n\")\n",
    "\n",
    "            file.write(f\"const float biases_fc{i}[] = {{\\n\")\n",
    "            file.write(\", \".join(map(float_to_str, biases_layer)))\n",
    "            file.write(\"\\n};\\n\\n\")\n",
    "\n",
    "            i+=1\n",
    "\n",
    "    # LINEAR LAYER\n",
    "    file.write(\"void nn_linear(const float* weights, const float* biases, const float* input, int in_features, int out_features, float* output) {\\n\")\n",
    "    file.write(\"    for (int i = 0; i < out_features; ++i) {\\n\")\n",
    "    file.write(\"        float neuron = biases[i];\\n\")\n",
    "    file.write(\"        for (int j = 0; j < in_features; ++j) {\\n\")\n",
    "    file.write(\"            neuron += input[j] * weights[i * in_features + j];\\n\")\n",
    "    file.write(\"        }\\n\")\n",
    "    file.write(\"        output[i] = neuron;\\n\")\n",
    "    file.write(\"    }\\n\")\n",
    "    file.write(\"}\\n\\n\")\n",
    "\n",
    "    # RELU LAYER\n",
    "    file.write(\"void nn_relu(float* input, int size) {\\n\")\n",
    "    file.write(\"    for (int i = 0; i < size; ++i) {\\n\")\n",
    "    file.write(\"        input[i] = fmaxf(0, input[i]);\\n\")\n",
    "    file.write(\"    }\\n\")\n",
    "    file.write(\"}\\n\\n\")\n",
    "\n",
    "    # TANH LAYER\n",
    "    file.write(\"void nn_tanh(float* input, int size) {\\n\")\n",
    "    file.write(\"    for (int i = 0; i < size; ++i) {\\n\")\n",
    "    file.write(\"        input[i] = tanh(input[i]);\\n\")\n",
    "    file.write(\"    }\\n\")\n",
    "    file.write(\"}\\n\\n\")\n",
    "\n",
    "    # FORWARD FUNCTION\n",
    "    file.write(\"void nn_forward(const float* input, float* output) {\\n\")\n",
    "    layer_size = network[0].out_features\n",
    "    num_layers = sum(isinstance(layer, nn.Linear) for layer in network)\n",
    "    i=0\n",
    "    input_array = \"input\"\n",
    "    for layer in network:\n",
    "        if isinstance(layer, nn.Linear):\n",
    "            i+=1\n",
    "            if i<num_layers:\n",
    "                file.write(f\"    float fc{i}_output[{layer.out_features}];\\n\")\n",
    "                file.write(f\"    nn_linear(weights_fc{i}, biases_fc{i}, {input_array}, {layer.in_features}, {layer.out_features}, fc{i}_output);\\n\")\n",
    "                input_array = f\"fc{i}_output\"\n",
    "            else:\n",
    "                file.write(f\"    nn_linear(weights_fc{i}, biases_fc{i}, {input_array}, {layer.in_features}, {layer.out_features}, output);\\n\")\n",
    "                input_array = \"output\"\n",
    "            layer_size = layer.out_features\n",
    "        elif isinstance(layer, nn.ReLU):\n",
    "            file.write(f\"    nn_relu({input_array}, {layer_size});\\n\")\n",
    "        elif isinstance(layer, nn.Tanh):\n",
    "            file.write(f\"    nn_tanh({input_array}, {layer_size});\\n\")\n",
    "        else:\n",
    "            raise Exception(f\"Unsupported layer: {layer}\")\n",
    "    file.write(\"}\\n\")\n",
    "\n",
    "# Generate the header file\n",
    "with open(header_file_path, \"w\") as header_file:\n",
    "    header_file.write(\"#ifndef NEURAL_NETWORK_H\\n\")\n",
    "    header_file.write(\"#define NEURAL_NETWORK_H\\n\\n\")\n",
    "    # Declare the forward function in the header file\n",
    "    header_file.write(\"void nn_forward(const float* input, float* output);\\n\")\n",
    "    header_file.write(\"\\n#endif // NEURAL_NETWORK_H\\n\")\n",
    "\n",
    "\n",
    "# Print the generated files\n",
    "# Print the generated files\n",
    "print(f\"Generated {source_file_path}\")\n",
    "print(f\"Generated {header_file_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "name = 'nn_controller'\n",
    "# Create the \"c_code\" folder if it doesn't exist\n",
    "output_folder = \"c_code\"\n",
    "os.makedirs(output_folder, exist_ok=True)\n",
    "\n",
    "# Generate the C file and the header file inside the \"c_code\" folder\n",
    "source_file_path = os.path.join(output_folder, f\"{name}.c\")\n",
    "header_file_path = os.path.join(output_folder, f\"{name}.h\")\n",
    "\n",
    "num_gates = test_env.num_gates\n",
    "gates_ahead = test_env.gates_ahead\n",
    "\n",
    "# Generate the header file\n",
    "with open(header_file_path, \"w\") as file:\n",
    "    file.write(f\"#ifndef {name.upper()}_H\\n\")\n",
    "    file.write(f\"#define {name.upper()}_H\\n\")\n",
    "    file.write(\"#define NN_INDI_CMDS\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(\"#include <stdint.h>\\n\")\n",
    "    file.write(\"#include <stdbool.h>\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(f'#define GATES_AHEAD {gates_ahead}\\n')\n",
    "    file.write(f'#define NUM_GATES {num_gates}\\n')\n",
    "    file.write(\"\\n\")\n",
    "    # include neural network code\n",
    "    file.write(\"// Include the neural network code\\n\")\n",
    "    file.write(\"#include \\\"neural_network.h\\\"\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(\"const float gate_pos[NUM_GATES][3];\\n\")\n",
    "    file.write(\"const float gate_yaw[NUM_GATES];\\n\")\n",
    "    file.write(\"const float start_pos[3];\\n\")\n",
    "    # file.write(\"const float gate_pos_rel[NUM_GATES][3];\\n\")\n",
    "    # file.write(\"const float gate_yaw_rel[NUM_GATES];\\n\")\n",
    "    file.write(\"uint8_t target_gate_index;\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    # nn_reset function that resets the target gate index\n",
    "    file.write(\"void nn_reset(void);\\n\")\n",
    "    # nn_control function that that takes as input a float array of size 16 (world_state) and outputs an array of size 4 (rpms)\n",
    "    file.write(\"void nn_control(const float world_state[13], float indi_cmd[4]);\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(\"#endif\\n\")\n",
    "\n",
    "# Generate the C file\n",
    "with open(source_file_path, \"w\") as file:\n",
    "    file.write(f\"#include \\\"{name}.h\\\"\\n\")\n",
    "    file.write(\"#include <math.h>\\n\")\n",
    "    file.write(\"#include <stdlib.h>\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    # define boolean to set controller to determistic\n",
    "    file.write(\"bool deterministic = false;\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(\"const float output_std[4] = {\\n\")\n",
    "    for i in range(4):\n",
    "        file.write(f\"    {network_std[i]},\\n\")\n",
    "    file.write(\"};\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    # define the gate positions and headings as const float arrays\n",
    "    file.write(\"const float gate_pos[NUM_GATES][3] = {\\n\")\n",
    "    for i in range(num_gates):\n",
    "        file.write(f\"    {{{test_env.gate_pos[i][0]}, {test_env.gate_pos[i][1]}, {test_env.gate_pos[i][2]}}},\\n\")\n",
    "    file.write(\"};\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(\"const float gate_yaw[NUM_GATES] = {\\n\")\n",
    "    for i in range(num_gates):\n",
    "        file.write(f\"    {test_env.gate_yaw[i]},\\n\")\n",
    "    file.write(\"};\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    # define the start pos as a const float array\n",
    "    file.write(\"const float start_pos[3] = {\\n\")\n",
    "    file.write(f\"    {test_env.start_pos[0]}, {test_env.start_pos[1]}, {test_env.start_pos[2]}\\n\")\n",
    "    file.write(\"};\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    # define the relative gate positions and headings as const float arrays\n",
    "    file.write(\"const float gate_pos_rel[NUM_GATES][3] = {\\n\")\n",
    "    for i in range(num_gates):\n",
    "        file.write(f\"    {{{test_env.gate_pos_rel[i][0]}, {test_env.gate_pos_rel[i][1]}, {test_env.gate_pos_rel[i][2]}}},\\n\")\n",
    "    file.write(\"};\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(\"const float gate_yaw_rel[NUM_GATES] = {\\n\")\n",
    "    for i in range(num_gates):\n",
    "        file.write(f\"    {test_env.gate_yaw_rel[i]},\\n\")\n",
    "    file.write(\"};\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    # define the target gate index and set it to 0\n",
    "    file.write(\"uint8_t target_gate_index = 0;\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(\"void nn_reset() {\\n\")\n",
    "    file.write(\"    target_gate_index = 0;\\n\")\n",
    "    file.write(\"}\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    \n",
    "    file.write(\"void nn_control(const float world_state[13], float indi_cmd[4]) {\\n\")\n",
    "    file.write(\"    // Get the current position, velocity and heading\\n\")\n",
    "    file.write(\"    float pos[3] = {world_state[0], world_state[1], world_state[2]};\\n\")\n",
    "    file.write(\"    float vel[3] = {world_state[3], world_state[4], world_state[5]};\\n\")\n",
    "    file.write(\"    float yaw = world_state[8];\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(\"    // Get the position and heading of the target gate\\n\")\n",
    "    file.write(\"    float target_pos[3] = {gate_pos[target_gate_index][0], gate_pos[target_gate_index][1], gate_pos[target_gate_index][2]};\\n\")\n",
    "    file.write(\"    float target_yaw = gate_yaw[target_gate_index];\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(\"    // Set the target gate index to the next gate if we passed through the current one\\n\")\n",
    "    file.write(\"    if (cosf(target_yaw) * (pos[0] - target_pos[0]) + sinf(target_yaw) * (pos[1] - target_pos[1]) > 0) {\\n\")\n",
    "    file.write(\"        target_gate_index++;\\n\")\n",
    "    file.write(\"        // loop back to the first gate if we reach the end\\n\")\n",
    "    file.write(\"        target_gate_index = target_gate_index % NUM_GATES;\\n\")\n",
    "    file.write(\"        // reset the target position and heading\\n\")\n",
    "    file.write(\"        target_pos[0] = gate_pos[target_gate_index][0];\\n\")\n",
    "    file.write(\"        target_pos[1] = gate_pos[target_gate_index][1];\\n\")\n",
    "    file.write(\"        target_pos[2] = gate_pos[target_gate_index][2];\\n\")\n",
    "    file.write(\"        target_yaw = gate_yaw[target_gate_index];\\n\")\n",
    "    file.write(\"    }\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(\"    // Get the position of the drone in gate frame\\n\")\n",
    "    file.write(\"    float pos_rel[3] = {\\n\")\n",
    "    file.write(\"        cosf(target_yaw) * (pos[0] - target_pos[0]) + sinf(target_yaw) * (pos[1] - target_pos[1]),\\n\")\n",
    "    file.write(\"        -sinf(target_yaw) * (pos[0] - target_pos[0]) + cosf(target_yaw) * (pos[1] - target_pos[1]),\\n\")\n",
    "    file.write(\"        pos[2] - target_pos[2]\\n\")\n",
    "    file.write(\"    };\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(\"    // Get the velocity of the drone in gate frame\\n\")\n",
    "    file.write(\"    float vel_rel[3] = {\\n\")\n",
    "    file.write(\"        cosf(target_yaw) * vel[0] + sinf(target_yaw) * vel[1],\\n\")\n",
    "    file.write(\"        -sinf(target_yaw) * vel[0] + cosf(target_yaw) * vel[1],\\n\")\n",
    "    file.write(\"        vel[2]\\n\")\n",
    "    file.write(\"    };\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(\"    // Get the heading of the drone in gate frame\\n\")\n",
    "    file.write(\"    float yaw_rel = yaw - target_yaw;\\n\")\n",
    "    file.write(\"    while (yaw_rel > M_PI) {yaw_rel -= 2*M_PI;}\\n\")\n",
    "    file.write(\"    while (yaw_rel < -M_PI) {yaw_rel += 2*M_PI;}\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(\"    // Get the neural network input\\n\")\n",
    "    file.write(\"    float nn_input[13+4*GATES_AHEAD];\\n\")\n",
    "    file.write(\"    // position and velocity\\n\")\n",
    "    file.write(\"    for (int i = 0; i < 3; i++) {\\n\")\n",
    "    file.write(\"        nn_input[i] = pos_rel[i];\\n\")\n",
    "    file.write(\"        nn_input[i+3] = vel_rel[i];\\n\")\n",
    "    file.write(\"    }\\n\")\n",
    "    file.write(\"    // attitude\\n\")\n",
    "    file.write(\"    nn_input[6] = world_state[6];\\n\")\n",
    "    file.write(\"    nn_input[7] = world_state[7];\\n\")\n",
    "    file.write(\"    nn_input[8] = yaw_rel;\\n\")\n",
    "    file.write(\"    // body rates\\n\")\n",
    "    file.write(\"    nn_input[9] = world_state[9];\\n\")\n",
    "    file.write(\"    nn_input[10] = world_state[10];\\n\")\n",
    "    file.write(\"    nn_input[11] = world_state[11];\\n\")\n",
    "    file.write(\"    // normalized thrust \\n\")\n",
    "    file.write(f'    float T_min = {T_min};\\n')\n",
    "    file.write(f'    float T_max = {T_max};\\n')\n",
    "    file.write(\"    nn_input[12] = (world_state[12] - T_min) / (T_max - T_min) * 2 - 1;\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(\"    // relative gate positions and headings\\n\")\n",
    "    file.write(\"    for (int i = 0; i < GATES_AHEAD; i++) {\\n\")\n",
    "    file.write(\"        uint8_t index = target_gate_index + i + 1;\\n\")\n",
    "    file.write(\"        // loop back to the first gate if we reach the end\\n\")\n",
    "    file.write(\"        index = index % NUM_GATES;\\n\")\n",
    "    file.write(\"        nn_input[13+4*i]   = gate_pos_rel[index][0];\\n\")\n",
    "    file.write(\"        nn_input[13+4*i+1] = gate_pos_rel[index][1];\\n\")\n",
    "    file.write(\"        nn_input[13+4*i+2] = gate_pos_rel[index][2];\\n\")\n",
    "    file.write(\"        nn_input[13+4*i+3] = gate_yaw_rel[index];\\n\")\n",
    "    file.write(\"    }\\n\")\n",
    "    file.write(\"    // Get the neural network output and write to the action array\\n\")\n",
    "    file.write(\"    float nn_output[4];\\n\")\n",
    "    file.write(\"    nn_forward(nn_input, nn_output);\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    # if determinstic is false, add gaussian noise to the output\n",
    "    file.write(\"    // add gaussian noise to the output\\n\")\n",
    "    file.write(\"    if (!deterministic) {\\n\")\n",
    "    file.write(\"        for (int i = 0; i < 4; i++) {\\n\")\n",
    "    # generate random gaussian variables using the Box–Muller transform\n",
    "    file.write(\"            // generate random gaussian variables using the Box–Muller transform\\n\")\n",
    "    file.write(\"            float u1 = (float)rand() / RAND_MAX;\\n\")\n",
    "    file.write(\"            float u2 = (float)rand() / RAND_MAX;\\n\")\n",
    "    file.write(\"            float rand_std = sqrtf(-2 * logf(u1)) * cosf(2 * M_PI * u2);\\n\")\n",
    "    file.write(\"            // add the noise to the output\\n\")\n",
    "    file.write(\"            nn_output[i] += output_std[i] * rand_std;\\n\")\n",
    "    file.write(\"        }\\n\")\n",
    "    file.write(\"    }\\n\")\n",
    "    file.write(\"\\n\")\n",
    "    file.write(\"    for (int i = 0; i < 4; i++) {\\n\")\n",
    "    file.write(\"        // clip the output to the range [-1, 1]\\n\")\n",
    "    file.write(\"        if (nn_output[i] > 1) {nn_output[i] = 1;}\\n\")\n",
    "    file.write(\"        if (nn_output[i] < -1) {nn_output[i] = -1;}\\n\")\n",
    "    file.write(\"    }\\n\")\n",
    "    file.write(\"    // map the output to the correct ranges\\n\")\n",
    "    file.write(f\"    float p_min = {p_min};\\n\")\n",
    "    file.write(f\"    float p_max = {p_max};\\n\")\n",
    "    file.write(f\"    float q_min = {q_min};\\n\")\n",
    "    file.write(f\"    float q_max = {q_max};\\n\")\n",
    "    file.write(f\"    float r_min = {r_min};\\n\")\n",
    "    file.write(f\"    float r_max = {r_max};\\n\")\n",
    "    file.write(\"    indi_cmd[0] = (nn_output[0] + 1) / 2 * (p_max - p_min) + p_min;\\n\")\n",
    "    file.write(\"    indi_cmd[1] = (nn_output[1] + 1) / 2 * (q_max - q_min) + q_min;\\n\")\n",
    "    file.write(\"    indi_cmd[2] = (nn_output[2] + 1) / 2 * (r_max - r_min) + r_min;\\n\")\n",
    "    file.write(\"    indi_cmd[3] = (nn_output[3] + 1) / 2 * (T_max - T_min) + T_min;\\n\")\n",
    "    file.write(\"}\\n\")\n",
    "\n",
    "# Print the generated files\n",
    "print(f\"Generated {source_file_path}\")\n",
    "print(f\"Generated {header_file_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test C code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import subprocess\n",
    "import ctypes\n",
    "import numpy as np\n",
    "import importlib\n",
    "importlib.reload(ctypes)\n",
    "\n",
    "# https://cu7ious.medium.com/how-to-use-dynamic-libraries-in-c-46a0f9b98270\n",
    "path = os.path.abspath('c_code')\n",
    "# Create object files\n",
    "subprocess.call('gcc -fPIC -c *.c', shell=True, cwd=path)\n",
    "# Create library\n",
    "subprocess.call('gcc -shared -Wl,-soname,libtools.so -o libtools.so *.o', shell=True, cwd=path)\n",
    "# Remove object files\n",
    "subprocess.call('rm *.o', shell=True, cwd=path)\n",
    "\n",
    "lib_path = os.path.abspath(\"c_code/libtools.so\")\n",
    "fun = ctypes.CDLL(lib_path)\n",
    "\n",
    "print(fun.nn_forward)\n",
    "print(fun.nn_control)\n",
    "\n",
    "# define argument types \n",
    "fun.nn_forward.argtypes = [ctypes.POINTER(ctypes.c_float), ctypes.POINTER(ctypes.c_float)]\n",
    "fun.nn_control.argtypes = [ctypes.POINTER(ctypes.c_float), ctypes.POINTER(ctypes.c_float)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def nn_control(x):\n",
    "    x = np.array(x, dtype=np.float32)\n",
    "\n",
    "    # scale to Thrust to T_min and T_max\n",
    "    x[12] = (x[12] + 1)/2*(T_max - T_min) + T_min\n",
    "\n",
    "    c_net_input = (ctypes.c_float*len(x))(*x)\n",
    "    c_net_output = (ctypes.c_float*4)()\n",
    "    \n",
    "    fun.nn_control(c_net_input, c_net_output)\n",
    "    out = np.array(c_net_output[:])\n",
    "    \n",
    "    # map back to [-1,1]\n",
    "    out[0] = (out[0] - p_min)/(p_max - p_min)*2 - 1\n",
    "    out[1] = (out[1] - q_min)/(q_max - q_min)*2 - 1\n",
    "    out[2] = (out[2] - r_min)/(r_max - r_min)*2 - 1\n",
    "    out[3] = (out[3] - T_min)/(T_max - T_min)*2 - 1\n",
    "    return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate C Network\n",
    "test_env = Quadcopter3DGates(num_envs=1, gates_pos=gate_pos, gate_yaw=gate_yaw, start_pos=start_pos, gates_ahead=1, pause_if_collision=True)\n",
    "test_env.reset()\n",
    "fun.nn_reset()\n",
    "\n",
    "action_list = []\n",
    "\n",
    "def run():\n",
    "    world_state = test_env.world_states[0]\n",
    "    action = nn_control(world_state.copy())\n",
    "    action_list.append(action)\n",
    "    \n",
    "    actions = np.array([action])\n",
    "\n",
    "    states, rewards, dones, infos = test_env.step(actions)\n",
    "\n",
    "    if dones[0]:\n",
    "        fun.nn_reset()\n",
    "    return test_env.render()\n",
    "\n",
    "animation.view(run, gate_pos=test_env.gate_pos, gate_yaw=test_env.gate_yaw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "action_list = np.array(action_list)\n",
    "plt.plot(action_list[:,0], label='p_cmd')\n",
    "plt.plot(action_list[:,1], label='q_cmd')\n",
    "plt.plot(action_list[:,2], label='r_cmd')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "plt.plot(action_list[:,3], label='T_cmd')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.10.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}