161 lines
5 KiB
Python
161 lines
5 KiB
Python
from functools import partial
|
|
|
|
import flax
|
|
import jax
|
|
import jax.numpy as jnp
|
|
|
|
|
|
# Chose to use a class as it seemed the easiest way to integrate the CleanRL code style
|
|
# with our need to seperate concerns
|
|
class PPO:
|
|
def __init__(self, args, sensor_apply, actor_apply, critic_apply, feature_extractor_apply, message_passer=None):
|
|
self.args = args
|
|
|
|
if not message_passer:
|
|
message_passer = identity
|
|
|
|
self.ppo_loss_grad_fn = jax.value_and_grad(
|
|
partial(
|
|
ppo_loss,
|
|
args=args,
|
|
sensor_apply=sensor_apply,
|
|
actor_apply=actor_apply,
|
|
critic_apply=critic_apply,
|
|
feature_extractor_apply=feature_extractor_apply,
|
|
message_passer=message_passer,
|
|
),
|
|
has_aux=True,
|
|
)
|
|
|
|
# This PPO class should be initialized only once,
|
|
# or this function will need to recompile
|
|
@partial(jax.jit, static_argnums=0)
|
|
def update_ppo(self, agent_state, storage, key):
|
|
args = self.args
|
|
ppo_loss_grad_fn = self.ppo_loss_grad_fn
|
|
|
|
def update_epoch(carry, _):
|
|
agent_state, key = carry
|
|
key, subkey = jax.random.split(key)
|
|
|
|
def flatten(x):
|
|
return x.reshape((-1,) + x.shape[2:])
|
|
|
|
def convert_data(x):
|
|
x = jax.random.permutation(subkey, x)
|
|
return jnp.reshape(x, (args.num_minibatches, -1) + x.shape[1:])
|
|
|
|
flatten_storage = jax.tree.map(flatten, storage)
|
|
shuffled_storage = jax.tree.map(convert_data, flatten_storage)
|
|
|
|
def update_minibatch(agent_state, minibatch):
|
|
(loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn(
|
|
agent_state.params,
|
|
minibatch.obs,
|
|
minibatch.actions,
|
|
minibatch.logprobs,
|
|
minibatch.advantages,
|
|
minibatch.returns,
|
|
)
|
|
agent_state = agent_state.apply_gradients(grads=grads)
|
|
return agent_state, (
|
|
loss,
|
|
pg_loss,
|
|
v_loss,
|
|
entropy_loss,
|
|
approx_kl,
|
|
grads,
|
|
)
|
|
|
|
agent_state, metrics = jax.lax.scan(update_minibatch, agent_state, shuffled_storage)
|
|
return (agent_state, key), metrics
|
|
|
|
(agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = jax.lax.scan(
|
|
update_epoch, (agent_state, key), (), length=args.update_epochs
|
|
)
|
|
return agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key
|
|
|
|
|
|
"""
|
|
Should be ok to use partial here, since the references to network,
|
|
actor and critic should not change at runtime
|
|
The cost of seperating concerns is to somehow pass these values
|
|
that are now not in the same scope
|
|
"""
|
|
|
|
|
|
@partial(jax.jit, static_argnums=(0, 1, 2, 3, 4))
|
|
def get_action_and_value(
|
|
sensor_apply,
|
|
actor_apply,
|
|
message_passer,
|
|
critic_apply,
|
|
feature_extractor_apply,
|
|
params: flax.core.FrozenDict,
|
|
x: jnp.ndarray,
|
|
action: jnp.ndarray,
|
|
):
|
|
hidden_sensor = sensor_apply(params["sensor_params"], x)
|
|
hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x)
|
|
hidden_sensor = message_passer(hidden_sensor)
|
|
mean, log_std = actor_apply(params["actor_params"], hidden_sensor)
|
|
log_std = jnp.clip(log_std, -5, 2)
|
|
std = jnp.exp(log_std)
|
|
|
|
logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
|
|
entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(-1)
|
|
value = critic_apply(params["critic_params"], hidden_critic).squeeze(-1)
|
|
|
|
return logprob, entropy, value
|
|
|
|
|
|
def ppo_loss(
|
|
params,
|
|
x,
|
|
a,
|
|
logp,
|
|
mb_advantages,
|
|
mb_returns,
|
|
args,
|
|
sensor_apply,
|
|
actor_apply,
|
|
message_passer,
|
|
critic_apply,
|
|
feature_extractor_apply,
|
|
):
|
|
newlogprob, entropy, newvalue = get_action_and_value(
|
|
sensor_apply,
|
|
actor_apply,
|
|
message_passer,
|
|
critic_apply,
|
|
feature_extractor_apply,
|
|
params,
|
|
x,
|
|
a,
|
|
)
|
|
logratio = newlogprob - logp
|
|
ratio = jnp.exp(logratio)
|
|
approx_kl = ((ratio - 1) - logratio).mean()
|
|
|
|
if args.norm_adv:
|
|
mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)
|
|
|
|
pg_loss1 = -mb_advantages * ratio
|
|
pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
|
|
pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean()
|
|
|
|
v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean()
|
|
entropy_loss = entropy.mean()
|
|
loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef
|
|
return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))
|
|
|
|
|
|
def identity(hidden):
|
|
"""
|
|
Used for seamless jax integration,
|
|
avoids having branching inside jitted function,
|
|
used as message_passer in case it is not given,
|
|
(in case of centralized lvl)
|
|
"""
|
|
|
|
return hidden
|