diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c86e484..b751166 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,4 +1,12 @@ repos: + - repo: local + hooks: + - id: ruff-format + name: ruff format (uv) + entry: uv run ruff format + language: system + types: [python] + - repo: https://github.com/alessandrojcm/commitlint-pre-commit-hook rev: v9.16.0 hooks: @@ -6,13 +14,6 @@ repos: stages: [commit-msg] additional_dependencies: ["@commitlint/config-conventional"] - - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.9.9 - hooks: - - id: ruff - args: [ --fix ] - - id: ruff-format - - repo: https://github.com/pre-commit/pre-commit-hooks rev: v5.0.0 hooks: diff --git a/ruff.toml b/ruff.toml index b19b931..cf728e1 100644 --- a/ruff.toml +++ b/ruff.toml @@ -394,4 +394,3 @@ extend-ignore = [ # "PLW1404", # implicit-str-concat ] - diff --git a/src/ppo.py b/src/ppo.py new file mode 100644 index 0000000..23c1e45 --- /dev/null +++ b/src/ppo.py @@ -0,0 +1,161 @@ +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, input_network, action_network, critic, critic_network, 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, + input_network_apply=input_network.apply, + action_network_apply=action_network.apply, + critic_apply=critic.apply, + critic_network_apply=critic_network.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_value2( + input_apply, + action_apply, + message_passer, + critic_apply, + critic_network_apply, + params: flax.core.FrozenDict, + x: jnp.ndarray, + action: jnp.ndarray, +): + hidden_network = input_apply(params["network_params"], x) + hidden_critic = critic_network_apply(params["critic_network_params"], x) + hidden_network = message_passer(hidden_network) + mean, log_std = action_apply(params["actor_params"], hidden_network) + 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, + input_network_apply, + action_network_apply, + message_passer, + critic_apply, + critic_network_apply, +): + newlogprob, entropy, newvalue = get_action_and_value2( + input_network_apply, + action_network_apply, + message_passer, + critic_apply, + critic_network_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