text: import torch
  from torch import nn
  import torch.nn.functional as F
  import torch.optim as optim
  from model import Net
  from torch.autograd import Variable
  import numpy as np
  
  test = False
  class Controller(nn.Module):
  
      def __init__(self):
          # Inherit from parent constructor
          super(Controller, self).__init__()
          self.nodeSize = {
              0: "term",
              1: 2,
              2: 4,
              3: 8,
              4: 16,
              5: 32,
              6: 64
          }
          self.activations = {
              0: "ReLU",
              1: "Tanh",
              2: "Sigmoid",
          }
          self.hidden_dim = 50
          self.max_depth = 12
          self.acc_tests = 5
          self.GRU = nn.GRUCell(input_size=self.hidden_dim,hidden_size = self.hidden_dim)
          
          self.decoder = []
          for i in range(self.max_depth):
              if i%2 == 0:
                  self.decoder.append(nn.Linear(self.hidden_dim,len(self.nodeSize)))
              else:
                  self.decoder.append(nn.Linear(self.hidden_dim,len(self.activations)))
  
          
          self.optimizer = optim.Adam(self.parameters(),lr=0.001)
          self.epsilon = 1
          self.baseline = 0.85
      def forward(self,x,h,depth):
          x = x.unsqueeze(0)
          h = h.unsqueeze(0)
          h = self.GRU(x,h)
          x = self.decoder[depth](h)
          
  
          return x.squeeze(),h.squeeze() 
  
  
      def step(self,state,depth):
          logits,new_state = self.forward(torch.zeros(self.hidden_dim),state,depth)
          self.probs = F.softmax(logits,dim=-1)
          log_probs = F.log_softmax(logits,dim=-1)
          choice = self.probs.multinomial(num_samples=1).data
          if depth %2 == 0:
              action = self.nodeSize[int(choice)]
          else:
              action = self.activations[int(choice)]
          act_log_prob = log_probs[int(choice)]
          return action, act_log_prob ,new_state
  
      def generate_rollout(self):
          # -- init state -- #
          self.states = []
          self.actions = []
          self.log_probs = []
          self.state_entropy = []
          self.reward=0
          state = torch.zeros(self.hidden_dim)
          self.states.append(state)
          depth = 0
          while True and self.max_depth > depth:
              action,log_prob,next_state = self.step(state,depth)
              state = next_state
            
              if action == "term":
                  self.log_probs.append(log_prob)
                  self.states.append(state)
                  self.state_entropy.append(torch.mul(log_prob.sum(),self.probs.sum()))            
                  break
  
              self.actions.append(action)
              self.log_probs.append(log_prob)
              self.states.append(state)
  
              self.state_entropy.append(-torch.mul(log_prob.sum(),self.probs.sum()))
              depth+=1
  
          net = Net(self.actions) 
          acc = 0
          for _ in range(self.acc_tests):
              acc += net.train()
          acc /= self.acc_tests
          self.reward += acc
          
          if self.reward > self.baseline:
              self.baseline*=1.01
  
          self.reward -= self.baseline    
          return self.optimize(),acc
         
  
  
      def optimize(self):
          # self.log_probs = torch.cat(self.log_probs)
          # self.state_entropy = torch.cat(self.state_entropy)
          # loss = -torch.mean(torch.mul(self.log_probs,(R + (self.epsilon * self.state_entropy) )))
          
          loss = 0
          R = torch.ones(1)*self.reward 
        
          for i in range(len(self.log_probs)):
               loss -= self.log_probs[i]*Variable(R) - self.epsilon*self.state_entropy[i]
               
  
          loss /= len(self.log_probs)
          self.optimizer.zero_grad()
          loss.backward()    
          self.optimizer.step()
          self.epsilon *= 0.99
          return loss.data
          
         
          
  
  
              
  
  
  
  
  if test:
      controller = Controller()
      controller.generate_rollout()
      
      