project/video/gif.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from evogym import sample_robot\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "import copy\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import animation\n",
    "import gym\n",
    "import evogym.envs\n",
    "from evogym import sample_robot\n",
    "from evogym.utils import get_full_connectivity\n",
    "from tqdm import tqdm\n",
    "from multiprocessing import Pool\n",
    "import json\n",
    "import time\n",
    "import imageio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Network(nn.Module):\n",
    "    def __init__(self, n_in, h_size, n_out):\n",
    "        super().__init__()\n",
    "        self.fc1 = nn.Linear(n_in, h_size)\n",
    "        self.fc2 = nn.Linear(h_size, h_size)\n",
    "        self.fc3 = nn.Linear(h_size, n_out)\n",
    " \n",
    "        self.n_out = n_out\n",
    "\n",
    "    def reset(self):\n",
    "        pass\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "\n",
    "        x = self.fc2(x)\n",
    "        x = F.relu(x)\n",
    "\n",
    "        x = self.fc3(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Agent:\n",
    "    def __init__(self, Net, config, genes = None):\n",
    "        self.config = config\n",
    "        self.Net = Net\n",
    "        self.model = None\n",
    "        self.fitness = None\n",
    "\n",
    "        self.device = torch.device(\n",
    "            \"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "        self.make_network()\n",
    "        if genes is not None:\n",
    "            self.genes = genes\n",
    "\n",
    "    def __repr__(self):  # pragma: no cover\n",
    "        return f\"Agent {self.model} > fitness={self.fitness}\"\n",
    "\n",
    "    def __str__(self):  # pragma: no cover\n",
    "        return self.__repr__()\n",
    "\n",
    "    def make_network(self):\n",
    "        n_in = self.config[\"n_in\"]\n",
    "        h_size = self.config[\"h_size\"]\n",
    "        n_out = self.config[\"n_out\"]\n",
    "        self.model = self.Net(n_in, h_size, n_out).to(self.device).double()\n",
    "        return self\n",
    "\n",
    "    @property\n",
    "    def genes(self):\n",
    "        if self.model is None:\n",
    "            return None\n",
    "        with torch.no_grad():\n",
    "            params = self.model.parameters()\n",
    "            vec = torch.nn.utils.parameters_to_vector(params)\n",
    "        return vec.cpu().double().numpy()\n",
    "\n",
    "    @genes.setter\n",
    "    def genes(self, params):\n",
    "        if self.model is None:\n",
    "            self.make_network()\n",
    "        assert len(params) == len(\n",
    "            self.genes), \"Genome size does not fit the network size\"\n",
    "        if np.isnan(params).any():\n",
    "            raise\n",
    "        a = torch.tensor(params, device=self.device)\n",
    "        torch.nn.utils.vector_to_parameters(a, self.model.parameters())\n",
    "        self.model = self.model.to(self.device).double()\n",
    "        self.fitness = None\n",
    "        return self\n",
    "\n",
    "    def mutate_ga(self):\n",
    "        genes = self.genes\n",
    "        n = len(genes)\n",
    "        f = np.random.choice([False, True], size=n, p=[1/n, 1-1/n])\n",
    "        \n",
    "        new_genes = np.empty(n)\n",
    "        new_genes[f] = genes[f]\n",
    "        noise = np.random.randn(n-sum(f))\n",
    "        new_genes[~f] = noise\n",
    "        return new_genes\n",
    "\n",
    "    def act(self, obs):\n",
    "        # continuous actions\n",
    "        with torch.no_grad():\n",
    "            x = torch.tensor(obs).double().unsqueeze(0).to(self.device)\n",
    "            actions = self.model(x).cpu().detach().numpy()\n",
    "        return actions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_cfg(env_name, robot=None):\n",
    "    env = make_env(env_name, robot=robot)\n",
    "    cfg = {\n",
    "        \"n_in\": env.observation_space.shape[0],\n",
    "        \"h_size\": 32,\n",
    "        \"n_out\": env.action_space.shape[0],\n",
    "    }\n",
    "    env.close()\n",
    "    return cfg\n",
    "\n",
    "def make_env(env_name, seed=None, robot=None, **kwargs):\n",
    "    if robot is None: \n",
    "        env = gym.make(env_name)\n",
    "    else:\n",
    "        connections = get_full_connectivity(robot)\n",
    "        env = gym.make(env_name, body=robot)\n",
    "    env.robot = robot\n",
    "    if seed is not None:\n",
    "        env.seed(seed)\n",
    "        \n",
    "    return env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(agent, env, max_steps=500, render=False):\n",
    "    obs = env.reset()\n",
    "    agent.model.reset()\n",
    "    reward = 0\n",
    "    steps = 0\n",
    "    done = False\n",
    "    if render:\n",
    "        imgs = []\n",
    "    while not done and steps < max_steps:\n",
    "        if render:\n",
    "            img = env.render(mode='img')\n",
    "            imgs.append(img)\n",
    "        action = agent.act(obs)\n",
    "        obs, r, done, _ = env.step(action)\n",
    "        reward += r\n",
    "        steps += 1\n",
    "        \n",
    "    if render:\n",
    "        return reward, imgs\n",
    "    return reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "walker = np.array([\n",
    "    [3, 3, 3, 3, 3],\n",
    "    [3, 3, 3, 0, 3],\n",
    "    [3, 3, 0, 3, 3],\n",
    "    [3, 3, 0, 3, 3],\n",
    "    [3, 3, 0, 3, 3]\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\n",
    "    \"env_name\": \"Walker-v0\",\n",
    "    \"robot\": walker,\n",
    "    \"generations\": 500, # To change: increase!\n",
    "    \"lambda\": 5,\n",
    "    \"max_steps\": 500, # to change to 500\n",
    "}\n",
    "\n",
    "cfg = get_cfg(config[\"env_name\"], robot=config[\"robot\"]) # Get network dims\n",
    "cfg = {**config, **cfg} # Merge configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_solution(cfg,name=\"solution.json\"):\n",
    "    with open(name, \"r\") as f:\n",
    "        cfg = json.load(f)\n",
    "    cfg[\"robot\"] = np.array(cfg[\"robot\"])\n",
    "    cfg[\"genes\"] = np.array(cfg[\"genes\"])\n",
    "    a = Agent(Network, cfg, genes=cfg[\"genes\"])\n",
    "    a.fitness = cfg[\"fitness\"]\n",
    "    return a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.7812507025710778\n"
     ]
    }
   ],
   "source": [
    "a = load_solution(cfg=cfg,name=\"solution.json\")\n",
    "cfg = a.config\n",
    "env = make_env(config[\"env_name\"], robot=config[\"robot\"])\n",
    "env.metadata.update({'render_modes': [\"rgb_array\"]})\n",
    "\n",
    "a.fitness, imgs = evaluate(a, env, render=True)\n",
    "env.close()\n",
    "print(a.fitness)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "imageio.mimsave(f'video/name.gif', imgs, duration=(1/50.0))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "evolutionaryClass",
   "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.7.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}