{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np\nimport matplotlib.pyplot as plt\nfrom random import randint, random\nimport sys\nimport gc\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom itertools import count","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"max_delta = 1\ngrid_size = 5\n\nif torch.cuda.is_available():  \n    dev = \"cuda:0\" \nelse:  \n    dev = \"cpu\"\ndevice = torch.device(dev)\n\ndtype = torch.float32","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def grid_images(images, save=None):\n    if not isinstance(images, np.ndarray):\n        images = np.stack(images)\n    assert len(images.shape) >= 2, \"pas assez de dimensions\"\n    assert len(images.shape) <= 4, \"trop de dimensions\"\n    if len(images.shape) == 2:\n        images = np.expand_dims(images, 0)\n    if len(images.shape) == 3:\n        images = np.expand_dims(images, 0)\n    plt.figure(figsize=(images.shape[1], images.shape[0]))\n    print(images.shape)\n\n    for j in range(images.shape[1]):\n        for i in range(images.shape[0]):\n            plt.subplot(images.shape[0], images.shape[1], i * images.shape[1] + j + 1)\n            plt.imshow(images[i, j])\n\n    plt.show()\n    if save is not None:\n        plt.savefig(save)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\n    Game of life functions \n\"\"\"\n\ndef get_padded_version_n(X):\n    X_pad = np.zeros((X.shape[0], X.shape[-2] + 2, X.shape[-1] + 2), dtype=X.dtype)\n    X_pad[:, 1:-1,1:-1] += X\n    \n    X_pad[:, 0, 1:-1] = X[:, -1, :]\n    X_pad[:, -1, 1:-1] = X[:, 0, :]\n    \n    X_pad[:, 1:-1, 0] = X[:, :, -1]\n    X_pad[:, 1:-1, -1] = X[:, :, 0]\n    \n    X_pad[:, 0, 0] = X[:, -1, -1]\n    X_pad[:, 0, -1] = X[:, -1, 0]\n    X_pad[:, -1, 0] = X[:, 0, -1]\n    X_pad[:, -1, -1] = X[:, 0, 0]\n    \n    return X_pad\n\ndef nConv2d_sw_3x3(X):\n    X_pad = get_padded_version_n(X)\n    N = np.zeros_like(X_pad)\n    \n    N[:, 1:, 1:] += X_pad[:,:-1,:-1]\n    N[:, 1:, :] += X_pad[:,:-1,:]\n    N[:, 1:, :-1] += X_pad[:,:-1,1:]\n\n    N[:, :, 1:] += X_pad[:,:,:-1]\n    N[:, :, :] += X_pad[:,:,:]\n    N[:, :, :-1] += X_pad[:,:,1:]\n\n    N[:, :-1, 1:] += X_pad[:,1:,:-1]\n    N[:, :-1, :] += X_pad[:,1:,:]\n    N[:, :-1, :-1] += X_pad[:,1:,1:]\n    \n    N = N[:,1:-1,1:-1]\n    \n    return N\n\ndef life_step(X):\n    if len(X.shape) == 2:\n        X = np.expand_dims(X, 0)\n    N =  nConv2d_sw_3x3(X) - X\n    return np.squeeze(np.logical_or(N == 3, np.logical_and(X, N==2)).astype(np.uint8))\n\nclass GRGoL_env:\n    def __init__(self, size, max_delta, max_step, h=1):\n        self.ydim, self.xdim = size, size\n        self.grid_size = self.ydim * self.xdim\n        self.max_delta = max_delta\n        self.h = h\n        self.max_step = max_step\n        \n    def reset(self):\n        valid = False\n        while not valid:\n            grid = (np.random.random((self.ydim, self.xdim)) < random() * 0.98 + 0.01).astype(np.uint8)\n\n            # Warmup steps\n            for i in range(5):\n                grid = life_step(grid)\n\n            # Generate delta\n            delta = randint(1, self.max_delta)\n\n            # Calculating final state\n            initial_grid = grid\n            for i in range(delta):\n                grid = life_step(grid)\n            target = grid\n            \n            valid = np.sum(target) > 0\n        self.initial_grid = initial_grid\n        self.target = target\n        self.delta = delta\n        self.solution_h = np.repeat(np.expand_dims(np.random.randint(0, 2, (self.ydim, self.xdim)), 0), self.h, axis=0)\n        res = life_step(self.solution_h[-1])\n        self.best_score = np.sum(res == self.target)\n        self.s = 0\n        \n        return self.get_state()\n        \n    def get_state(self):\n        return self.solution_h, self.best_score, self.delta, self.target\n    \n    def step(self, action):\n        initial_score = self.best_score\n        \n        if self.s < self.max_step and self.best_score < self.grid_size:\n            new = np.expand_dims(np.copy(self.solution_h[-1]), 0)\n            new = np.expand_dims(np.copy(self.solution_h[-1]), 0)\n            new.reshape(-1)[action] = 1 - new.reshape(-1)[action]\n            \n            self.solution_h = np.append(self.solution_h, new, axis=0)\n            self.solution_h = np.delete(self.solution_h, 0, 0)\n\n            self.best_score = max(self.best_score, np.sum(life_step(new) == self.target))\n\n            self.s += 1\n            \n        reward = self.best_score - initial_score\n        done = self.s >= self.max_step or self.best_score >= self.grid_size\n        return self.get_state(), reward, done","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Playing\neps_start = 0.99\neps_decay = 0.999\nnum_episodes = 10000\ngamma = 0.999\n#Replay buffer\nbuffer_size = 10000\nmin_buffer_size = 128\n# Training\nlr = 0.001\nOptimizer = torch.optim.Adam\nbatch_size = 32\ntarget_update_freq = 128\n# Model\nn_layers = 5\nn_filters = 32\n\ndef format_res_input(I, s, d, F):\n    if len(I.shape) == 3:\n        I = np.expand_dims(I, 0)\n        d = np.expand_dims(d, 0)\n        F = np.expand_dims(F, 0)\n    s = np.repeat(s / grid_size / grid_size, grid_size * grid_size).reshape((-1, 1, grid_size, grid_size))\n    d_indic = np.repeat(np.eye(max_delta)[d-1], grid_size * grid_size).reshape((-1, max_delta, grid_size, grid_size))\n    F = np.expand_dims(F, axis=1)\n    z = np.concatenate((I, s, d_indic, F), axis=1)\n    return torch.tensor(z, dtype=dtype, device=device) * 2 - 1\n        \ndef count_parameters(model):\n    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n\nclass InBlock(nn.Module):\n    def __init__(self, infilters, filters):\n        super(InBlock, self).__init__()\n        self.conv = nn.Conv2d(infilters, filters, 3, stride=1, padding=1, padding_mode = 'circular')\n        self.bn = nn.BatchNorm2d(filters)\n\n    def forward(self, s):\n        return F.relu(self.bn(self.conv(s)))\n\nclass ResBlock(nn.Module):\n    def __init__(self, infilters, filters):\n        super(ResBlock, self).__init__()\n        self.conv1 = nn.Conv2d(infilters, filters, kernel_size=3, stride=1,\n                     padding=1, bias=False, padding_mode = 'circular')\n        self.bn1 = nn.BatchNorm2d(filters)\n        self.conv2 = nn.Conv2d(filters, filters, kernel_size=3, stride=1,\n                     padding=1, bias=False, padding_mode = 'circular')\n        self.bn2 = nn.BatchNorm2d(filters)\n\n    def forward(self, x):\n        residual = x\n        out = self.conv1(x)\n        out = F.relu(self.bn1(out))\n        out = self.conv2(out)\n        out = self.bn2(out)\n        out += residual\n        out = F.relu(out)\n        return out\n    \nclass OutBlock(nn.Module):\n    def __init__(self, infilters):\n        super(OutBlock, self).__init__()\n        self.conv = nn.Conv2d(infilters, 1, kernel_size=1, padding_mode = 'circular')\n    \n    def forward(self,s):\n        return torch.relu(self.conv(s))\n    \nclass ResNet(nn.Module):\n    def __init__(self, n_filters, res_layers):\n        super(ResNet, self).__init__()\n        self.convi = InBlock(max_delta + 3, n_filters)\n        \n        self.i_blocks = []\n        for _ in range(res_layers):\n            self.i_blocks.append(ResBlock(n_filters, n_filters))\n        self.i_blocks = nn.ModuleList(self.i_blocks)\n        \n        self.convo = OutBlock(n_filters)\n        \n    def forward(self, y):\n        s = self.convi(y)\n        \n        for b in self.i_blocks:\n            s = b(s)\n            \n        return self.convo(s)\n    \nclass ReplayBuffer:\n    def __init__(self, max_size):\n        self.max_size = max_size\n        self.states = np.zeros((max_size, max_delta+3, grid_size, grid_size))\n        self.actions = np.zeros(max_size)\n        self.next_states = np.zeros((max_size, max_delta+3, grid_size, grid_size))\n        self.rewards = np.zeros(max_size)\n        self.cursor = 0\n        \n    def add(self, state, action, next_state, reward):\n        if self.cursor + 1 == self.max_size:\n            self.states = np.concatenate([self.states[1:], np.zeros((1, max_delta + 3, grid_size, grid_size))], axis=0)\n            self.actions = np.concatenate([self.actions[1:], np.zeros(1)], axis=0)\n            self.next_states = np.concatenate([self.next_states[1:], np.zeros((1, max_delta + 3, grid_size, grid_size))], axis=0)\n            self.rewards = np.concatenate([self.rewards[1:], np.zeros(1)], axis=0)\n        self.states[self.cursor] = state\n        self.actions[self.cursor] = action\n        self.next_states[self.cursor] = next_state\n        self.rewards[self.cursor] = reward\n        self.cursor = min(self.max_size - 1, self.cursor + 1)\n        \n    def sample(self, batch_size):\n        idxs = np.random.randint(0, self.cursor, batch_size)\n        return (\n            torch.tensor(self.states[idxs], dtype=dtype, device=device),\n            torch.tensor(self.actions[idxs], dtype=torch.int64, device=device),\n            torch.tensor(self.next_states[idxs], dtype=dtype, device=device),\n            torch.tensor(self.rewards[idxs], dtype=dtype, device=device)\n        )\n        \n    def __len__(self):\n        return self.cursor\n    \ndef optimize_model(buffer):\n    if len(buffer) < min_buffer_size:\n        return\n    state_batch, action_batch, next_state_batch, reward_batch = buffer.sample(batch_size)\n    # Compute Q(s_t, a) - the model computes Q(s_t), then we select the\n    # columns of actions taken. These are the actions which would've been taken\n    # for each batch state according to policy_net\n    qs = policy_net(state_batch).view((-1, grid_size * grid_size))\n    state_action_values = qs.gather(1, torch.unsqueeze(action_batch, axis=-1))\n\n    # Compute V(s_{t+1}) for all next states.\n    # Expected values of actions for non_final_next_states are computed based\n    # on the \"older\" target_net; selecting their best reward with max(1)[0].\n    # This is merged based on the mask, such that we'll have either the expected\n    # state value or 0 in case the state was final.\n    next_state_values = target_net(next_state_batch).view((-1, grid_size * grid_size)).max(1)[0].detach()\n    # Compute the expected Q values\n    expected_state_action_values = (next_state_values * gamma) + reward_batch\n\n    # Compute Huber loss\n    loss = F.smooth_l1_loss(state_action_values, expected_state_action_values.unsqueeze(1))\n\n    # Optimize the model\n    optimizer.zero_grad()\n    loss.backward()\n    for param in policy_net.parameters():\n        param.grad.data.clamp_(-1, 1)\n    optimizer.step()\n    \n    return loss.item()\n    \n\npolicy_net = ResNet(n_filters, n_layers).to(device)\ntarget_net = ResNet(n_filters, n_layers).to(device)\ntarget_net.load_state_dict(policy_net.state_dict())\ntarget_net.eval()\n\noptimizer = Optimizer(policy_net.parameters(), lr=lr)\nenv = GRGoL_env(grid_size, max_delta, grid_size * grid_size)\nbuffer = ReplayBuffer(buffer_size)\ntot_r_h = []\nfor i_episode in range(num_episodes):\n    # Initialize the environment and state\n    state = env.reset()\n    for t in count():\n        # Select and perform an action\n        eps = eps_start * eps_decay ** i_episode\n        nn_state = format_res_input(*state)\n        if random() < eps:\n            action = randint(0, grid_size ** 2-1)\n        else:\n            with torch.no_grad():\n                qs = policy_net(nn_state)\n                action = qs.view(-1, grid_size * grid_size).max(1)[1].view(1, 1).item()\n        new_state, reward, done = env.step(action)\n        \n        # Store the transition in memory\n        buffer.add(nn_state.cpu().numpy(), action, format_res_input(*new_state).cpu().numpy(), reward)\n        \n        # Move to the next state\n        state = new_state\n\n        # Perform one step of the optimization (on the policy network)\n        optimize_model(buffer)\n        \n        if done:\n            tot_r_h.append(env.best_score)\n            print(\"Episode {}: epsilon={}, Solution MAE={}\".format(i_episode, eps, 1 - env.best_score / grid_size / grid_size))\n            break\n    # Update the target network, copying all weights and biases in DQN\n    if i_episode % target_update_freq == target_update_freq-1:\n        print(\"Updating target net\")\n        target_net.load_state_dict(policy_net.state_dict())\n        torch.save(policy_net, 'policy.pt')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.plot(tot_r_h)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}