feat: add build system, mypy and editable source

src/brittle_star_project/MLPs/mlps.py

@@ -0,0 +1,63 @@
+from dataclasses import dataclass, fields, field
+
+import flax
+import flax.linen as nn
+import jax.numpy as jnp
+import jax.tree_util
+from typing import Sequence, Callable
+from flax.linen.initializers import constant, orthogonal
+
+
+# semi generic so we can easily make a config for it in experiments
+class GenericDenseLayersWithActivation(nn.Module):
+    layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64])
+    activation: Callable = nn.tanh
+
+    @nn.compact
+    def __call__(self, x):
+        for size in self.layer_sizes:
+            x = nn.Dense(size, kernel_init=orthogonal(jnp.sqrt(2)))(x)
+            x = self.activation(x)
+        return x
+
+
+class OneDenseLayerMLP(nn.Module):
+    @nn.compact
+    def __call__(self, x):
+        return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x)
+
+
+class Actor(nn.Module):
+    action_dim: int
+
+    @nn.compact
+    def __call__(self, x):
+        mean = nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x)
+        log_std = self.param("log_std", nn.initializers.zeros, (self.action_dim,))
+        return mean, log_std
+
+
+@jax.tree_util.register_dataclass
+@dataclass
+class AgentParams:
+    sensor_params: flax.core.FrozenDict
+    actor_params: flax.core.FrozenDict
+    critic_params: flax.core.FrozenDict
+    feature_extractor_params: flax.core.FrozenDict
+
+
+@jax.tree_util.register_dataclass
+@dataclass
+class Storage:
+    obs: jnp.array
+    actions: jnp.array
+    logprobs: jnp.array
+    dones: jnp.array
+    values: jnp.array
+    advantages: jnp.array
+    returns: jnp.array
+    rewards: jnp.array
+
+    def replace(self, **kwargs) -> "Storage":
+        fs = fields(self)
+        return Storage(**{f.name: kwargs.get(f.name, getattr(self, f.name)) for f in fs})

src/brittle_star_project/ppo.py

@@ -0,0 +1,159 @@
+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, actor, critic, feature_extractor, 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)
+    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

src/brittle_star_project/trainers/PPOTrainer.py

@@ -16,14 +16,14 @@ from experiment_logger import get_logger
 
 from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs
 from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
-from MLPs.mlps import (
+from brittle_star_project.MLPs.mlps import (
     Actor,
     AgentParams,
     GenericDenseLayersWithActivation,
     OneDenseLayerMLP,
     Storage,
 )
-from ppo import PPO
+from brittle_star_project.ppo import PPO
 
 
 @jax.jit