73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
import torch
|
|
from src.core.sampling import rollout
|
|
from src.core.value_estimation import gae
|
|
|
|
def ppo(env_fn, value, policy, epochs, rollout_episodes, rollout_steps, gamma, gae_lambda, clip_ratio, pi_opt, pi_iters, v_opt, v_iters, target_kl=None, max_grad_norm=None):
|
|
|
|
for epoch in range(epochs):
|
|
policy.eval()
|
|
states, actions, rewards, dones = rollout(env_fn, policy, rollout_episodes, rollout_steps)
|
|
|
|
print('mean', states[~dones].mean(0))
|
|
print('std', states[~dones].std(0))
|
|
|
|
print(f'Iteration {epoch} mean episode length {(~dones).sum() / states.shape[0]}')
|
|
print(f'Iteration {epoch} mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
|
|
|
|
policy.train()
|
|
value.train()
|
|
value, policy = ppo_step(value, policy, states, actions, rewards, dones, clip_ratio, gamma, gae_lambda, pi_opt, pi_iters, v_opt, v_iters, target_kl, max_grad_norm)
|
|
|
|
return value, policy
|
|
|
|
def ppo_step(value, policy, states, actions, rewards, dones, clip_ratio, gamma, gae_lambda, pi_opt, pi_iters, v_opt, v_iters, target_kl, max_grad_norm):
|
|
|
|
states = states.detach()
|
|
actions = actions.detach()
|
|
rewards = rewards.detach()
|
|
dones = dones.detach()
|
|
|
|
advantages, returns, valid = gae(states, rewards, value(states), dones, gamma, gae_lambda)
|
|
advantages = advantages.detach()
|
|
returns = returns.detach()
|
|
|
|
# update value function
|
|
|
|
for _ in range(v_iters):
|
|
v_opt.zero_grad()
|
|
value_loss = (value(states) - returns).pow(2)[valid].mean()
|
|
value_loss.backward()
|
|
v_opt.step()
|
|
|
|
# update policy
|
|
|
|
old_dist = policy(states).detach()
|
|
old_logprob = policy.log_prob(old_dist, actions).detach()
|
|
|
|
def g(advantages, clip_ratio):
|
|
return torch.where(advantages >= 0, (1 + clip_ratio) * advantages, (1 - clip_ratio) * advantages)
|
|
|
|
def L(states, actions, advantages, clip_ratio):
|
|
return torch.minimum(
|
|
(policy.log_prob(policy(states), actions) - old_logprob).exp() * advantages,
|
|
g(advantages, clip_ratio)
|
|
)[valid].mean()
|
|
|
|
for _ in range(pi_iters):
|
|
pi_opt.zero_grad()
|
|
ppo_loss = -L(states, actions, advantages, clip_ratio)
|
|
ppo_loss.backward()
|
|
|
|
if max_grad_norm:
|
|
torch.nn.utils.clip_grad_norm(policy.parameters(), max_grad_norm)
|
|
|
|
pi_opt.step()
|
|
|
|
kl = policy.kl_divergence(policy(states), old_dist)[valid].mean()
|
|
if target_kl and kl > target_kl:
|
|
break
|
|
|
|
print('KL', kl.item())
|
|
|
|
return value, policy
|