<jupyter_start><jupyter_text><jupyter_code>import torch as T
from torch import nn
import torch.distributions as D 
from torch import optim as O
import numpy as np
import gym
device = T.device('cpu')

if(T.cuda.is_available()): 
    device = T.device('cuda:0') 
    T.cuda.empty_cache()
    print("Device set to : " + str(T.cuda.get_device_name(device)))
else:
    print("Device set to : cpu")
env = gym.make('Pendulum-v0')
state_shape = env.observation_space.shape[0]
action_shape = env.observation_space.shape[0]
action_bound = [env.action_space.low, env.action_space.high]
print(f'state shape: {state_shape}')
print(f'action shape: {action_shape}')
print(f'action bound: {action_bound}')
class PolicyNW(nn.Module):

    def __init__(self, state_shape, action_shape):
        super(PolicyNW, self).__init__()
        self.fc1 = nn.Linear(state_shape, 128)
        self.relu1 = nn.ReLU()
        self.fcmu = nn.Linear(128, action_shape)
        self.mu_act = nn.Tanh()
        self.fcsigma = nn.Linear(128, action_shape)
        self.sigma_act = nn.Softplus()

    def forward(self, x):

        x = self.relu1(self.fc1(x))
        mu = 2 * self.mu_act(self.fcmu(x))
        sigma = self.sigma_act(self.fcsigma(x))

        return mu, sigma

pi = PolicyNW(state_shape, action_shape).to(device)
old_pi = PolicyNW(state_shape, action_shape).to(device)
def update_old_pi():
    old_pi.load_state_dict(pi.state_dict())

update_old_pi()
v = nn.Sequential(nn.Linear(state_shape, 128), nn.ReLU(), nn.Linear(128, 1)).to(device)
pi_optim = O.Adam(pi.parameters(), lr=1e-3)
v_optim = O.Adam(v.parameters(), lr=2e-3)
#Hyper Parameters
num_episodes = 2000
num_timesteps = 200
gamma = 0.9
delta = 0.3
beta = 0.2
epsilon = 0.2
batch_size = 32
epochs = 10
def policy(state):
    with T.no_grad():
        state = T.tensor(state, dtype=T.float32).unsqueeze(0).to(device)
        mu, sigma = pi(state)
        dist = D.Normal(mu[0], sigma[0])
        action = dist.sample().unsqueeze(0)
        clipped = T.clip(action, min=-2, max=2)[0]
    
    return clipped.cpu().numpy()
def value(state):
    with T.no_grad():
        state = T.tensor(state, dtype=T.float32).unsqueeze(0).to(device)
        out = v(state)[0, 0]
    return out.cpu().numpy()
def train_step(state, action, reward, _beta):

    update_old_pi()

    state = T.tensor(state, dtype=T.float32).to(device)
    action = T.tensor(action, dtype=T.float32).to(device)
    reward = T.tensor(reward, dtype=T.float32).to(device)

    with T.no_grad():
        advantage_cost = reward - v(state)

    for _ in range(epochs):

        mu, sigma = pi(state)
        dist = D.Normal(mu, sigma)
        pi_prob = dist.log_prob(action)

        with T.no_grad():
            mu_, sigma_ = old_pi(state)
            dist_ = D.Normal(mu_, sigma_)
            old_pi_prob = dist_.log_prob(action)

        kl_div = nn.KLDivLoss()(pi_prob, old_pi_prob)

        ratio = pi_prob/(old_pi_prob + 1e-8)
        objective = ratio * advantage_cost

        clipped = T.minimum(objective, T.clip(ratio, 1 - epsilon, 1+epsilon) * advantage_cost)

        pi_loss = -T.mean(clipped - _beta * kl_div)

        pi_optim.zero_grad()
        pi_loss.backward()
        pi_optim.step()
    

    mean_kl = T.mean(kl_div)
    if mean_kl > 1.5 * delta:
        _beta *= 2.0
    elif mean_kl < delta/1.5:
        _beta *= 0.5
    
    advantage = reward - v(state)
    v_loss = T.mean(T.square(advantage))

    v_optim.zero_grad()
    v_loss.backward()
    v_optim.step()

    return _beta
for i in range(1, num_episodes + 1):

    state = env.reset()
    episode_states, episode_actions, episode_rewards = [], [], []
    Return = 0

    for t in range(num_timesteps):

        action = policy(state)

        next_state, reward, done, _ = env.step(action)

        episode_states.append(state)
        episode_rewards.append(reward)
        episode_actions.append(action)

        state = next_state

        Return += reward

        if (t+1) % batch_size == 0 or t == num_timesteps - 1:

            v_s_ = value(state)

            discounted_r = []
            for reward in episode_rewards[::-1]:
                v_s_ = reward + gamma * v_s_
                discounted_r.append(v_s_)
            
            discounted_r.reverse()

            es, ea, er = np.vstack(episode_states), np.vstack(episode_actions), np.array(discounted_r, np.float32)[:, np.newaxis]

            beta = train_step(es, ea, er, beta)

            episode_states, episode_actions, episode_rewards = [], [], []
    
    if i % 10 == 0:
        print('Episode: {}, Return: {}'.format(i, Return))<jupyter_output>/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py:2742: UserWarning: reduction: 'mean' divides the total loss by both the batch size and the support size.'batchmean' divides only by the batch size, and aligns with the KL div math definition.'mean' will be changed to behave the same as 'batchmean' in the next major release.
  "reduction: 'mean' divides the total loss by both the batch size and the support size."
{'fraction_non_alphanumeric': 0.1325516693163752, 'fraction_numerical': 0.019475357710651828, 'mean_word_length': 2.80999242997729, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, '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': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '171612', 'license': 'CC0-1.0', 'n_tokens_mistral': 1864, 'n_tokens_neox': 1763, 'n_words': 387, 'path': 'PPO_Pytorch.ipynb', 'repo_name': 'dude123studios/AdvancedReinforcementLearning', 'size': '5032'}