<jupyter_start><jupyter_text>Gym Environment<jupyter_code>import caviar_tools
from beamselect_env import BeamSelectionEnv<jupyter_output><empty_output><jupyter_text>Hyper Params<jupyter_code>reward_type = 'test'    # 'test' or 'train'
epi = [0,10] #[start,end] 
epi_val = [490,499]


gym_env = BeamSelectionEnv(epi, reward_type)
gym_env_val = BeamSelectionEnv(epi_val)

n_steps_per_epi = []

for i in range(epi[1]+1):
    n_steps_per_epi.append(caviar_tools.linecount([i, i]))

for i in range(1, len(n_steps_per_epi)):
    n_steps_per_epi[i] += n_steps_per_epi[i-1]

n_steps = n_steps_per_epi[-1]
n_steps_val = caviar_tools.linecount(epi_val)
print(n_steps, n_steps_per_epi)
n_steps<jupyter_output><empty_output><jupyter_text>Observation Space : X,Y,Z,pkts_dropped,pkts_transmitted,pkts_buffered,bit_rateAction Space : [3,64] -> [UE,Possible beams]  Replay Memory<jupyter_code># Transition = namedtuple('Transition',
#                         ('state', 'action', 'next_state', 'reward'))


# class ReplayMemory(object):

#     def __init__(self, capacity):
#         self.memory = deque([],maxlen=capacity)

#     def push(self, *args):
#         """Save a transition"""
#         self.memory.append(Transition(*args))

#     def sample(self, batch_size):
#         return random.sample(self.memory, batch_size)

#     def __len__(self):
#         return len(self.memory)<jupyter_output><empty_output><jupyter_text>A2C<jupyter_code>class A2C(nn.Module):

    def __init__(self, inputs:int=7, outputs:int=64*3):
        super().__init__()

        self.affine =  nn.Sequential(
            self.create_linear(inputs,16),
            self.create_linear(16, 32),
            self.create_linear(32,64),
            self.create_linear(64,256)
        )
        
        self.actor_linear = nn.Sequential(
            self.create_linear(256,outputs)
        )

        self.critic_linear = nn.Sequential(
            self.create_linear(256,1)
        )


    
    def create_linear(self,inp:int,out:int)-> nn.Module:
        return nn.Sequential(
            nn.Linear(inp,out),
            nn.ELU()
            # nn.BatchNorm1d(out)
        )
    
    def forward(self, x):
        x = x.to(device)
        x = self.affine(x)

        # Actor - A2C implementation
        x_action = self.actor_linear(x)
        x_action = nn.Softmax(dim=-1)(x_action)

        #  Critic
        x_critic = self.critic_linear(x) 

        return x_action,x_critic<jupyter_output><empty_output><jupyter_text>Storage of actions and policy<jupyter_code>SavedAction = namedtuple('SavedAction', ['log_prob', 'value'])
class Policy():
    def __init__(self):
         # Episode policy and reward history 
        self.rewards = []
        self.saved_actions = []
policy = Policy()
tmp = A2C()
sample = torch.rand((1,7))
out = tmp(sample)
print(out[0].shape,out[1].shape)<jupyter_output><empty_output><jupyter_text>Hyperparams<jupyter_code>BATCH_SIZE = 128
GAMMA = 0.99
# EPS_START = 0.9
# EPS_END = 0.3

# # It depends on overall number of steps, basic intitution is that
# # once steps_done == EPS_DECAY then the probablity of choosing 
# # random action is 33%; considering EPS_END is zero
# # As for ep = [0,10]; approx ep is 80k therefore exploration can be reduced to 33% around 50k
# # Also because of this factor smoothed accuracy matters more for training then seeing the average
# EPS_DECAY = n_steps*0.3

TARGET_UPDATE = 1000
VAL_STEP = 60000<jupyter_output><empty_output><jupyter_text>Action<jupyter_code>policy_net = A2C(11, 192).to(device)

optimizer = optim.Adam(policy_net.parameters(),lr=3e-2)

n_actions = 64*3
steps_done = 0

def select_action(state:torch.Tensor,val:bool=False):
    if not val:
        action_array, state_value = policy_net(state)
        action_dist = Categorical(action_array)
        action_value = action_dist.sample()

        # Add log probability of our chosen action to our history    
        # if policy.policy_history.dim() != 0:
        policy.saved_actions.append(SavedAction(action_dist.log_prob(action_value).float(), state_value.float()))
        # else:
        #     policy.policy_history = (action_dist.log_prob(action_value))
        return action_value
    else:
        with torch.no_grad():
            action = torch.max(policy_net(state)[0],dim=1)
            return action.indices

# Tensorboard
log_dir = './mini_logs/AC' 

writer = SummaryWriter(log_dir=log_dir)<jupyter_output><empty_output><jupyter_text>Optimize Model<jupyter_code>def update_policy(n_step:int):

    """
    Training code. Calculates actor and critic loss and performs backprop.
    """
    R = 0
    saved_actions = policy.saved_actions
    policy_losses = [] # list to save actor (policy) loss
    value_losses = [] # list to save critic (value) loss
    returns = [] # list to save the true values

    # calculate the true value using rewards returned from the environment
    for r in policy.rewards[::-1]:
        # calculate the discounted value
        R = r + GAMMA * R
        returns.insert(0, R)

    returns = torch.tensor(returns).float()
    
    returns = (returns - returns.mean()) / (returns.std() + np.finfo(np.float32).eps)
    for (log_prob, value), R in zip(saved_actions, returns):
        advantage = R - value.item()
        advantage = advantage.float()

        # calculate actor (policy) loss 
        policy_losses.append(-log_prob * advantage)

        # calculate critic (value) loss using L1 smooth loss
        value_losses.append(F.smooth_l1_loss(value, torch.tensor([R])))

    # reset gradients
    optimizer.zero_grad()

    # sum up all the values of policy_losses and value_losses
    loss = torch.stack(policy_losses).sum() + torch.stack(value_losses).sum()
    loss = loss.float()

    writer.add_scalar('Loss',loss.item(),n_step)
    
    # perform backprop
    loss.backward()
    optimizer.step()

    # reset rewards and action buffer
    del policy.rewards[:]
    del policy.saved_actions[:]<jupyter_output><empty_output><jupyter_text>Training loop<jupyter_code>def val(train_step:int,gym_env_val:gym.Env,n_steps_val:int=n_steps_val,
        writer:SummaryWriter=writer):

    state = torch.zeros((1,11), dtype=torch.float32)
    running_reward = 0.0
    for episode in tqdm.tqdm_notebook(range(n_steps_val),desc='Val'):
        # Reset environment and record the starting state
        done = False       

        action = select_action(state, val = 'True')
        # Step through environment using chosen action
        state, reward, done, _ = gym_env_val.step([action.item()//64,action.item()%64])

        state = state.astype(np.float32).reshape(1, state.shape[0])
        state = torch.tensor(state)
        running_reward +=reward
        
        # if done:
        #     break
    
    writer.add_scalar('Val Reward',running_reward/n_steps_val,train_step)
    gym_env_val.close()
    print(f'Validation Reward {running_reward/n_steps_val}')

state = torch.zeros((1,11), dtype=torch.float32)
ovr_reward = 0.0
episode_cnt = 0
for episode in tqdm.tqdm_notebook(range(n_steps),desc='Train'):
     # Reset environment and record the starting state
    done = False       

    action = select_action(state)

    # Step through environment using chosen action
    state, reward, done, _ = gym_env.step([action.data[0].item()//64,action.data[0].item()%64])

    state = state.astype(np.float32).reshape(1, state.shape[0])
    state = torch.tensor(state)

    # Save reward
    policy.rewards.append(reward)
    writer.add_scalar('Rewards',reward.item(),episode)
    ovr_reward+=reward.item()
    

    if episode == n_steps_per_epi[episode_cnt]:
        update_policy(episode)
        episode_cnt += 1

    if (episode)%VAL_STEP == 0 or (episode == n_steps-1):
        val(episode, copy.deepcopy(gym_env_val))


writer.add_hparams(
    {
        'BATCH_SIZE' : BATCH_SIZE,
        'GAMMA' : GAMMA,
        'TARGET_UPDATE' :TARGET_UPDATE
    },
    {
        'Overall Reward':ovr_reward,
        'Average Reward': ovr_reward/n_steps
    }
)

print(f'Overall Train reward = {ovr_reward:.2f}. ' \
    f'Average Reward = {ovr_reward/n_steps:.4f}')
gym_env.close()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11993616498895163, 'fraction_numerical': 0.017800147311563957, 'mean_word_length': 3.2969409282700424, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '236362', 'license': 'MIT', 'n_tokens_mistral': 2909, 'n_tokens_neox': 2760, 'n_words': 532, 'path': 'A2C.ipynb', 'repo_name': 'mehtamohit013/ITU-Challenge', 'size': '8146'}