Repository: epignatelli/helx (language: python)

Issue:
wrap jax utils into a debugging module



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 06be0c2..27f1822 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -1,27 +1,42 @@
-name: build
+name: Import
 
 on: [push]
 
 jobs:
-  build-linux:
-    runs-on: ubuntu-latest
+  score-code:
+    runs-on: ${{ matrix.os }}-latest
+    defaults:
+      run:
+        shell: bash -l {0}
     strategy:
       fail-fast: false
       max-parallel: 5
       matrix:
-        python-version: [3.7, 3.8]
+        os: ["ubuntu", "macos"]
+    continue-on-error: true
 
     steps:
+
     - uses: actions/checkout@v2
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+
+    - uses: conda-incubator/setup-miniconda@v2
       with:
-        python-version: ${{ matrix.python-version }}
-    - name: Install dependencies
+        miniforge-version: latest
+        environment-file: environment.yml
+        activate-environment: helx
+        auto-activate-base: false
+
+    - name: Check code quality
+      id: linting
       run: |
-        # $CONDA is an environment variable pointing to the root of the miniconda directory
-        $CONDA/bin/conda install python=3.7
-        $CONDA/bin/pip install -e .
-    - name: Test with pytest
+        pip install pylint
+        MESSAGE=$(pylint -ry $(git ls-files '*.py') ||:)
+        echo "$MESSAGE"
+
+    - name: Check for reproducibility
       run: |
-        $CONDA/bin/pytest test/
+        conda info --envs
+        conda env list
+        conda list
+        pip install -e .
+        python -c "import helx"
diff --git a/environment.yml b/environment.yml
new file mode 100644
index 0000000..86741dc
--- /dev/null
+++ b/environment.yml
@@ -0,0 +1,19 @@
+name: helx
+
+dependencies:
+  - python=3.9
+  - pip
+  - pip:
+     - black
+     - flake8
+     - pytest
+     - jax
+     - chex
+     - optax
+     - rlax
+     - flax
+     - dm_env
+     - wandb
+     - absl-py
+     - gym[all]
+     - bsuite
\ No newline at end of file
diff --git a/experiments/dqn/run.py b/experiments/dqn/run.py
deleted file mode 100644
index 1d9e416..0000000
--- a/experiments/dqn/run.py
+++ /dev/null
@@ -1,139 +0,0 @@
-import jax
-from helx.jax import batch
-import helx
-from helx.rl import baselines, environment
-import wandb
-from absl import app, flags
-
-
-def main(argv):
-    flags.DEFINE_string(
-        "env",
-        "MiniGrid-Empty-6x6-v0",
-        "Environment name",
-    )
-    flags.DEFINE_integer(
-        "seed",
-        0,
-        "Random seed to control the experiment",
-    )
-    flags.DEFINE_integer(
-        "num_episodes",
-        10,
-        "Number of episodes to train on",
-    )
-    flags.DEFINE_integer(
-        "num_eval_episodes",
-        10,
-        "Number of episodes to train on",
-    )
-    flags.DEFINE_integer(
-        "batch_size",
-        32,
-        "Number of training cases over which each stochastic gradient descent (SGD) update is computed",
-    )
-    flags.DEFINE_integer(
-        "replay_memory_size",
-        5000,
-        "SGD updates are sampled from this number of most recent frames",
-    )
-    flags.DEFINE_integer(
-        "agent_history_length",
-        1,
-        "The number of most recent frames experienced by the agent that are given as  input to the Q network",
-    )
-    flags.DEFINE_integer(
-        "target_network_update_frequency",
-        10000,
-        "The frequency (measured in the number of parameters update) with which the target network is updated (this corresponds to the parameter C from Algorithm 1)",
-    )
-    flags.DEFINE_float(
-        "discount",
-        0.99,
-        "Discount factor gamma used in the Q-learning update",
-    )
-    flags.DEFINE_integer(
-        "action_repeat",
-        4,
-        "Repeat each action selected by the agent this many times. Using a value of 4 results in the agent seeing only every 4th input frame",
-    )
-    flags.DEFINE_integer(
-        "update_frequency",
-        4,
-        "The number of actions selected by the agent between successive SGD updates. Using a value of 4 results in the agent selecting 4 actions between each pair of successive updates.",
-    )
-    flags.DEFINE_float(
-        "learning_rate",
-        0.00025,
-        "The learning rate used by the RMSProp",
-    )
-    flags.DEFINE_float(
-        "gradient_momentum",
-        0.95,
-        "Gradient momentum used by the RMSProp",
-    )
-    flags.DEFINE_float(
-        "squared_gradient_momentum",
-        0.95,
-        "Squared gradient (denominator) momentum used by the RMSProp",
-    )
-    flags.DEFINE_float(
-        "min_squared_gradient",
-        0.01,
-        "Constant added to the squared gradient in the denominator of the RMSProp update",
-    )
-    flags.DEFINE_float(
-        "initial_exploration",
-        1.0,
-        "Initial value of ɛ in ɛ-greedy exploration",
-    )
-    flags.DEFINE_float(
-        "final_exploration",
-        0.01,
-        "Final value of ɛ in ɛ-greedy exploration",
-    )
-    flags.DEFINE_integer(
-        "final_exploration_frame",
-        1000000,
-        "The number of frames over which the initial value of ɛ is linearly annealed to its final value",
-    )
-    flags.DEFINE_integer(
-        "replay_start",
-        5000,
-        "A uniform random policy is run for this number of frames before learning starts and the resulting experience is used to populate the replay memory",
-    )
-    flags.DEFINE_integer(
-        "no_op_max",
-        30,
-        'Maximum numer of "do nothing" actions to be performed by the agent at the start of an episode',
-    )
-    flags.DEFINE_integer(
-        "hidden_size",
-        512,
-        "Dimension of last linear layer for value regression",
-    )
-    flags.DEFINE_integer(
-        "n_steps",
-        1,
-        "Number of timesteps for multistep return",
-    )
-
-    FLAGS = flags.FLAGS
-    hparams = baselines.dqn.HParams(
-        replay_memory_size=5000,
-        replay_start=5000,
-        batch_size=32,
-        final_exploration_frame=100000,
-    )
-
-    wandb.init()
-    env = environment.make_minigrid("MiniGrid-Empty-6x6-v0")
-    preprocess = jax.jit(lambda x: x / 255, backend="cpu")
-    dqn = baselines.dqn.Dqn(
-        (56, 56, 3), env.action_spec().num_values, hparams, preprocess=preprocess
-    )
-    return dqn.run(env, 1000000)
-
-
-if __name__ == "__main__":
-    app.run(main)
diff --git a/experiments/feature_ca/agent.py b/experiments/feature_ca/agent.py
deleted file mode 100644
index 106fe8d..0000000
--- a/experiments/feature_ca/agent.py
+++ /dev/null
@@ -1,82 +0,0 @@
-from logging import Logger
-from helx.logging import NullLogger
-from helx.rl.agent import IAgent
-from helx.rl.policy import EGreedy
-import jax
-import jax.numpy as jnp
-from dm_env import Environment, TimeStep
-
-
-class SarsaLambda:
-    """True online Sarsa(λ) algorithm for control, implemented using
-    accumulating traces and linear function approximation
-
-    Args:
-    lamda (float): trace decay hyperparameter
-    n_features (int): number of components to represent a state. Do not include actions.
-    """
-
-    def __init__(
-        self,
-        env: Environment,
-        alpha: float,
-        lamda: float,
-        n_features: int,
-        logger: Logger = NullLogger(),
-    ):
-        super().__init__()
-        self.alpha = alpha
-        self.lamda = lamda
-        self.w = jnp.zeros(n_features + 1)
-        self.z = jnp.zeros(n_features + 1)
-        self.logger = logger
-        self.policy = EGreedy(env.action_spec, 0.1)
-
-        # temporary value estimate for the starting state
-        self._actions = jnp.arange(env.action_spec().num_actions)
-        self._q_old = jnp.zeros(env.action_spec().num_values)
-
-    def phi(self, x):
-        return jnp.dot(x, self.w)
-
-    def q_values(self, s: jnp.ndarray):
-        x = jnp.vmap(
-            lambda a: jnp.concatenate([s, jnp.ones(s.shape[:2]) * a], axis=-1)
-        )(self._actions)
-        return jnp.dot(self.w.T, x)
-
-    def select_action(self, timestep: TimeStep):
-        q = self.q_values(timestep.observation)
-        return self.policy.sample(q)
-
-    def update(self, timestep: TimeStep, action: int, new_timestep: TimeStep) -> float:
-        """True online Sarsa(λ) update using accumulating traces.
-        See Sutton R., Barto, G. 2018. Reinforcement Learning: an Introduction, pp. 300-306"""
-        g, l, a = timestep.discount, self.lamda, self.alpha
-        x_m1, x = timestep.observation, new_timestep.observation
-        q_m1, q = self.q_values(x_m1), self.q_values(x)
-
-        # as we are using TD(λ) for control
-        td = new_timestep.reward + g * q - q_m1
-
-        # update traces using accumulation
-        self.z = g * l * self.z + (1 - a * g * l * jnp.dot(self.z.T, x_m1)) * x_m1
-
-        # update weights
-        self.w = (
-            self.w
-            + a(td + q_m1 - self._q_old) * self.z
-            + a * (q_m1 - self._q_old) * x_m1
-        )
-
-        # update value estimate
-        error = q - self._q_old
-        self._q_old = int(not new_timestep.last()) * q
-
-        # log
-        self.logger.log({"loss": error})
-
-        return error
-
-    def log(self, timestep, action, new_timestep, loss):
-        return
diff --git a/experiments/feature_ca/environment.py b/experiments/feature_ca/environment.py
deleted file mode 100644
index 5a2e4aa..0000000
--- a/experiments/feature_ca/environment.py
+++ /dev/null
@@ -1,131 +0,0 @@
-import gym
-from gym.spaces.box import Box
-from gym_minigrid.wrappers import ImgObsWrapper
-import numpy as np
-from gym import spaces
-from gym_minigrid.minigrid import OBJECT_TO_IDX, Goal, Grid, Lava, MiniGridEnv
-
-from helx.rl import agent
-
-
-class PartialObsWrapper(gym.core.ObservationWrapper):
-    """
-    Fully observable grid with a symbol for the state of each cell.
-    The symbol is a triple of (X, Y, IDX), where IDX is
-    the id of the object in the cell. Agent direction is removed.
-    """
-
-    def __init__(self, env, agent_view_size=3):
-        super().__init__(env)
-
-        assert (
-            "image" in self.observation_space.spaces
-        ), "Observation does not contain an image field."
-
-        # set agent view size
-        self.unwrapped.agent_view_size = agent_view_size
-        topX, topY, botX, botY = self.unwrapped.get_view_exts()
-
-        # set partial observation space
-        current = self.unwrapped.observation_space["image"]
-        self.env.observation_space["image"] = Box(
-            low=current.low[topX:botX, topY:botY],
-            high=current.high[topX:botX, topY:botY],
-            shape=(agent_view_size, agent_view_size, *current.shape[2:]),
-            dtype=current.dtype,
-        )
-
-    def observation(self, obs):
-        topX, topY, botX, botY = self.unwrapped.get_view_exts()
-        obs["image"] = obs["image"][topX:botX, topY:botY]
-        return obs
-
-
-class SymbolicObsWrapper(gym.core.ObservationWrapper):
-    """
-    Fully observable grid with a symbol for the state of each cell.
-    The symbol is a triple of (X, Y, IDX), where IDX is
-    the id of the object in the cell. Agent direction is removed.
-    """
-
-    def __init__(self, env):
-        super().__init__(env)
-
-        current = self.env.observation_space["image"]
-        self.observation_space.spaces["image"] = spaces.Box(
-            low=min(OBJECT_TO_IDX.values()),
-            high=max(OBJECT_TO_IDX.values()),
-            shape=(self.env.width, self.env.height),
-            dtype="uint8",
-        )
-
-    def observation(self, obs):
-        w, h = self.width, self.height
-        objects = np.array(
-            [OBJECT_TO_IDX[o.type] if o is not None else -1 for o in self.grid.grid]
-        ).reshape(w, h)
-        # grid = np.mgrid[:w, :h]
-        # grid = np.concatenate([grid, objects])
-        # grid = np.transpose(grid, (1, 2, 0))
-        obs["image"] = objects
-        return obs
-
-
-class EmptyMultigoal(MiniGridEnv):
-    def __init__(
-        self,
-        size=8,
-        agent_start_pos=None,
-        agent_start_dir=None,
-        n_goals=2,
-        n_traps=1,
-    ):
-        self.n_goals = n_goals
-        self.n_traps = n_traps
-        self.agent_start_pos = agent_start_pos
-        self.agent_start_dir = agent_start_dir
-
-        size += 2
-        super().__init__(
-            grid_size=size,
-            max_steps=4 * size * size,
-            # Set this to True for maximum speed
-            see_through_walls=True,
-            agent_view_size=size * 2 + 1,  # init as fully observable
-        )
-
-    def _gen_grid(self, width, height):
-        # Create an empty grid
-        self.grid = Grid(width, height)
-
-        # Generate the surrounding walls
-        self.grid.wall_rect(0, 0, width, height)
-
-        # Place the goals
-        for _ in range(self.n_goals):
-            self.place_obj(Goal())
-
-        # Place the traps
-        for _ in range(self.n_traps):
-            self.place_obj(Lava())
-
-        # Place the agent
-        if self.agent_start_pos is not None:
-            self.agent_pos = self.agent_start_pos
-            self.agent_dir = self.agent_start_dir
-        else:
-            self.place_agent()
-
-        self.mission = "get to the green goal square, avoid the lava"
-
-
-if __name__ == "__main__":
-    # debugging
-    env = EmptyMultigoal(size=5, n_goals=1, n_traps=1)
-    env = SymbolicObsWrapper(env)
-    env = PartialObsWrapper(env, agent_view_size=1)
-    env = ImgObsWrapper(env)
-    o = env.reset()
-    print(o.shape)
-    env.render()
-    print(o)
diff --git a/experiments/feature_ca/main.py b/experiments/feature_ca/main.py
deleted file mode 100644
index 5d90550..0000000
--- a/experiments/feature_ca/main.py
+++ /dev/null
@@ -1,49 +0,0 @@
-import helx.logging
-import helx.rl
-import helx.rl.experiment
-import jax.numpy as jnp
-from absl import app, flags
-from gym_minigrid.wrappers import ImgObsWrapper
-
-import agent as a_lib
-import environment as e_lib
-
-
-flags.DEFINE_integer("size", 8, "Size of the minigrid", 3)
-flags.DEFINE_integer(
-    "n_goals", 2, "Number of goals that will spawn in the environment in each episode"
-)
-flags.DEFINE_integer(
-    "n_traps", 1, "Number of traps that will spawn in the environment in each episode"
-)
-flags.DEFINE_float("alpha", 0.1, "Step size parameter")
-flags.DEFINE_float(
-    "lamda", 0.9, "Trace decay paremterer for the Sarsa(λ) algorithm", 0, 1
-)
-flags.DEFINE_integer(
-    "train_episodes", 1_000_000, "Max number of episodes to learn from"
-)
-flags.DEFINE_integer(
-    "eval_episodes", 1_000_000, "Max number of episodes to evaluate on"
-)
-FLAGS = flags.FLAGS
-
-
-def main(argv):
-    env = e_lib.EmptyMultigoal(
-        size=FLAGS.size, n_goals=FLAGS.n_goals, n_traps=FLAGS.n_traps
-    )
-    env = e_lib.SymbolicObsWrapper(env)
-    env = ImgObsWrapper(env)
-    env = helx.rl.environment.from_gym(env)
-
-    n_features = jnp.prod(env.observation_spec().shape)
-    logger = helx.logging.TerminalLogger()
-    agent = a_lib.SarsaLambda(env, FLAGS.alpha, FLAGS.lamda, n_features, logger)
-
-    helx.rl.experiment.run(env, agent, FLAGS.train_episodes)
-    helx.rl.experiment.run(env, agent, FLAGS.eval_episodes, True)
-
-
-if __name__ == "__main__":
-    app.run(main)
diff --git a/helx/__init__.py b/helx/__init__.py
index 8e1032b..d17c320 100644
--- a/helx/__init__.py
+++ b/helx/__init__.py
@@ -1,9 +1,8 @@
-from . import typing
-from . import jax
-from . import logging
-from . import random
+from .agents import agent, dqn, sac, sacd
+from . import environment
+from . import experiment
 from . import image
-from . import distributed
-from . import nn
-from . import rl
-from . import optimise
+from .mdp import Episode
+from .memory import ReplayBuffer
+from . import random
+from . import stax
diff --git a/helx/agents/__init__.py b/helx/agents/__init__.py
new file mode 100644
index 0000000..71a20bf
--- /dev/null
+++ b/helx/agents/__init__.py
@@ -0,0 +1,4 @@
+from .agent import Agent
+from .dqn import DQN, DQNhparams
+from .sac import SAC, SAChparams
+from .sacd import SACD
\ No newline at end of file
diff --git a/helx/agents/agent.py b/helx/agents/agent.py
new file mode 100644
index 0000000..3786772
--- /dev/null
+++ b/helx/agents/agent.py
@@ -0,0 +1,122 @@
+"""Base class for a Reinforcement Learning agent."""
+# pyright: reportPrivateImportUsage=false
+from __future__ import annotations
+
+import abc
+from typing import Any, NamedTuple, Tuple
+
+import jax
+from chex import Array, PyTreeDef
+from dm_env import Environment
+
+from ..mdp import Episode
+
+
+class Hparams(NamedTuple):
+    ...
+
+
+class Agent():
+    def __init__(self):
+        self.loss = jax.jit(self.loss)
+
+        if not hasattr(self, "hparams"):
+            self.hparams = Hparams()
+        if not hasattr(self, "iteration"):
+            self.iteration = 0
+
+    @abc.abstractmethod
+    def loss(
+        self, params: PyTreeDef, sarsa: Episode, *args, **kwargs
+    ) -> Tuple[Array, Any]:
+        """The loss function to differentiate through for Deep RL agents.
+            This function can be used for heavy computation and can be xla-compiled,
+            and is always *jitted by default*.
+            A possible design pattern is with decorators:
+        Example:
+        ```
+            >>> @partial(jax.value_and_grad, argnums=1, has_aux=True)
+            >>> def loss(self, params, sarsa, *, **kwargs):
+            >>>     s, a, r_tp1, s_tp1, d_tp1, a_tp1 = sarsa
+            >>>     return rlax.l2_loss(self.network.apply(params, s)).mean()
+        ```
+        Args:
+            params (PyTreeDef): A pytree of function parameters.
+            sarsa (Tuple): A tuple of 6 elements (sₜ, aₜ, rₜ₊₁, sₜ₊₁, dₜ₊₁, aₜ₊₁)
+                that defines MDP transitions from a state sₜ through action aₜ,
+                to a state sₜ₊₁, while collecting reward rₜ₊₁. The term dₜ₊₁
+                indicates whether the state sₜ₊₁ is terminal,
+                and aₜ₊₁ is an optional on-policy action.
+
+        Returns:
+            Returns a tuple of two elements containing, respectively:
+            the calculated loss as a scalar, and any auxiliary variable to carry forward
+            in the computation. Notice that once wrapped around `jax.value_and_grad`, the return
+            type becomes a nested tuple `Tuple[Tuple[float, Any], PyTreeDef]`, with the last term
+            being a `PyTree` of gradients with the same structure of `params`
+        """
+
+    @abc.abstractmethod
+    def policy(
+        self, observation: Array, eval: bool = False, **kwargs
+    ) -> Tuple[Array, Array]:
+        """Collects a sample from the policy distribution conditioned on the provided observation π(s), and
+            returns both the sampled action, and a (possibly empty) array of the correspondent log probabilities.
+        Args:
+            observation (Array): The observation to condition onto.
+            keys (PRNGKey): a random key used to sample the action
+            eval (bool): optional boolean if the agent follows a different policy at evaluation time
+        Returns:
+            Tuple[Array, Array]: a tuple containing
+                i) the sampled action and
+                ii) the log probability of the sampled action, for one-action-out architectures
+                    the log probabilities for all actions for all-actions-out architectures"""
+
+    @abc.abstractmethod
+    def update(self, episode: Episode) -> Array:
+        """Updates the agent state at the end of each episode and returns the loss as a scalar.
+        This function can used to update both the agent's parameters and its memory.
+        This function is usually not jittable, as we can not ensure that the agent memory, and other
+        properties are jittable. This is also a good place to perform logging."""
+
+    def new_key(self, n: int = 1):
+        self.key, k = jax.random.split(self.key)
+        return jax.random.split(k, n).squeeze()
+
+    def unroll(
+        self, env: Environment, eval: bool = False, max_steps: int = int(2e9)
+    ) -> Episode:
+        """Deploys the agent in the environment for a full episode.
+            In case of batched environments, the each property of the episode
+            has an additional `batch` axis at index `0`.
+            The episode terminates only for absorbing states, and the number of
+            maximum steps can be specified in the environment itself.
+            While using `experiment.run` the result of this function is passed
+            to the `update` method.
+        Args:
+            env (dm_env.Environment): the environment to interact with implements
+                the `dm_env.Environment` interface, and NOT a `gym` interface.
+                You can wrap the former around the latter using  the
+                `bsuite.utils.gym_wrapper.DMFromGymWrapper` wrapper.
+            eval (bool): eval flag passed to the `policy` method.
+        Returns:
+            (Episode): an full episode following the current policy.
+                Each property of the episode has an additional `batch` axis at index `0`.
+                in case of batched environments.
+        """
+        t = 0
+        timestep = env.reset()
+        episode = Episode.start(timestep)
+        while (not timestep.last()) and t < max_steps:
+            t += 1
+            action, _ = self.policy(timestep.observation, eval=eval)
+            timestep = env.step(action.item())
+            episode.add(timestep, action)
+        return episode
+
+    def save(self, path):
+        raise NotImplementedError
+
+    @classmethod
+    def load(cls, path):
+        raise NotImplementedError
diff --git a/helx/agents/dqn.py b/helx/agents/dqn.py
new file mode 100644
index 0000000..c210c33
--- /dev/null
+++ b/helx/agents/dqn.py
@@ -0,0 +1,140 @@
+# pyright: reportPrivateImportUsage=false
+from __future__ import annotations
+
+from functools import partial
+from typing import NamedTuple, Tuple
+
+import distrax
+import jax
+import jax.numpy as jnp
+import optax
+import rlax
+import wandb
+from chex import Array, Shape
+from flax import linen as nn
+from optax import GradientTransformation
+
+from .agent import Agent
+from ..mdp import Episode
+from ..memory import ReplayBuffer
+
+
+class DQNhparams(NamedTuple):
+    # network
+    input_shape: Shape
+    hidden_size: int = 128
+    # rl
+    initial_exploration: float = 1.0
+    final_exploration: float = 0.01
+    final_exploration_frame: int = 1000000
+    replay_start: int = 1000
+    replay_memory_size: int = 1000
+    update_frequency: int = 1
+    target_network_update_frequency: int = 10000
+    discount: float = 0.99
+    n_steps: int = 1
+    # sgd
+    batch_size: int = 32
+    learning_rate: float = 0.00025
+    gradient_momentum: float = 0.95
+    squared_gradient_momentum: float = 0.95
+    min_squared_gradient: float = 0.01
+
+
+class DQN(Agent):
+    """Implements a Deep Q-Network:
+    Mnih, Volodymyr, et al. "Human-level control through deep reinforcement learning."
+    Nature 518.7540 (2015): 529-533.
+    https://www.nature.com/articles/nature14236"""
+
+    def __init__(
+        self,
+        network: nn.Module,
+        optimiser: GradientTransformation,
+        hparams: DQNhparams,
+        seed: int,
+    ):
+        key = jax.random.PRNGKey(seed)
+        key, k1 = jax.random.split(key)
+        input_shape = hparams.input_shape
+        params = network.init(k1, jnp.ones(input_shape))
+
+        # const:
+        self.key = key
+        self.network = network
+        self.optimiser = optimiser
+        self.hparams = hparams
+        self.memory = ReplayBuffer(hparams.replay_memory_size)
+        self.iteration = 0
+        self.params = params
+        self.params_target = params.copy({})
+        self.opt_state = optimiser.init(params)
+
+    def epsilon(self, eval=False) -> Array:
+        x0, y0 = (self.hparams.replay_start, self.hparams.initial_exploration)
+        x1 = self.hparams.final_exploration_frame
+        y1 = self.hparams.final_exploration
+        x = self.iteration
+        y = ((y1 - y0) * (x - x0) / (x1 - x0)) + y0
+        eps = jnp.clip(
+            y, self.hparams.final_exploration, self.hparams.initial_exploration
+        )
+        return jax.lax.select(eval, 0.0, eps)
+
+    def policy(self, observation: Array, eval=False) -> Tuple[Array, Array]:
+        """Selects an action using an e-greedy policy"""
+        q_values = self.network.apply(self.params, observation)  # type: ignore
+        action, log_probs = distrax.EpsilonGreedy(
+            q_values, self.epsilon(eval)  # type: ignore
+        ).sample_and_log_prob(seed=self.new_key())
+        return action, log_probs
+
+    @partial(jax.value_and_grad, argnums=1, has_aux=True)
+    def loss(self, params, transitions_batch, params_target):
+        Q = jax.vmap(self.network.apply, in_axes=(None, 0))  # type: ignore
+        s_0, a_0, r_1, s_1, d = transitions_batch
+        q_1 = jax.lax.stop_gradient(jnp.max(Q(params_target, s_1), axis=-1))  # type: ignore
+        q_0 = Q(params, s_0)
+        td_error = r_1 + (1 - d) * self.hparams.discount * jnp.max(q_1) - q_0[a_0]
+        loss = jnp.mean(rlax.l2_loss(td_error))
+        return loss, ()
+
+    def update(self, episode: Episode) -> Array:
+        # increment iteration
+        self.iteration += 1
+        wandb.log({"Iteration": self.iteration})
+
+        # update memory
+        transitions = episode.sars()
+        self.memory.add_range(transitions)
+        wandb.log({"Buffer size": len(self.memory)})
+
+        # if replay buffer is smaller than the minimum size, there is nothing else to do
+        if len(self.memory) < self.hparams.replay_start:
+            return jnp.asarray([])
+
+        # update every `update_frequency` steps
+        if self.iteration % self.hparams.update_frequency != 0:
+            return jnp.asarray([])
+
+        episode_batch = self.memory.sample(self.hparams.batch_size)
+        (loss, _), grads = self.loss(
+            self.params,
+            episode_batch,
+            self.params_target,
+        )
+        updates, opt_state = self.optimiser.update(grads, self.opt_state)
+
+        # update dqn state
+        self.opt_state = opt_state
+        self.params = optax.apply_updates(self.params, updates)
+        self.params_target = rlax.periodic_update(
+            self.params,
+            self.params_target,
+            jnp.asarray(self.iteration),
+            self.hparams.target_network_update_frequency,
+        )
+
+        wandb.log({"train/total_loss": loss.item()})
+        wandb.log({"train/Return": jnp.sum(episode.r).item()})  # type: ignore
+        return loss
diff --git a/helx/agents/sac.py b/helx/agents/sac.py
new file mode 100644
index 0000000..075d0c9
--- /dev/null
+++ b/helx/agents/sac.py
@@ -0,0 +1,199 @@
+# pyright: reportPrivateImportUsage=false
+from __future__ import annotations
+
+from typing import NamedTuple, Tuple
+
+import distrax
+import jax
+import jax.numpy as jnp
+import optax
+import rlax
+import wandb
+from chex import Array, PRNGKey, Shape
+from flax import linen as nn
+from jax.lax import stop_gradient
+from optax import GradientTransformation
+
+from .agent import Agent
+from ..mdp import Episode
+from ..memory import ReplayBuffer
+
+
+class SAChparams(NamedTuple):
+    # network
+    input_shape: Shape
+    tau: float = 0.005
+    # rl
+    replay_start: int = 1000
+    replay_memory_size: int = 1000
+    update_frequency: int = 1
+    target_network_update_frequency: int = 10000
+    discount: float = 0.99
+    n_steps: int = 1
+    entropy_rewards: bool = False
+    # sgd
+    batch_size: int = 32
+    learning_rate: float = 3e-4
+
+
+class Temperature(nn.Module):
+    initial_temperature: float = 1.0
+
+    @nn.compact
+    def __call__(self) -> Array:
+        log_temperature = self.param(
+            "log_temperature",
+            init_fn=lambda key: jnp.full((), jnp.log(self.initial_temperature)),
+        )
+        return jnp.exp(log_temperature)
+
+
+class SAC(Agent):
+    """Implements Soft Actor-Critic with gaussian policy for continuous action spaces.
+    Haarnoja, Tuomas, et al. "Soft actor-critic: Off-policy maximum entropy deep
+        reinforcement learning with a stochastic actor."
+        International conference on machine learning. PMLR, 2018.
+        https://arxiv.org/abs/1801.01290 for the base SAC.
+    Haarnoja, Tuomas, et al. "Soft actor-critic algorithms and applications."
+        arXiv preprint arXiv:1812.05905 (2018).
+        https://arxiv.org/abs/1812.05905 for the tunable entropy temperature
+    Haonanm Yu, et al. "Do you need reward entropy (in practice)?".
+        arXiv preprint arXiv:2201.12434. 2022 Jan 28.
+        https://arxiv.org/abs/2201.12434
+    """
+    def __init__(
+        self,
+        actor: nn.Module,
+        critic: nn.Module,
+        optimiser: GradientTransformation,
+        hparams: SAChparams,
+        seed: int,
+    ):
+        key = jax.random.PRNGKey(seed)
+        key, k1, k2, k3 = jax.random.split(key, 4)
+        temperature = Temperature()
+        input_tensor = jnp.ones(hparams.input_shape)
+        output_tensor, params_actor = actor.init_with_output(k1, input_tensor)
+        params_critic = critic.init(k2, input_tensor)
+        params_temperature = temperature.init(k3)
+
+        self.key = key
+        self.actor = actor
+        self.critic = critic
+        self.optimiser = optimiser
+        self.temperature = temperature
+        self.hparams = hparams
+        self.memory = ReplayBuffer(hparams.replay_memory_size)
+        self.iteration = 0
+        self.dim_A = jnp.ndim(output_tensor)
+        self.params = (params_actor, params_critic, params_temperature)
+        self.opt_state = optimiser.init(self.params)
+        self.params_critic_target = params_critic.copy({})
+        super().__init__()
+
+    def policy(self, observation: Array, eval: bool = False) -> Tuple[Array, Array]:
+        """Selects an action using a parameterise gaussian policy"""
+        (params_actor, _, _) = self.params
+        return self._policy(params_actor, observation, self.new_key())  # type: ignore
+
+    def _policy(self, params_actor: nn.FrozenDict, observation: Array, key: PRNGKey):
+        mu, logvar = self.actor.apply(params_actor, observation)
+        noise = jax.random.normal(key, (self.dim_A,))
+        action = jnp.tanh(mu + logvar * noise)
+        logprob = distrax.Normal(mu, jnp.exp(logvar)).log_prob(action)  # type: ignore
+        return action, logprob
+
+    def loss(self, params, transitions_batch, keys, params_critic_target):
+        print("{} compiles".format(self.__class__.__name__))
+        s_0, _, r_1, s_1, d = transitions_batch
+        (params_actor, params_critic, params_temperature) = params
+
+        _policy = jax.vmap(self._policy(p, o, k), in_axes=(None, 0, 0))  # type: ignore
+        _value = jax.vmap(self.critic.apply, in_axes=(None, 0))  # type: ignore
+
+        # current estimates
+        temperature = self.temperature.apply(params_temperature)
+        alpha = stop_gradient(temperature)
+        _, logprobs_a_0 = _policy(params_actor, s_0, keys)
+        _, logprobs_a_1 = _policy(params_actor, s_1, keys)
+        probs_a_0 = jnp.exp(logprobs_a_0)
+        qA_0, qB_0 = _value(params_critic, s_0)
+        qA_1, qB_1 = _value(params_critic_target, s_1)
+
+        # augment reward with policy entropy
+        policy_entropy = -jnp.sum(probs_a_0 * logprobs_a_0, axis=-1)  # type: ignore
+        # SACLite: https://arxiv.org/abs/2201.12434
+        policy_entropy = policy_entropy * self.hparams.entropy_rewards
+        r_1 = r_1 + alpha * policy_entropy
+
+        # actor loss
+        q_0 = stop_gradient(jnp.min(jnp.stack([qA_0, qB_0], axis=0), axis=0))
+        actor_loss = alpha * logprobs_a_0 - q_0
+
+        # critic losses
+        q_1 = jnp.min(jnp.stack([qA_1, qB_1], axis=0), axis=0)
+        v_1 = q_1 - alpha * logprobs_a_1
+        td_target = stop_gradient(r_1 + (1 - d) * self.hparams.discount * v_1)
+        critic_loss = rlax.l2_loss(qA_0, td_target) + rlax.l2_loss(qB_0, td_target)  # type: ignore
+
+        # temperature loss
+        target_entropy = -self.dim_A
+        logprobs_a_0 = stop_gradient(logprobs_a_0)
+        temperature_loss = -(temperature * (logprobs_a_0 + target_entropy))
+
+        # average over batches
+        actor_loss = actor_loss.mean()
+        critic_loss = critic_loss.mean()
+        temperature_loss = temperature_loss.mean()
+        policy_entropy = policy_entropy.mean()
+        loss = actor_loss + critic_loss + temperature_loss
+        aux = (actor_loss, critic_loss, temperature_loss, policy_entropy, alpha)
+        return jnp.array(loss, dtype=float), aux
+
+    def update(self, episode: Episode) -> Array:
+        # update iteration
+        self.iteration += 1
+        wandb.log({"Iteration": self.iteration})
+
+        # update memory
+        transitions = episode.sars()
+        self.memory.add_range(transitions)  # makes self.loss recompile each time, why??
+        wandb.log({"Buffer size": len(self.memory)})
+
+        # if replay buffer is smaller than the minimum size, there is nothing else to do
+        if len(self.memory) < self.hparams.replay_start:
+            return jnp.asarray([])
+
+        # check update period
+        if self.iteration % self.hparams.update_frequency != 0:
+            return jnp.asarray([])
+
+        episode_batch = self.memory.sample(self.hparams.batch_size)
+        keys = self.new_key(self.hparams.batch_size)
+        (loss, aux), grads = self.loss(
+            self.params,
+            episode_batch,
+            keys,
+            self.params_critic_target,
+        )
+        updates, opt_state = self.optimiser.update(grads, self.opt_state)
+        actor_loss, critic_loss, temperature_loss, policy_entropy, alpha = aux
+
+        # update dqn state
+        self.opt_state = opt_state
+        self.params = optax.apply_updates(self.params, updates)
+        self.params_critic_target = jax.tree_map(
+            lambda theta, theta_: theta * self.hparams.tau
+            + (1 - self.hparams.tau) * theta_,
+            self.params[1],  # type: ignore
+            self.params_critic_target,
+        )
+
+        wandb.log({"train/total_loss": loss.item()})
+        wandb.log({"train/actor_loss": actor_loss.item()})
+        wandb.log({"train/critic_loss": critic_loss.item()})
+        wandb.log({"train/temperature_loss": temperature_loss.item()})
+        wandb.log({"train/policy_entropy": policy_entropy.item()})
+        wandb.log({"train/alpha": alpha.item()})
+        wandb.log({"train/Return": jnp.sum(episode.r).item()})  # type: ignore
+        return loss
diff --git a/helx/agents/sacd.py b/helx/agents/sacd.py
new file mode 100644
index 0000000..76a160b
--- /dev/null
+++ b/helx/agents/sacd.py
@@ -0,0 +1,90 @@
+# pyright: reportPrivateImportUsage=false
+from __future__ import annotations
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+import rlax
+from jax.lax import stop_gradient
+
+from .sac import SAC
+
+
+class SACD(SAC):
+    """Soft Actor-Critic agent with softmax policy for discrete action spaces
+    and a twinned Q-network for value approximation with a separate target network
+    updated by polyak averaging the online network parameters.
+    See:
+        1. Haarnoja T., 2018 (https://arxiv.org/abs/1801.01290)
+            for the base SAC
+        2. Haarnoja T., 2019 (https://arxiv.org/abs/1812.05905)
+            for the tunable entropy temperature and the ablation of the value network
+        3. Christodoulou P., 2019 (https://arxiv.org/abs/1910.07207)
+            for the adaptation to discrete action spaces
+    """
+
+    def _policy(self, params_actor, observation, key):
+        logits = self.actor.apply(params_actor, observation)
+        logprobs = jax.nn.log_softmax(logits)
+        action = jax.random.categorical(key, logits)  # type: ignore
+        return action, logprobs
+
+    @partial(jax.value_and_grad, argnums=1, has_aux=True)
+    def loss(
+        self,
+        params,
+        sarsd_batch,
+        keys,
+        params_critic_target,
+    ):
+        print(f"{self.__class__.__name__}.loss compiles")  # prints only if compiled
+        s_0, _, r_1, s_1, d = sarsd_batch
+        (params_actor, params_critic, params_temperature) = params
+
+        _policy = jax.vmap(self._policy, in_axes=(None, 0, 0))  # type: ignore
+        _value = jax.vmap(self.critic.apply, in_axes=(None, 0))  # type: ignore
+        batch_matmul = lambda a, b: jnp.sum(a * b, axis=-1, keepdims=True)
+
+        # current estimates
+        temperature = self.temperature.apply(params_temperature)
+        alpha = stop_gradient(temperature)
+        _, logprobs_a_0 = _policy(params_actor, s_0, keys)
+        _, logprobs_a_1 = _policy(params_actor, s_1, keys)
+        probs_a_0 = jnp.exp(logprobs_a_0)
+        probs_a_1 = jnp.exp(logprobs_a_1)
+        qA_0, qB_0 = _value(params_critic, s_0)
+        qA_1, qB_1 = _value(params_critic_target, s_1)
+
+        # augment reward with policy entropy
+        policy_entropy = -jnp.sum(probs_a_0 * logprobs_a_0, axis=-1)  # type: ignore
+        # SACLite: https://arxiv.org/abs/2201.12434
+        policy_entropy = policy_entropy * self.hparams.entropy_rewards
+        r_1 = r_1 + alpha * policy_entropy
+
+        # actor target: V(sₜ) = π(sₜ)ᵀ • [αlog(π(sₜ)) + Q(sₜ)]
+        q_0 = stop_gradient(jnp.min(jnp.stack([qA_0, qB_0], axis=0), axis=0))
+        actor_loss = batch_matmul(probs_a_0, alpha * logprobs_a_0 - q_0)  # type: ignore
+
+        # critic target: V(sₜ₊₁) = π(sₜ₊₁)ᵀ • [Q(sₜ₊₁) - αlog(π(sₜ₊₁))]
+        probs_a_1 = jnp.exp(logprobs_a_1)
+        q_1 = jnp.min(jnp.stack([qA_1, qB_1], axis=0), axis=0)
+        v_1 = batch_matmul(probs_a_1, (q_1 - alpha * logprobs_a_1))  # type: ignore
+        q_target = stop_gradient(r_1 + (1 - d) * self.hparams.discount * v_1)
+        critic_loss = rlax.l2_loss(qA_0, q_target) + rlax.l2_loss(qB_0, q_target)  # type: ignore
+
+        # temperature target: π(sₜ)ᵀ • [- αlog(π(sₜ)) + H]
+        target_entropy = -jnp.log(1 / self.dim_A)
+        logprobs_a_0 = stop_gradient(logprobs_a_0)
+        probs_a_0 = stop_gradient(probs_a_0)
+        temperature_loss = -jnp.log(temperature) * (logprobs_a_0 + target_entropy)  # type: ignore
+        temperature_loss = batch_matmul(probs_a_0, temperature_loss)
+
+        # average over batches
+        actor_loss = actor_loss.mean()  # type: ignore
+        critic_loss = critic_loss.mean()
+        temperature_loss = temperature_loss.mean()  # type: ignore
+        policy_entropy = -jnp.sum(probs_a_0 * logprobs_a_0, axis=-1).mean()  # type: ignore
+        loss = actor_loss + critic_loss + temperature_loss
+        aux = (actor_loss, critic_loss, temperature_loss, policy_entropy, alpha)
+        return loss, aux
diff --git a/helx/distributed.py b/helx/distributed.py
deleted file mode 100644
index e2bfcd2..0000000
--- a/helx/distributed.py
+++ /dev/null
@@ -1,84 +0,0 @@
-import jax
-import jax.numpy as jnp
-import threading
-
-
-def distribute_tree(tree):
-    """
-    This is usually used to prepare the params for `pmap`ped functions.
-    Adds an additional axis to the tree, by duplicating the tree as many times
-    as the number returned by `jax.local_device_count()`.
-    """
-    return jax.tree_map(lambda x: jnp.array([x] * jax.local_device_count()), tree)
-
-
-def gather_tree(tree):
-    """
-    This is usually used to reduce back a tree that has been distributed
-    using `distribute_tree` or `redistribute_tree`.
-    Assumes that the tree is duplicated along the axis 0 an arbitrary number of times,
-    and selects only one of the replicas of the pyTree
-    """
-    return jax.tree_map(lambda x: x[0], tree)
-
-
-def redistribute_tree(tree):
-    """
-    This is usually used to prepare the params for `pmap`ped functions when the params
-    are coming from a machine with a different number of devices.
-    Assumes that the tree is duplicated along the axis 0 an arbitrary number of times.
-    The tree usually is a result of a the `distributed_tree` function
-    on a machine with a different number of devices.
-    This function sets the number of duplicates equal to the
-     number returned by `jax.local_device_count()`.
-    """
-    return distribute_tree(gather_tree(tree))
-
-
-def distribute_array(array):
-    """
-    Prepares an array for `jax.pmap` by adding a leading axis
-    whose dimension are equal to `jax.local_device_count()`.
-    As a results the dimensions of the first axis pre-modification are divided by
-    `jax.local_device_count()`. Note that the first dimension needs to be
-    divisible by `jax.local_device_count()`.
-    """
-    n = jax.local_device_count()
-    s = array.shape
-    assert (
-        len(s) > 1 and s[0] % n == 0
-    ), "Array dimension {} must be divisible for the number of local devices {}".format(
-        s[0], n
-    )
-    new_shape = (n,) + (s[0] // n,) + s[1:]
-    return array.reshape(new_shape)
-
-
-def gather_array(array):
-    """
-    Gathers an array that has previously been distributed using `distribute_array`,
-    by removing the leading axis and stacking them on the second axis.
-    """
-    s = array.shape
-    assert len(s) > 2, "Array must have at least 3 axes"
-    return array.reshape(s[0] + s[1], s[2:])
-
-
-def redistribute_array(array):
-    """
-    Gathers and then distributes again a given array.
-    See `helx.distributed.distribute_array` and `helx.distributed.gather_array`
-    for more
-    """
-    return distribute_array(gather_array(array))
-
-
-def asynchronous(f):
-    """Decorator to run a function asynchronously"""
-
-    def run(*k, **kw):
-        t = threading.Thread(target=f, args=k, kwargs=kw, name=f.__name__)
-        t.start()
-        return
-
-    return run
diff --git a/helx/rl/environment.py b/helx/environment.py
similarity index 76%
rename from helx/rl/environment.py
rename to helx/environment.py
index 639753e..a67688d 100644
--- a/helx/rl/environment.py
+++ b/helx/environment.py
@@ -1,30 +1,46 @@
+"""A set of functions to define RL environments."""
 import logging
 import multiprocessing as mp
 from copy import deepcopy
 from multiprocessing.connection import Connection
-from typing import Sequence
+from typing import Any, List, Sequence
 import matplotlib.pyplot as plt
 
 import dm_env
 import gym
+import gym.utils.seeding
 import jax
 import jax.numpy as jnp
 from bsuite.utils.gym_wrapper import DMEnvFromGym
-from gym_minigrid.wrappers import *
 from helx.image import greyscale, imresize
 from helx.random import PRNGSequence
-from helx.typing import Size
+from chex import Shape
+
+
+class Environment(DMEnvFromGym):
+    """A wrapper around a gym environment to make it compatible with dm_env
+    with two additional methods: i.e.
+      - `seed(self, seed: int) -> None` and
+      - `render(self)`.
+    """
+
+    def seed(self, seed: int) -> None:
+        if seed is not None:
+            self.gym_env.np_random, seed = gym.utils.seeding.np_random(seed)
+        return
+
+    def render(self):
+        return self.gym_env.render()
 
 
 def from_gym(gym_env):
+    """Convert a gym environment to dm_env.Environment"""
     #  Convert to dm_env.Environment
-    env = DMEnvFromGym(gym_env)
-    # fix gym's env.render function
-    env.render = lambda: plt.imshow(env.gym_env.render("rgb_array"))
-    return env
+    return Environment(gym_env)
 
 
 def make(name):
+    """Create an helx.Environment from a gym environment name"""
     env = gym.make(name)
     return from_gym(env)
 
@@ -40,14 +56,22 @@ def preprocess_atari(x):
     return greyscale(imresize(x, (84, 84)))
 
 
-def preprocess_minigrid(x, size: Size = (56, 56)):
+def preprocess_minigrid(x, size: Shape = (56, 56)):
     """Refer to the minigrid implementation at:
-    https://github.com/maximecb/gym-minigrid
+    https://github.com/Farama-Foundation/Minigrid
     """
     return imresize(x / 255, size=size, channel_first=False)
 
 
-def actor(server: Connection, client: Connection, env: dm_env.Environment):
+def actor(server: Connection, client: Connection, env: Environment):
+    """Actor definition for Actor-Learner architectures.
+
+    Args:
+        server (Connection): server connection
+        client (Connection): client connection
+        env (Environment): environment to interact with
+    """
+
     def _step(env, a: int):
         timestep = env.step(a)
         if timestep.last():
@@ -76,7 +100,7 @@ def _step_async(env, a: int, buffer: mp.Queue):
             elif cmd == "reset":
                 client.send(env.reset())
             elif cmd == "render":
-                client.send(env.render(data))
+                client.send(env.render())
             elif cmd == "close":
                 client.send(env.close())
                 break
@@ -92,7 +116,6 @@ class MultiprocessEnv(dm_env.Environment):
     """
     This class is inspired by openai's SubprocEnv.
     https://github.com/openai/baselines/blob/master/baselines/common/vec_env/subproc_vec_env.py
-
     An environment that allows concurrent interactions to improve experience collection throughput, as used in:
     https://arxiv.org/abs/1602.01783, https://arxiv.org/abs/1802.01561 and https://arxiv.org/abs/1707.06347
     The class runs multiple subproceses and communicates with them via pipes.
@@ -100,7 +123,7 @@ class MultiprocessEnv(dm_env.Environment):
 
     def __init__(
         self,
-        env: dm_env.Environment,
+        env: Environment,
         n_actors: int,
         context: str = "spawn",
         seed: int = 0,
@@ -112,11 +135,9 @@ def __init__(
         )
         #  public:
         self.n_actors: int = n_actors
-        self.clients: Sequence[Connection] = None
-        self.servers: Sequence[Connection] = None
-        self.envs: Sequence[dm_env.Environment] = [
-            deepcopy(env) for _ in range(n_actors)
-        ]
+        self.clients: Sequence[Connection] = []
+        self.servers: Sequence[Connection] = []
+        self.envs: Sequence[Environment] = [deepcopy(env) for _ in range(n_actors)]
         self.processes = []
 
         #  setup parallel workers
@@ -125,9 +146,9 @@ def __init__(
         pipes = zip(*[ctx.Pipe() for _ in range(self.n_actors)])
         self.clients, self.servers = pipes
         for server, client, env in zip(self.servers, self.clients, self.envs):
-            env.gym_env.seed(int(next(rng)[0]))
+            env.seed(int(next(rng)[0]))
             self.processes.append(
-                ctx.Process(
+                ctx.Process(  # type: ignore
                     target=actor,
                     args=(server, client, env),
                     daemon=True,
@@ -158,12 +179,12 @@ def action_spec(self):
     def discount_spec(self):
         return self.envs[0].discount_spec()
 
-    def reset(self) -> None:
+    def reset(self) -> List[Any]:
         for server in self.servers:
             server.send(("reset", None))
         return self._receive()
 
-    def step(self, actions: Sequence[int]) -> dm_env.TimeStep:
+    def step(self, actions: Sequence[int]) -> List[dm_env.TimeStep]:
         self._check_actions(actions)
         for a, server in zip(actions, self.servers):
             server.send(("step", a))
@@ -175,7 +196,7 @@ def step_async(self, actions: Sequence[int], queue: mp.Queue) -> None:
             server.send(("step_async", (a, queue)))
         return
 
-    def close(self) -> None:
+    def close(self) -> List[Any]:
         for server in self.servers:
             server.send(("close", None))
         return self._receive()
@@ -188,7 +209,7 @@ def render(self, mode: str = "rgb_array"):
     def is_waiting(self):
         return any(server.poll() for server in self.servers)
 
-    def _receive(self):
+    def _receive(self) -> List[Any]:
         return [server.recv() for server in self.servers]
 
     def _check_actions(self, actions: Sequence[int]):
diff --git a/helx/experiment.py b/helx/experiment.py
new file mode 100644
index 0000000..7d14224
--- /dev/null
+++ b/helx/experiment.py
@@ -0,0 +1,88 @@
+import logging
+
+import dm_env
+import jax.numpy as jnp
+import numpy as np
+
+import wandb
+
+from .agents.agent import Agent
+
+
+def run(
+    agent: Agent,
+    env: dm_env.Environment,
+    num_episodes: int,
+    num_eval_episodes: int = 5,
+    eval_frequency: int = 1000,
+    print_frequency=10,
+    project="",
+    experiment_name: str = "",
+    debug: bool = False
+):
+    # init logger
+    nameof = lambda x: type(x).__name__
+    env_name, agent_name = list(map(nameof, (env, agent)))
+    mode = ("online", "disabled")[int(debug)]
+    wandb.init(
+        project=project,
+        group="{}-{}".format(env_name, agent_name),
+        tags=(env_name, agent_name, experiment_name),
+        config=agent.hparams._asdict(),
+        mode=mode
+    )
+
+    logging.info(
+        "Starting {} agent {} on environment {}.\nThe scheduled number of episode is {}".format(
+            "evaluating" if eval else "training", agent_name, env_name, num_episodes
+        )
+    )
+    logging.info(
+        "The hyperparameters for the current experiment are {}".format(
+            agent.hparams._asdict()
+        )
+    )
+    logging.info("Start logging experiment on wandb project {}".format(experiment_name))
+
+    # train
+    for i in range(num_episodes):
+
+        #  experience a new episode
+        episode = agent.unroll(env)
+        returns = jnp.sum(episode.r).item()
+
+        #  update policy
+        loss = agent.update(episode)
+
+        # log episode
+        if i % print_frequency == 0:
+            logging.info(
+                "Iteration: {}/{} - Return: {} - Loss: {}".format(
+                    agent.iteration, num_episodes - 1, returns, loss
+                )
+            )
+
+        if i % eval_frequency:
+            continue
+
+        # evaluate
+        # nameof = lambda x: type(x).__name__
+        # agent.save("{}-{}.pickle".format(nameof(agent), nameof(env)))
+
+        expected_returns = 0.0
+        for j in range(num_eval_episodes):
+            logging.info("Evaluating episode {} at iteration {}".format(j, i))
+            #  experience a new episode
+            episode = agent.unroll(env, eval=True)
+            returns = jnp.sum(episode.r).item()
+            video = np.array(episode.s).transpose((0, 3, 1, 2))
+            wandb.log({f"val/policy-{j}": wandb.Video(video, format="mp4")})
+
+            # log episode
+            logging.info(
+                "Episode: {}/{} - Return: {}".format(j, num_eval_episodes - 1, returns)
+            )
+            expected_returns += returns
+
+        expected_returns /= num_eval_episodes
+        wandb.log({"val/Return": expected_returns})
diff --git a/helx/image.py b/helx/image.py
index 6d652f1..ff5a383 100644
--- a/helx/image.py
+++ b/helx/image.py
@@ -1,11 +1,11 @@
 import jax
 import jax.numpy as jnp
 
-from .typing import Size
+from chex import Shape
 
 
 def imresize(
-    img, size: Size, method: str = "bilinear", channel_first: bool = True, **kwargs
+    img, size: Shape, method: str = "bilinear", channel_first: bool = True, **kwargs
 ):
     assert (
         len(size) == 2
@@ -26,4 +26,4 @@ def greyscale(img):
         (1,) * (img.ndim - 1) + (3,)
     )
     img = jnp.sum(img * luminance_mask, axis=-1).squeeze()
-    return img
+    return img
\ No newline at end of file
diff --git a/helx/jax.py b/helx/jax.py
deleted file mode 100644
index 164078f..0000000
--- a/helx/jax.py
+++ /dev/null
@@ -1,93 +0,0 @@
-from functools import wraps
-
-import jax
-import jax.numpy as jnp
-
-
-def batch(fun, in_axes=0, out_axes=0, axis_name=None, **jit_kwargs):
-    """
-    A utility wrapper around `jax.vmap` + `jax.jit`.
-    """
-    return jax.jit(
-        jax.vmap(fun, in_axes=in_axes, out_axes=out_axes, axis_name=axis_name),
-        **jit_kwargs
-    )
-
-
-def inject(fun, **kwargs):
-    """
-    It is a common pattern to define closures of classes with JAX.
-    `inject` adds the closure to the outer class as a static method, and jits it,
-    The function after `inject` is pure, and can be used as any other function in JAX.
-    """
-    cls = fun.__globals__[fun.__qualname__.split(".")[0]]
-    f_jit = jax.jit(fun, **kwargs)
-    setattr(cls, fun.__name__, staticmethod(f_jit))
-    return inject
-
-
-def pure(fun, **kwargs):
-    """
-    Puryfies a class function from the class instance parameter
-    in order to be used as any jax function, and finally `jit`s it.
-    """
-    f_jit = jax.jit(fun, **kwargs)
-
-    #  the actual computation tu run
-    @wraps(fun)
-    def wrapped(*a, **k):
-        return f_jit(*a, **k)
-
-    #  staticmethod does not pass the `self` argument when wrapper is called
-    return staticmethod(wrapped)
-
-
-def ppure(fun, **kwargs):
-    """
-    Parallelised version of `pure`.
-    Puryfies a class function from the class instance parameter
-    in order to be used as any jax function, and finally `pmap`s it.
-    """
-    f_pmap = jax.pmap(fun, **kwargs)
-
-    #  the actual computation tu run
-    @wraps(fun)
-    def wrapped(*a, **k):
-        return f_pmap(*a, **k)
-
-    #  staticmethod does not pass the `self` argument when wrapper is called
-    return staticmethod(wrapped)
-
-
-def tree_vmap(f, lst):
-    stacked = jax.tree_map(lambda args: jnp.stack(args), lst)
-    out_stacked = jax.vmap(f)(stacked)
-    _, outer_treedef = jax.tree_flatten([None] * len(lst))
-    _, inner_treedef = jax.tree_flatten(out_stacked)
-    out_unstacked_transposed = jax.tree_map(list, out_stacked)
-    out_unstacked = jax.tree_transpose(
-        outer_treedef, inner_treedef, out_unstacked_transposed
-    )
-    return out_unstacked
-
-
-def fori_scan(start, stop, body_fun, init, reversed=False):
-    """Uses `jax.lax.scan` to return all the intermediate
-    computations of the `jax.lax.fori_loop` function.
-    This function is also backward differentiable, unlike `fori_loop`.
-    """
-
-    def f(carry, i):
-        y = body_fun(i, carry)
-        return y, y
-
-    indices = jnp.arange(start, stop)
-    indices = jax.lax.cond(
-        reversed, lambda _: jnp.flip(indices), lambda _: indices, operand=None
-    )
-    return jax.lax.scan(f, init, xs=indices)[1]
-
-
-def device_array(numpy_array, device, *args, **kwargs):
-    """Creates a DeviceArray directly on a specified device"""
-    return jnp.array(jax.device_put(numpy_array, device), *args, **kwargs)
diff --git a/helx/logging.py b/helx/logging.py
deleted file mode 100644
index 4abd257..0000000
--- a/helx/logging.py
+++ /dev/null
@@ -1,104 +0,0 @@
-import logging
-from abc import ABC, abstractmethod
-from logging import handlers
-from datetime import datetime
-from typing import Sequence, Tuple
-
-import wandb
-
-
-class Logger(ABC):
-    @abstractmethod
-    def log(self, *args, **kwargs):
-        return
-
-
-class NullLogger(Logger):
-    def log(self):
-        return None
-
-
-class TerminalLogger(Logger):
-    def __init__(self, level: int = 20):
-        """
-        Uses the python default logging module to print
-        information on the screen.
-
-        Args:
-
-        level (int): minimum logging level, between:
-            CRITICAL = 50
-            ERROR = 40
-            WARNING = 30
-            INFO = 20
-            DEBUG = 10
-            NOTSET = 0
-        """
-        logger = logging.getLogger("helx-terminal")
-        logger.setLevel(level)
-        self._logger = logger
-
-    def log(self, *args, **kwargs):
-        return self._logger.info(*args, **kwargs)
-
-
-class FileLogger(Logger):
-    def __init__(self, filepath: str, level: int = 20):
-        """
-        Uses the python default logging module to print
-        information on the screen.
-
-        Args:
-
-        filepath (str): The path of the file to save the logging to
-        level (int): minimum logging level, between:
-            CRITICAL = 50
-            ERROR = 40
-            WARNING = 30
-            INFO = 20
-            DEBUG = 10
-            NOTSET = 0
-        """
-        logger = logging.getLogger("helx-file")
-        logger.setLevel(level)
-        logger.addHandler(handlers.RotatingFileHandler(filepath, maxBytes=2e8))
-        self._logger = logger
-
-    @property
-    def filepath(self):
-        if self._logger.hasHandlers():
-            return None
-        return self._logger.handlers[0].baseFilename
-
-    def log(self, *args, **kwargs):
-        return self._logger.info(*args, **kwargs)
-
-
-class EmailLogger(Logger):
-    def __init__(
-        self,
-        mailhost: Tuple[str, int],
-        sender: str,
-        recipients: Sequence[str],
-        credentials: Tuple[str, str],
-        level=20,
-    ):
-        raise NotImplementedError
-
-
-class HttpLogger(Logger):
-    def __init__(self, host: str, method: str = "POST", level: int = 20):
-        raise NotImplementedError
-
-
-class WandBLogger(Logger):
-    def __init__(self, *args, **kwargs):
-        wandb.init(*args, **kwargs)
-
-    def log(self, *args, **kwargs):
-        return wandb.log(*args, **kwargs)
-
-
-class TensorboardLogger(Logger):
-    def __init__(self):
-        raise NotImplementedError
diff --git a/helx/mdp.py b/helx/mdp.py
new file mode 100644
index 0000000..cc42212
--- /dev/null
+++ b/helx/mdp.py
@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+from functools import partial
+from typing import List, Tuple
+
+import dm_env
+import jax
+import jax.numpy as jnp
+from chex import Array
+from dm_env import TimeStep
+from jax.tree_util import register_pytree_node_class
+
+
+def tree_stack(pytree, axis=0):
+    return jax.tree_util.tree_map(lambda *x: jnp.stack(x, axis=axis), *pytree)
+
+
+@register_pytree_node_class
+class Episode:
+    """A collection of (s)
+    The class uses the following data structure:
+        >>> s₀, s₁, s₂, s₃, ..., sₜ,
+        >>> a₀, a₁, a₂, a₃, ..., aₜ,
+        >>> /  , r₁, r₂, r₃, ..., rₜ,
+        >>> /  , d₁, d₂, d₃, ..., dₜ,
+    ```
+    And corresponds to the following sarsa unroll:
+    (s₀, a₀, r₁, s₁, d₁, a₁)
+
+    Where, a_0 corresponds to the action taken with respect to
+    observation s_0, and r_1 is the reward received during the transition
+    from s_0 to s_1.
+    """
+
+    def __init__(self, s: List[Array], a: List[Array], r: List[Array], d: List[Array]):
+        self._s: List[Array] = s
+        self._a: List[Array] = a
+        self._r: List[Array] = r
+        self._d: List[Array] = d
+
+    @property
+    def s(self) -> Array:
+        return jnp.stack(self._s)
+
+    @property
+    def a(self) -> Array:
+        return jnp.stack(self._a)
+
+    @property
+    def r(self) -> Array:
+        return jnp.stack(self._r)
+
+    @property
+    def d(self) -> Array:
+        return jnp.stack(self._d)
+
+    def __getitem__(self, idx):
+        return jax.tree_util.tree_map(lambda x: x[idx], self.sarsa())
+
+    def __len__(self):
+        # the number of transitions is 1 minus the number of stored states
+        # since states contain also the state for the next timestep
+        return len(self.s) - 1
+
+    def tree_flatten(self):
+        return (self.s, self.a, self.r, self.d), ()
+
+    @classmethod
+    def tree_unflatten(
+        cls,
+        aux: Tuple[int],
+        children: Tuple[List[Array], List[Array], List[Array], List[Array]],
+    ) -> Episode:
+        """Decodes the PyTree into the `Episode` python object"""
+        return cls(*children)
+
+    @classmethod
+    def start(cls, timestep: TimeStep):
+        return cls([timestep.observation], [], [], [])
+
+    def add(
+        self,
+        timestep: dm_env.TimeStep,
+        action: Array,
+    ) -> None:
+        """Adds a new timestep to the episode.
+            Args:
+                timestep (dm_env.TimeStep): The timestep to add
+                action (Array): the action taken at `timestep.observation`
+            Returns:
+                None
+            """
+        # add new transition to the trajectory
+        self._s.append(jnp.asarray(timestep.observation, dtype=jnp.float32))
+        self._a.append(action)
+        self._r.append(jnp.asarray([timestep.reward], dtype=jnp.float32))
+        # `self.d` stores whether the timestep is *terminal*
+        # reaching env.max_steps does not cause termination, but truncation
+        # γ(truncation) != 0, while γ(termination) = 0
+        truncated = bool(timestep.discount)
+        self._d.append(jnp.array([timestep.last() * int(truncated)], dtype=jnp.int32))
+        return
+
+    def sars(self, axis=0):
+        """Computes a (s₀, a₀, r₁, s₁, d₁) unroll of the episode.
+            Args:
+                axis (int): The temporala axis to index into
+            Returns:
+                (Tuple[Array, Array, Array, Array, Array]) a 5D-tuple containing
+                each transition in the episode, where the first axis it the
+                temporal axis
+        """
+        assert len(self.s) - 1== len(self.r) == len(self.d) == len(self.a)
+        take = partial(jax.lax.slice_in_dim, axis=axis)
+        pairs = []
+        for t in range(0, len(self.s) - 1):
+            transition = (
+                take(self.s, t, t + 1),
+                take(self.a, t, t + 1),
+                take(self.r, t, t + 1),
+                take(self.s, t + 1, t + 2),
+                take(self.d, t, t + 1),
+            )
+            pairs.append(transition)
+        return pairs
+
+    def sarsa(self, axis=0):
+        """Returns a (s₀, a₀, r₁, s₁, d₁, a₁) unroll of the episode
+            Args:
+                axis (int): The temporala axis to index into
+            Returns:
+                (Tuple[Array, Array, Array, Array, Array, Array]) a 6D-tuple containing
+                each transition in the episode, where the first axis it the
+                temporal axis
+        """
+        assert len(self.s) == len(self.r) + 1 == len(self.d) + 1 == len(self.a)
+        take = partial(jax.lax.slice_in_dim, axis=axis)
+        pairs = []
+        for t in range(0, len(self.s) - 1):
+            transition = (
+                take(self.s, t, t + 1),
+                take(self.a, t, t + 1),
+                take(self.r, t, t + 1),
+                take(self.s, t + 1, t + 2),
+                take(self.d, t, t + 1),
+                take(self.a, t + 1, t + 2),
+            )
+            pairs.append(transition)
+        return pairs
diff --git a/helx/memory.py b/helx/memory.py
new file mode 100644
index 0000000..1621ede
--- /dev/null
+++ b/helx/memory.py
@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from collections import deque
+
+import jax
+
+from .mdp import Episode, tree_stack
+
+
+class ReplayBuffer:
+    """A replay buffer used for Experience Replay (ER):
+    Li, L., 1993, https://apps.dtic.mil/sti/pdfs/ADA261434.pdf.
+    This type of buffer is usually used
+    with off-policy methods, such as DQN or ACER.
+    Note that, to save memory, it stores only the two extremes
+    of the trajectory and accumulates the discunted rewards
+    at each time step, to calculate the value target.
+    However, this does not allow for off-policy corrections with nstep methods
+    at consumption time. To perform off-policy corrections, please store
+    the action probabilities foreach time step in the buffer.
+    """
+
+    def __init__(self, capacity: int, seed: int = 0):
+        self.capacity = capacity
+        self.key = jax.random.PRNGKey(seed)
+        self.elements = [None] * capacity
+        self.length = 0
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, idx):
+        return self.elements[idx]
+
+    def full(self) -> bool:
+        return len(self) == self.capacity
+
+    def add(
+        self,
+        transition: object,
+    ) -> None:
+        idx = self.length % self.capacity
+        self.elements[idx] = transition  # type: ignore
+        self.length += 1
+
+    def add_range(self, elements) -> None:
+        for element in elements:
+            self.add(element)
+
+    def sample(self, n: int) -> Episode:
+        self.key, k = jax.random.split(self.key)
+        n = min(n, self.length)
+        indices = jax.random.randint(k, (n,), 0, len(self))
+        return tree_stack([self.elements[i] for i in indices])  # type: ignore
diff --git a/helx/nn/__init__.py b/helx/nn/__init__.py
deleted file mode 100644
index 1fc6b86..0000000
--- a/helx/nn/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-from . import module
-from . import rnn
\ No newline at end of file
diff --git a/helx/nn/module.py b/helx/nn/module.py
deleted file mode 100644
index 720bb6d..0000000
--- a/helx/nn/module.py
+++ /dev/null
@@ -1,43 +0,0 @@
-import functools
-from typing import Callable, NamedTuple
-
-import jax.numpy as jnp
-
-from helx.typing import Apply, Init, factory
-from helx.distributed import distribute_tree
-
-
-InitState = Callable[[], jnp.ndarray]
-
-
-class Module(NamedTuple):
-    init: Init
-    apply: Apply
-
-
-def module(fun):
-    """
-    Decorator method to define `Module`s.
-    The decorator wraps the desired function and returns a Module object.
-    """
-    return factory(fun, Module)
-
-
-def pmodule(fun):
-    """
-    Same as `helx.methods.module` with an additional axis for the parameters
-    to be used with `jax.pmap`.
-    """
-
-    @functools.wraps(fun)
-    def wrapper(*args, **kwargs):
-        module = Module(*fun(*args, **kwargs))
-
-        def init(input_shape, rng):
-            output_shape, params = module.init(input_shape, rng)
-            params = distribute_tree(params)
-            return output_shape, params
-
-        return module._replace(init=init)
-
-    return wrapper
diff --git a/helx/optimise/__init__.py b/helx/optimise/__init__.py
deleted file mode 100644
index e3e7906..0000000
--- a/helx/optimise/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-from . import optimisers
-from . import schedulers
\ No newline at end of file
diff --git a/helx/optimise/optimisers.py b/helx/optimise/optimisers.py
deleted file mode 100644
index d9d13e0..0000000
--- a/helx/optimise/optimisers.py
+++ /dev/null
@@ -1,46 +0,0 @@
-import pickle
-from typing import Callable, NamedTuple
-
-import jax.numpy as jnp
-from jax.experimental import optimizers
-from jax.experimental.optimizers import OptimizerState
-
-from ..typing import HParams, Params
-from ..distributed import redistribute_tree
-
-
-class Optimiser(NamedTuple):
-    init: Callable[[Params], OptimizerState]
-    update: Callable[[int, jnp.ndarray, OptimizerState], OptimizerState]
-    params: Callable[[OptimizerState], Params]
-
-
-
-class TrainState(NamedTuple):
-    iteration: int
-    opt_state: OptimizerState
-    rng: jnp.ndarray = None
-    hparams: HParams = None
-
-    def serialise(self):
-        state = self._replace(
-            opt_state=optimizers.unpack_optimizer_state(self.opt_state)
-        )
-        return pickle.dumps(state)
-
-    @staticmethod
-    def deserialise(obj):
-        state = pickle.loads(obj)
-        opt_state = optimizers.pack_optimizer_state(state.opt_state)
-        state = state._replace(opt_state=redistribute_tree(opt_state))
-        return state
-
-    def save(self, filepath):
-        with open(filepath, "wb") as f:
-            state = self.serialise()
-            f.write(state)
-
-    @staticmethod
-    def load(filepath):
-        with open(filepath, "rb") as f:
-            return TrainState.deserialise(f.read())
diff --git a/helx/optimise/schedulers.py b/helx/optimise/schedulers.py
deleted file mode 100644
index e165f73..0000000
--- a/helx/optimise/schedulers.py
+++ /dev/null
@@ -1,18 +0,0 @@
-from typing import Deque, Callable, NamedTuple
-
-from ..typing import factory
-
-class SchedulerState(NamedTuple):
-    value: float
-    metric_history: Deque[float]
-
-
-class Scheduler(NamedTuple):
-    step: Callable[[float, SchedulerState], SchedulerState]
-    state: SchedulerState
-
-
-def scheduler(fun):
-    return factory(fun, Scheduler)
-
-
diff --git a/helx/random.py b/helx/random.py
index 0173dc0..4a2b3d8 100644
--- a/helx/random.py
+++ b/helx/random.py
@@ -1,8 +1,9 @@
+"""Random number generation utilities."""
 from typing import Generator
 import jax
 import jax.numpy as jnp
 
-from .typing import Key
+from jax.random import KeyArray as Key
 
 
 KeySequence = Generator[Key, Key, None]
@@ -34,5 +35,5 @@ def PRNGSequence(seed: int) -> KeySequence:
         yield rng
 
 
-def new_key(seq: PRNGSequence) -> Key:
+def new_key(seq: KeySequence) -> Key:
     return next(seq)
diff --git a/helx/rl/__init__.py b/helx/rl/__init__.py
deleted file mode 100644
index 9cabdb5..0000000
--- a/helx/rl/__init__.py
+++ /dev/null
@@ -1,7 +0,0 @@
-from . import policy
-from . import agent
-from . import memory
-from . import environment
-from . import pg
-from . import td
-from . import baselines
diff --git a/helx/rl/agent.py b/helx/rl/agent.py
deleted file mode 100644
index 43d048c..0000000
--- a/helx/rl/agent.py
+++ /dev/null
@@ -1,107 +0,0 @@
-import abc
-import logging
-
-import dm_env
-import jax
-from helx.nn.module import Module
-from helx.optimise.optimisers import Optimiser
-from helx.typing import Action, HParams, Loss
-
-
-class IAgent:
-    network: Module
-    optimiser: Optimiser
-    hparams: HParams
-
-    def __init__(
-        self,
-        network: Module,
-        optimiser: Optimiser,
-        hparams: HParams,
-        logging: bool = False,
-    ):
-        IAgent.network = network
-        IAgent.optimiser = optimiser
-        IAgent.hparams = hparams
-        self.logging = logging
-        self._iteration = 0
-
-    @abc.abstractmethod
-    def observe(
-        self, env: dm_env.Environment, timestep: dm_env.TimeStep, action: int
-    ) -> dm_env.TimeStep:
-        """The agent's observation function defines how  it interacts with the enviroment"""
-
-    @abc.abstractmethod
-    def policy(self, timestep: dm_env.TimeStep) -> int:
-        """The agent's policy function that maps an observation to an action"""
-
-    @abc.abstractmethod
-    def update(
-        self, timestep: dm_env.TimeStep, action: int, new_timestep: dm_env.TimeStep
-    ) -> float:
-        """The agent's policy function that maps an observation to an action"""
-
-    @abc.abstractmethod
-    def log(
-        self,
-        timestep: dm_env.TimeStep,
-        action: Action,
-        new_timestep: dm_env.TimeStep,
-        loss: Loss = None,
-        log_frequency: int = 1,
-    ):
-        """Empty by default, can be overriden to record learning statistics"""
-
-    def run(
-        self,
-        env: dm_env.Environment,
-        num_episodes: int,
-        eval: bool = False,
-    ) -> Loss:
-        logging.info(
-            "Starting {} agent {} on environment {}.\nThe scheduled number of episode is {}".format(
-                "evaluating" if eval else "training", self, env, num_episodes
-            )
-        )
-        logging.info(
-            "The hyperparameters for the current experiment are {}".format(
-                self.hparams._asdict()
-            )
-        )
-        for episode in range(num_episodes):
-            print(
-                "Episode {}/{}\t\t\t".format(episode, num_episodes - 1),
-                end="\r",
-            )
-            #  initialise environment
-            episode_reward = 0.0
-            timestep = env.reset()
-            while not timestep.last():
-                #  apply policy
-                action = self.policy(timestep)
-                #  observe new state
-                new_timestep = self.observe(env, timestep, action)
-                episode_reward += new_timestep.reward
-                print(
-                    "Episode reward {}\t\t".format(episode_reward),
-                    end="\r",
-                )
-                #  update policy
-                loss = None
-                if not eval:
-                    loss = self.update(timestep, action, new_timestep)
-                #  log update
-                if self.logging:
-                    self.log(timestep, action, new_timestep, loss)
-                # prepare next iteration
-                timestep = new_timestep
-        return loss
-
-    def rand_prob(self):
-        k, self.rng = jax.random.split(self.rng)
-        return jax.random.uniform(k)
-
-    def rand_action(self, low=0, high=10):
-        k, self.rng = jax.random.split(self.rng)
-        return jax.random.randint(k, (1,), 0, self.action_spec.num_values)
diff --git a/helx/rl/baselines/__init__.py b/helx/rl/baselines/__init__.py
deleted file mode 100644
index 92f35c1..0000000
--- a/helx/rl/baselines/__init__.py
+++ /dev/null
@@ -1,5 +0,0 @@
-from . import (
-    a2c,
-    dqn,
-    ppo,
-)
diff --git a/helx/rl/baselines/a2c.py b/helx/rl/baselines/a2c.py
deleted file mode 100644
index 62a4bf4..0000000
--- a/helx/rl/baselines/a2c.py
+++ /dev/null
@@ -1,171 +0,0 @@
-from typing import Callable, NamedTuple, Tuple
-
-import dm_env
-import jax
-import jax.numpy as jnp
-import wandb
-from dm_env import TimeStep, specs
-from helx.jax import pure
-from helx.nn.module import Module, module
-from helx.optimise.optimisers import Optimiser
-from helx.typing import Action, Loss, Params, Size
-from jax.experimental import stax
-from jax.experimental.optimizers import OptimizerState, rmsprop_momentum
-
-from .. import pg, td
-from ..agent import IAgent
-from ..memory import OnlineBuffer, Trajectory
-
-
-class HParams(NamedTuple):
-    seed: int = 0
-    discount: float = 1.0
-    trace_decay: float = 1.0
-    n_steps: int = 1
-    beta: float = 0.001
-    learning_rate: float = 0.001
-    gradient_momentum: float = 0.95
-    squared_gradient_momentum: float = 0.95
-    min_squared_gradient: float = 0.01
-
-
-@module
-def Cnn(n_actions: int, hidden_size: int = 512) -> Module:
-    return stax.serial(
-        stax.Conv(32, (8, 8), (4, 4), "VALID"),
-        stax.Relu,
-        stax.Conv(64, (4, 4), (2, 2), "VALID"),
-        stax.Relu,
-        stax.Conv(64, (3, 3), (1, 1), "VALID"),
-        stax.Relu,
-        stax.Flatten,
-        stax.Dense(hidden_size),
-        stax.Relu,
-        stax.FanOut(2),
-        stax.parallel(
-            stax.serial(
-                stax.Dense(n_actions),
-            ),  #  actor
-            stax.serial(
-                stax.Dense(1),
-            ),  # critic
-        ),
-    )
-
-
-class A2C(IAgent):
-    """Synchronous on-policy, online actor-critic algorithm
-    with n-step advantage baseline.
-    See:
-    Mnih V., 2016, https://arxiv.org/pdf/1602.01783.pdf and
-    Sutton R., Barto. G., 2018, http://incompleteideas.net/book/RLbook2020.pdf
-
-    This algorithm implements Mnih, V. 2016, that is,
-    we ran the policy for T steps, where T = batch_size << len_episode
-    and use the collected samples to update.
-    """
-
-    def __init__(
-        self,
-        obs_spec: specs.Array,
-        action_spec: specs.DiscreteArray,
-        hparams: HParams,
-        preprocess: Callable = lambda x: x,
-        logging: bool = False,
-    ):
-        # public:
-        self.obs_spec = obs_spec
-        self.action_spec = action_spec
-        self.rng = jax.random.PRNGKey(hparams.seed)
-        self.memory = OnlineBuffer(1, hparams.seed)
-        self.preprocess = preprocess
-        network = Cnn(action_spec.num_values)
-        optimiser = Optimiser(
-            *rmsprop_momentum(
-                step_size=hparams.learning_rate,
-                gamma=hparams.squared_gradient_momentum,
-                momentum=hparams.gradient_momentum,
-                eps=hparams.min_squared_gradient,
-            )
-        )
-        super().__init__(network, optimiser, hparams, logging)
-
-        # private:
-        _, params = self.network.init(self.rng, (-1, *obs_spec.shape))
-        self._opt_state = self.optimiser.init(params)
-
-    @pure
-    def loss(
-        params: Params,
-        trajectory: Trajectory,
-    ) -> Loss:
-        """We run the policy π for T timesteps,
-        to calculate the advantage"""
-        #  inference
-        logits, values = A2C.network.apply(params, trajectory.observations)
-        #   Critic loss (Bellman regression)
-        advantages = td.advantages(trajectory, values[:-1], values[1:])
-        critic_loss = jnp.mean(jnp.square(advantages))  #  bellman mse
-        #  Actor loss (Policy gradient loss)
-        actor_loss = pg.a2c_softmax_loss(logits, advantages)
-        entropy = pg.softmax_entropy(logits)
-        return critic_loss + actor_loss - A2C.hparams.beta * entropy
-
-    @pure
-    def sgd_step(
-        iteration: int,
-        opt_state: OptimizerState,
-        transition: Trajectory,
-    ) -> Tuple[Loss, OptimizerState]:
-        params = A2C.optimiser.params(opt_state)
-        backward = jax.value_and_grad(A2C.loss, argnums=2)
-        error, grads = backward(A2C.network, params, transition)
-        return error, A2C.optimiser.update(iteration, grads, opt_state)
-
-    def observe(
-        self, env: dm_env.Environment, timestep: dm_env.TimeStep, action: int
-    ) -> dm_env.TimeStep:
-        #  iterate over the number of steps
-        for t in range(self.hparams.n_steps):
-            #  get new MDP state
-            new_timestep = env.step(action)
-            #  store transition into the replay buffer
-            self.memory.add(timestep, action, new_timestep, preprocess=self.preprocess)
-            timestep = new_timestep
-        return timestep
-
-    def policy(self, timestep: dm_env.TimeStep) -> int:
-        """Selects an action using a softmax policy"""
-        params = self.optimiser.params(self._opt_state)
-        logits, _ = self.network.apply(params, timestep.observation)
-        action = jax.random.categorical(self.rng, logits).squeeze()  # on-policy action
-        return int(action)
-
-    def update(
-        self, timestep: dm_env.TimeStep, action: int, new_timestep: dm_env.TimeStep
-    ) -> None:
-
-        loss = None
-        if self.memory.full():
-            transition = self.memory.sample(self.hparams.batch_size)
-            loss, self.opt_state = self.sgd_step(
-                self.network, self.optimiser, transition
-            )
-        return loss
-
-    def log(
-        self,
-        timestep: TimeStep,
-        action: Action,
-        new_timestep: TimeStep,
-        loss: Loss = None,
-        log_frequency: int = 1,
-    ):
-        wandb.log({"Iteration": float(self._iteration)})
-        wandb.log({"Action": int(action)})
-        wandb.log({"Reward": float(new_timestep.reward)})
-        if self._iteration % 100 == 0:
-            wandb.log({"Observation": wandb.Image(timestep.observation)})
-        if loss is not None:
-            wandb.log({"Loss": float(loss)})
-        return
diff --git a/helx/rl/baselines/dqn.py b/helx/rl/baselines/dqn.py
deleted file mode 100644
index f086dbc..0000000
--- a/helx/rl/baselines/dqn.py
+++ /dev/null
@@ -1,202 +0,0 @@
-from typing import Callable, NamedTuple, Tuple
-
-import dm_env
-import jax
-import jax.numpy as jnp
-import wandb
-from dm_env import Environment, TimeStep
-from helx.jax import pure
-from helx.nn.module import Module, module
-from helx.optimise.optimisers import Optimiser
-from helx.random import PRNGSequence
-from helx.typing import Action, Loss, Observation, Params, Shape
-from jax.experimental.optimizers import OptimizerState, rmsprop_momentum, sgd
-from jax.experimental.stax import Conv, Dense, Flatten, Relu, serial
-
-from .. import td
-from ..agent import IAgent
-from ..memory import ReplayBuffer, Trajectory
-
-
-class HParams(NamedTuple):
-    batch_size: Shape = 32
-    replay_memory_size: int = 1000
-    target_network_update_frequency: int = 10000
-    discount: float = 0.99
-    learning_rate: float = 0.00025
-    gradient_momentum: float = 0.95
-    squared_gradient_momentum: float = 0.95
-    min_squared_gradient: float = 0.01
-    initial_exploration: float = 1.0
-    final_exploration: float = 0.01
-    final_exploration_frame: int = 1000000
-    replay_start: int = 1000
-    hidden_size: int = 512
-    n_steps: int = 1
-
-
-@module
-def Cnn(n_actions: int, hidden_size) -> Module:
-    return serial(
-        Conv(32, (8, 8), (4, 4), "VALID"),
-        Relu,
-        Conv(64, (4, 4), (2, 2), "VALID"),
-        Relu,
-        Conv(64, (3, 3), (1, 1), "VALID"),
-        Relu,
-        Flatten,
-        Dense(hidden_size),
-        Relu,
-        Dense(n_actions),
-    )
-
-
-class Dqn(IAgent):
-    def __init__(
-        self,
-        obs_shape: Shape,
-        n_actions: int,
-        hparams: HParams,
-        seed: int = 0,
-        preprocess: Callable[[Observation], Observation] = lambda x: x,
-        logging: bool = False,
-    ):
-        # public:
-        self.n_actions = n_actions
-        self.obs_shape = obs_shape
-        self.memory = ReplayBuffer(hparams.replay_memory_size)
-        self.rng = PRNGSequence(seed)
-        self.preprocess = preprocess
-        network = Cnn(n_actions, hidden_size=hparams.hidden_size)
-        optimiser = Optimiser(
-            *rmsprop_momentum(
-                step_size=hparams.learning_rate,
-                gamma=hparams.squared_gradient_momentum,
-                momentum=hparams.gradient_momentum,
-                eps=hparams.min_squared_gradient,
-            )
-        )
-        super().__init__(network, optimiser, hparams, logging)
-
-        # private:
-        k = next(self.rng)
-        _, params_target = self.network.init(k, (-1, *obs_shape))
-        _, params_online = self.network.init(k, (-1, *obs_shape))
-        self._opt_state = self.optimiser.init(params_online)
-        self._params_target = params_target
-
-    @pure
-    def loss(
-        params_target: Params,
-        params_online: Params,
-        transition: Trajectory,
-    ) -> Loss:
-        q_online = Dqn.network.apply(params_online, transition.observations[:, 0])
-        q_target = Dqn.network.apply(params_target, transition.observations[:, -1])
-        q_target = jnp.max(q_target, axis=-1).reshape(-1, 1)
-        returns = td.nstep_return(transition, q_target)
-        return jnp.sqrt(
-            jnp.mean(jnp.square(returns - q_online[:, transition.actions[0]]))
-        )
-
-    @pure
-    def sgd_step(
-        iteration: int,
-        params_target: Params,
-        opt_state: OptimizerState,
-        transition: Trajectory,
-    ) -> Tuple[Loss, OptimizerState]:
-        params_online = Dqn.optimiser.params(opt_state)
-        backward = jax.value_and_grad(Dqn.loss, argnums=1)
-        error, gradients = backward(params_target, params_online, transition)
-        return error, Dqn.optimiser.update(iteration, gradients, opt_state)
-
-    def epsilon(self):
-        x0, y0 = (self.hparams.replay_start, self.hparams.initial_exploration)
-        x1 = self.hparams.final_exploration_frame
-        y1 = self.hparams.final_exploration
-        x = self._iteration
-        y = ((y1 - y0) * (x - x0) / (x1 - x0)) + y0
-        return jnp.clip(
-            y, self.hparams.final_exploration, self.hparams.initial_exploration
-        )
-
-    def observe(self, env: Environment, timestep: TimeStep, action: int) -> Trajectory:
-        #  iterate over the number of steps
-        for t in range(self.hparams.n_steps):
-            #  get new MDP state
-            new_timestep = env.step(action)
-            #  store transition into the replay buffer
-            self.memory.add(timestep, action, new_timestep, preprocess=self.preprocess)
-            timestep = new_timestep
-        return timestep
-
-    def policy(
-        self,
-        timestep: dm_env.TimeStep,
-    ) -> int:
-        """Selects an action using an e-greedy policy"""
-        k = next(self.rng)
-        #  use random policy with epsilon probability
-        if jax.random.uniform(k, (1,)) < self.epsilon():
-            return jax.random.randint(k, (1,), 0, self.n_actions)
-
-        #  use greedy policy otherwise
-        params = self.optimiser.params(self._opt_state)
-        obs = self.preprocess(timestep.observation)[None, ...]  # batching
-        q_values = self.network.apply(params, obs)
-        action = jnp.argmax(q_values)
-        return int(action)
-
-    def update(
-        self,
-        timestep: dm_env.TimeStep,
-        action: int,
-        new_timestep: dm_env.TimeStep,
-    ) -> float:
-        """Dqn uses a replay buffer. The three inputs
-        `timestep`, `action` and `new_timesdtep` are ignored."""
-        # increment iteration
-        self._iteration += 1
-
-        # if replay buffer is smaller than the minimum size, there is nothing else to do
-        if len(self.memory) < self.hparams.replay_start:
-            return None
-
-        # update the online parameters
-        transition = self.memory.sample(self.hparams.batch_size)
-        loss, opt_state = self.sgd_step(
-            self._iteration,
-            self._params_target,
-            self._opt_state,
-            transition,
-        )
-        self._opt_state = opt_state
-
-        # update the target network parameters every `target_network_update_frequency` steps
-        if self._iteration % self.hparams.target_network_update_frequency == 0:
-            params = self.optimiser.params(opt_state)
-            self._params_target = params
-        return loss
-
-    def log(
-        self,
-        timestep: TimeStep,
-        action: Action,
-        new_timestep: TimeStep,
-        loss: Loss = None,
-        log_frequency: int = 1,
-    ):
-        if self._iteration % log_frequency:
-            return
-        wandb.log({"Iteration": float(self._iteration)})
-        wandb.log({"Action": int(action)})
-        wandb.log({"Reward": float(new_timestep.reward)})
-        if self._iteration % 100 == 0:
-            wandb.log({"Observation": wandb.Image(timestep.observation)})
-        if loss is not None:
-            wandb.log({"Loss": float(loss)})
-
-        wandb.log({"Buffer size": len(self.memory)})
-        wandb.log({"Epsilon": float(self.epsilon())})
-        return
diff --git a/helx/rl/baselines/ppo.py b/helx/rl/baselines/ppo.py
deleted file mode 100644
index 1dc5606..0000000
--- a/helx/rl/baselines/ppo.py
+++ /dev/null
@@ -1,155 +0,0 @@
-from typing import NamedTuple, Tuple
-
-import dm_env
-import jax
-import jax.numpy as jnp
-import wandb
-from dm_env import specs
-from helx.jax import pure
-from helx.nn.module import Module, module
-from helx.optimise.optimisers import Optimiser
-from helx.typing import Loss, Params, Size
-from jax.experimental import stax
-from jax.experimental.optimizers import OptimizerState, rmsprop_momentum
-
-from .. import pg, td
-from ..agent import IAgent
-from ..memory import OnlineBuffer, Trajectory
-
-
-class HParams(NamedTuple):
-    seed: int = 0
-    discount: float = 1.0
-    trace_decay: float = 1.0
-    n_steps: int = 1
-    n_actors: int = 32
-    beta: float = 0.001
-    epsilon: float = 0.2
-    learning_rate: float = 0.001
-    gradient_momentum: float = 0.95
-    squared_gradient_momentum: float = 0.95
-    min_squared_gradient: float = 0.01
-
-
-@module
-def Cnn(n_actions: int, hidden_size: int = 512) -> Module:
-    return stax.serial(
-        stax.Conv(32, (8, 8), (4, 4), "VALID"),
-        stax.Relu,
-        stax.Conv(64, (4, 4), (2, 2), "VALID"),
-        stax.Relu,
-        stax.Conv(64, (3, 3), (1, 1), "VALID"),
-        stax.Relu,
-        stax.Flatten,
-        stax.Dense(hidden_size),
-        stax.Relu,
-        stax.FanOut(2),
-        stax.parallel(
-            stax.serial(
-                stax.Dense(n_actions),
-                stax.Softmax,
-            ),  #  actor
-            stax.serial(
-                stax.Dense(1),
-            ),  # critic
-        ),
-    )
-
-
-class Ppo(IAgent):
-    """Proximal Policy Optiomisation algorithm.
-    See:
-    Schulman, J., 2017, https://arxiv.org/abs/1707.06347.
-    The implementation is
-    """
-
-    def __init__(
-        self,
-        obs_spec: specs.Array,
-        action_spec: specs.DiscreteArray,
-        hparams: HParams,
-        preprocess: lambda x: x,
-        logging: bool = False,
-    ):
-        # public:
-        self.obs_spec = obs_spec
-        self.action_spec = action_spec
-        self.memory = OnlineBuffer(obs_spec, hparams.n_steps, hparams.n_actors)
-        self.rng = jax.random.PRNGKey(hparams.seed)
-        self.preprocess = preprocess
-        network = Cnn(action_spec.num_values)
-        optimiser = Optimiser(
-            *rmsprop_momentum(
-                step_size=hparams.learning_rate,
-                gamma=hparams.squared_gradient_momentum,
-                momentum=hparams.gradient_momentum,
-                eps=hparams.min_squared_gradient,
-            )
-        )
-        super().__init__(network, optimiser, hparams, logging)
-
-        # private:
-        _, params = self.network.init(self.rng, (-1, *obs_spec.shape))
-        self._opt_state = self.optimiser.init(params)
-        self.policy_old = params
-
-    @pure
-    def loss(
-        params: Params,
-        trajectory: Trajectory,
-        policy_old: Params,
-    ) -> Loss:
-        logits, values = Ppo.network.apply(params, trajectory.x_0)
-        logits_old, _ = Ppo.network.apply(policy_old, trajectory.observations)
-        #  Critic loss (Bellman regression)
-        gae = td.lambda_returns(trajectory, values)
-        critic_loss = jnp.mean(jnp.square(gae))  #  bellman mse
-        #  PPO clipped objective
-        actor_loss = pg.ppo_softmax_loss(logits, logits_old, gae, Ppo.hparams)
-        entropy = pg.softmax_entropy(logits)
-        return actor_loss + Ppo.hparams.c1 * critic_loss - Ppo.hparams.c2 * entropy
-
-    @pure
-    def sgd_step(
-        iteration: int,
-        opt_state: OptimizerState,
-        trajectory: Trajectory,
-        policy_old: Params,
-    ) -> Tuple[Loss, OptimizerState]:
-        params = Ppo.optimiser.params(opt_state)
-        backward = jax.value_and_grad(Ppo.loss, argnums=0)
-        error, grads = backward(Ppo.network, params, trajectory, policy_old)
-        return error, Ppo.optimiser.update(iteration, grads, opt_state)
-
-    def observe(
-        self, env: dm_env.Environment, timestep: dm_env.TimeStep, action: int
-    ) -> dm_env.TimeStep:
-        #  iterate over the number of steps
-        for t in range(self.hparams.n_steps):
-            #  get new MDP state
-            new_timestep = env.step(action)
-            #  store transition into the replay buffer
-            self.memory.add(timestep, action, new_timestep, preprocess=self.preprocess)
-            timestep = new_timestep
-        return timestep
-
-    def policy(self, timestep: dm_env.TimeStep) -> int:
-        """Selects an action using a softmax policy"""
-        params = self.optimiser.params(self._opt_state)
-        logits, _ = self.network.apply(params, timestep.observation)
-        action = jax.random.categorical(self.rng, logits).squeeze()  # on-policy action
-        return int(action)
-
-    def update(
-        self, timestep: dm_env.TimeStep, action: int, new_timestep: dm_env.TimeStep
-    ) -> None:
-        trajectory = self.buffer.sample()
-        loss, opt_state = self.sgd_step(self.network, self.optimiser, trajectory)
-        return loss, opt_state
-
-    def log(self, reward: float, loss: Loss):
-        wandb.log({"Iteration": float(self._iteration)})
-        wandb.log({"Reward": reward})
-        if loss is not None:
-            wandb.log({"Loss": float(loss)})
-        return
diff --git a/helx/rl/memory.py b/helx/rl/memory.py
deleted file mode 100644
index f35f891..0000000
--- a/helx/rl/memory.py
+++ /dev/null
@@ -1,291 +0,0 @@
-import abc
-import logging
-from collections import deque
-from typing import Callable, NamedTuple
-
-import dm_env
-import jax
-import jax.numpy as jnp
-from dm_env import specs
-from helx.jax import device_array
-from helx.random import PRNGSequence
-from helx.typing import Action, Batch, Discount, Key, Observation, Reward, TraceDecay
-from jaxlib.xla_extension import Device
-
-
-class Transition(NamedTuple):
-    """A (s, a, r, s', a', γ, λ) transition with discount and lambda factors"""
-
-    s: Observation  #  observation at t=0
-    a: Action  #  actions at t=0
-    r: Reward  #  rewards at t=1
-    s1: Observation  # observatin at t=1 (note multistep)
-    a1: Action  # action at t=1
-    g: Discount = 1.0  #  discount factor
-    l: TraceDecay = 1.0  # trace decay for lamba returns
-
-
-class Trajectory(NamedTuple):
-    """A set of batched transitions"""
-
-    observations: Batch[Observation]  #  [T + 1, *obs.shape]
-    actions: Batch[Action]  #  [T, 1] if off-policy, [T + 1, 1] otherwise
-    rewards: Batch[Reward]  #  [T, 1]
-    discounts: Batch[Discount] = None  #  [T, 1]
-    trace_decays: Batch[TraceDecay] = None  #  [T, 1]
-
-
-class IBuffer(abc.ABC):
-    @abc.abstractmethod
-    def add(
-        self,
-        timestep: dm_env.TimeStep,
-        action: int,
-        new_timestep: dm_env.TimeStep,
-        preprocess=lambda x: x,
-    ) -> None:
-        raise NotImplementedError
-
-    @abc.abstractmethod
-    def full(self) -> bool:
-        raise NotImplementedError
-
-    @abc.abstractmethod
-    def add(
-        self,
-        timestep: dm_env.TimeStep,
-        action: int,
-        new_timestep: dm_env.TimeStep,
-        preprocess=lambda x: x,
-    ) -> None:
-        raise NotImplementedError
-
-    @abc.abstractmethod
-    def sample(self, n: int, rng: Key = None) -> Trajectory:
-        raise NotImplementedError
-
-
-class ReplayBuffer(IBuffer):
-    """A replay buffer used for Experience Replay (ER):
-    Li, L., 1993, https://apps.dtic.mil/sti/pdfs/ADA261434.pdf.
-    This type of buffer is usually used
-    with off-policy methods, such as DQN or ACER.
-    Note that, to save memory, it stores only the two extremes
-    of the trajectory and accumulates the discunted rewards
-    at each time step, to calculate the value target.
-    However, this does not allow for off-policy corrections with nstep methods
-    at consumption time. To perform off-policy corrections, please store
-    the action probabilities foreach time step in the buffer.
-    """
-
-    def __init__(
-        self,
-        capacity: int,
-        n_steps: int = 1,
-        seed: int = 0,
-        device: Device = None,
-    ):
-        #  public:
-        self.capacity = capacity
-        self.n_steps = n_steps
-        self.seed = seed
-        self.device = device
-        self.trajectories = deque(maxlen=capacity)
-
-        #  private:
-        self._rng = jax.random.PRNGKey(seed)
-        self._reset()
-
-    def __len__(self):
-        return len(self.trajectories)
-
-    def __getitem__(self, idx):
-        return self.trajectories[idx]
-
-    def full(self) -> bool:
-        return len(self) == self.capacity
-
-    def collecting(self) -> bool:
-        return self._t < self.n_steps
-
-    def add(
-        self,
-        timestep: dm_env.TimeStep,
-        action: int,
-        new_timestep: dm_env.TimeStep,
-        preprocess=lambda x: x,
-    ) -> None:
-        #  start of a new episode
-        if not self.collecting():
-            self.trajectories.append(self._current)
-            self._reset()
-
-        # add new transition to the trajectory
-        store = lambda x: device_array(x, device=self.device)
-        self._current.observations[self._t] = preprocess(store(timestep.observation))
-        self._current.actions[self._t] = store(int(action))
-        self._current.rewards[self._t] = store(float(new_timestep.reward))
-        self._current.discounts[self._t] = store(float(new_timestep.discount))
-        self._t += 1
-
-        # ready to store, just add final observation
-        if not self.collecting():
-            self._current.observations[self._t] = preprocess(
-                jnp.array(timestep.observation, dtype=jnp.float32)
-            )
-        # if not enough samples, and we can't sample the env anymore, reset
-        elif new_timestep.last():
-            self._reset()
-        return
-
-    def sample(self, n: int, rng: Key = None, device: Device = None) -> Trajectory:
-        high = len(self) - n
-        if high <= 0:
-            logging.warning(
-                "The buffer contains less elements than requested: {} <= {}\n"
-                "Returning all the available elements".format(len(self), n)
-            )
-            indices = range(len(self))
-        elif rng is None:
-            rng, _ = jax.random.split(self._rng)
-        indices = jax.random.randint(rng, (n,), 0, high)
-
-        collate = lambda x: device_array(x, device=device)
-        obs = collate([self.trajectories[idx].observations for idx in indices])
-        actions = collate([self.trajectories[idx].actions for idx in indices])
-        rewards = collate([self.trajectories[idx].rewards for idx in indices])
-        discounts = collate([self.trajectories[idx].discounts for idx in indices])
-        # traces = collate([self.trajectories[idx].trace_decays for idx in indices])
-        return Trajectory(obs, actions, rewards, discounts)
-
-    def _reset(self):
-        self._t = 0
-        self._current = Trajectory(
-            observations=[None] * (self.n_steps + 1),
-            actions=[None] * self.n_steps,
-            rewards=[None] * self.n_steps,
-            discounts=[None] * self.n_steps,
-            trace_decays=[None] * self.n_steps,
-        )
-
-
-class OnlineBuffer(IBuffer):
-    """A replay buffer that stores a single n-step trajectory
-    of experience.
-    This type of buffer is most commonly used with online methods,
-    generally on-policy methods, such as A2C.
-    """
-
-    def __init__(
-        self,
-        observation_spec: specs.Array,
-        n_steps: int = 1,
-        batch_size: int = 1,
-    ):
-        #  public:
-        self.observation_spec = observation_spec
-        self.n_steps = n_steps
-        self.batch_size = batch_size
-
-        #  private:
-        self._reset()
-
-    def full(self) -> bool:
-        return self._t == self.n_steps - 1
-
-    def add(
-        self,
-        timestep: dm_env.TimeStep,
-        action: int,
-        new_timestep: dm_env.TimeStep,
-        trace_decay: TraceDecay = 1.0,
-        preprocess: Callable[[Observation], Observation] = lambda x: x,
-    ) -> None:
-        #  if buffer is full, prepare for new trajectory
-        if self.full():
-            self._reset()
-
-        # add new transition to the trajectory
-        self.trajectory.observations[self._t] = preprocess(
-            jnp.array(timestep.observation, dtype=jnp.float32)
-        )
-        self.trajectory.actions[self._t] = action
-        self.trajectory.rewards[self._t] = new_timestep.reward
-        self.trajectory.discounts[self._t] = new_timestep.discount
-        self.trajectory.trace_decays[self._t] = trace_decay
-        self._t += 1
-
-        #  if we have enough transitions, add last obs and return
-        if self.full():
-            self.trajectory.observations[self._t] = preprocess(
-                jnp.array(new_timestep.observation, dtype=jnp.float32)
-            )
-        #  if we do not have enough transitions, and can't sample more, retry
-        elif new_timestep.last():
-            self._reset()
-        return
-
-    def sample(self, n: int = 1, rng: Key = None) -> Trajectory:
-        return self.trajectory
-
-    def _reset(self):
-        self._t = 0
-        self.trajectory = Trajectory(
-            observations=jnp.empty(
-                self.n_steps + 1, self.batch_size, *self.observation_spec.shape
-            ),
-            actions=jnp.empty(self.n_steps, self.batch_size, 1),
-            rewards=jnp.empty(self.n_steps, self.batch_size, 1),
-            discounts=jnp.empty(self.n_steps, self.batch_size, 1),
-            trace_decays=jnp.empty(self.n_steps, self.batch_size, 1),
-        )
-        return
-
-
-class EpisodicMemory(IBuffer):
-    def __init__(self, seed: int = 0):
-        #  public:
-        self.states = []
-
-        #  private:
-        self._terminal = False
-        self._rng = PRNGSequence(seed)
-        self._reset()
-
-    def __len__(self):
-        return len(self.states)
-
-    def __getitem__(self, idx):
-        return self.states[idx]
-
-    def full(self) -> bool:
-        return self._terminal
-
-    def add(
-        self,
-        timestep: dm_env.TimeStep,
-        action: int,
-        new_timestep: dm_env.TimeStep,
-        preprocess: Callable = lambda x: x,
-    ) -> None:
-        #  if buffer is full, prepare for new trajectory
-        if self.full():
-            self._reset()
-
-        #  collect experience
-        self.states.append(preprocess(timestep.observation))
-        self._terminal = new_timestep.last()
-
-        #   if transition is terminal, append last state
-        if self.full():
-            self.states.append(preprocess(new_timestep.observation))
-        return
-
-    def sample(self, n: int, rng: Key = None) -> Trajectory:
-        key = next(self._rng)
-        indices = jax.random.randint(key, (n,), 0, len(self))
-        return [self.states[idx] for idx in indices]
-
-    def _reset(self):
-        self._terminal = False
-        self.states = []
diff --git a/helx/rl/pg.py b/helx/rl/pg.py
deleted file mode 100644
index 1390355..0000000
--- a/helx/rl/pg.py
+++ /dev/null
@@ -1,20 +0,0 @@
-import jax
-import jax.numpy as jnp
-
-
-def softmax_entropy(logits):
-    return -jnp.sum(jnp.mean(logits) * jnp.log(logits))
-
-
-def a2c_softmax_loss(logits, advantages):
-    return jax.nn.log_softmax(logits) * advantages
-
-
-def trpo_softmax_loss(logits, advantages, hparams):
-    raise NotImplementedError
-
-
-def ppo_softmax_loss(logits, logits_old, advantages, hparams):
-    r_t = logits / logits_old
-    r_t_clipped = jnp.clip(r_t, 1 - hparams.epsilon, 1 + hparams.epsilon)
-    return min(r_t, r_t_clipped) * advantages
diff --git a/helx/rl/policy.py b/helx/rl/policy.py
deleted file mode 100644
index 46d3e5a..0000000
--- a/helx/rl/policy.py
+++ /dev/null
@@ -1,72 +0,0 @@
-from abc import ABC, abstractmethod
-from dm_env import TimeStep
-from dm_env.specs import BoundedArray
-import jax
-import jax.numpy as jnp
-
-from helx.random import PRNGSequence
-
-
-class Policy(ABC):
-    @abstractmethod
-    def sample(self, probs: jnp.ndarray, **kwargs):
-        return
-
-    def __call__(self, probs: jnp.ndarray, **kwargs):
-        return self.sample(probs, **kwargs)
-
-
-class EGreedy(Policy):
-    def __init__(
-        self,
-        action_spec: BoundedArray,
-        eps: float = 0.1,
-        start_eps: float = None,
-        end_eps: float = None,
-        start_iter: int = None,
-        end_iter: int = None,
-        seed: int = 0,
-    ):
-        self.action_spec = action_spec
-        self.eps = eps
-        self.start_eps = start_eps
-        self.end_eps = end_eps
-        self.start_iter = start_iter
-        self.end_iter = end_iter
-        self.rng = PRNGSequence(seed)
-
-    def epsilon(self, iteration: int = None):
-        if iteration is None:
-            return self._eps
-
-        x0, y0 = self.start_iter, self.start_eps
-        x1, y1 = self.end_iter, self.end_eps
-        msg = "If `iteration` is provided, the properties `start_iter`, `end_iter`, `start_eps` and `end_eps` must be set"
-        assert None not in (x0, x1, y0, y1), msg
-
-        y = ((y1 - y0) * (iteration - x0) / (x1 - x0)) + y0
-        return min(max(y0, y), y1)
-
-    def rand_prob(self):
-        return jax.random.uniform(next(self.rng))
-
-    def rand_action(self):
-        return jax.random.randint(
-            next(self.rng), (1,), self.action_spec.minimum, self.action_spec.maximum
-        )
-
-    def sample(self, q, iteration):
-        """Draws a random action with ɛ probability,
-        a greedy action with respect to q otherwise"""
-        if self.rand_prob() < self.epsilon(iteration):
-            return self.rand_action()
-        else:
-            return jnp.argmax(q, axis=-1)
-
-
-class Gaussian(Policy):
-    pass
-
-
-class Softmax(Policy):
-    pass
\ No newline at end of file
diff --git a/helx/rl/td.py b/helx/rl/td.py
deleted file mode 100644
index 05ef322..0000000
--- a/helx/rl/td.py
+++ /dev/null
@@ -1,99 +0,0 @@
-from functools import partial
-from typing import Sequence
-import jax
-import jax.numpy as jnp
-
-from helx.typing import Value, Return
-from helx.jax import fori_scan
-from .memory import Trajectory
-
-
-def nstep_return(trajectory: Trajectory, value: Value) -> Return:
-    """n-step return as of:
-    Sutton R., Barto. G., 2018, http://incompleteideas.net/book/RLbook2020.pdf
-    """
-    body_fun = (
-        lambda i, g: trajectory.rewards[-i - 1] + trajectory.discounts[-i - 1] * g
-    )
-    n, init = len(trajectory.rewards), jnp.zeros_like(trajectory.rewards)
-    return (
-        jax.lax.fori_loop(0, n - 1, body_fun, init)
-        + value * trajectory.discounts[0] ** n
-    )
-
-
-def nstep_returns(trajectory: Trajectory, value: Value) -> Return:
-    """n-step return as of:
-    Sutton R., Barto. G., 2018, http://incompleteideas.net/book/RLbook2020.pdf
-    """
-    body_fun = lambda i, x: x * trajectory.discounts[i] ** i
-    n, init = len(trajectory.rewards), trajectory.rewards[-1] * 0.0
-    return (
-        fori_scan(body_fun, 0, n - 1, body_fun, init)
-        + value * trajectory.discounts[-1] ** n
-    )
-
-
-def _lambda_return_t(t, v, i, g):
-    return t.rewards[i] + t.discounts[i] * (
-        t.trace_decays[i] * g + (1 - t.trace_decays[i]) * v[i]
-    )
-
-
-def lambda_return(trajectory: Trajectory, values: Value) -> jnp.ndarray:
-    """λ-returns as of:
-    Sutton R., Barto. G., 2018, http://incompleteideas.net/book/RLbook2020.pdf
-    The off-line λ-return algorithm is a forward view rl algorithm.
-    This function returns all the λ-returns for each time step.
-    For the backward version see td_lambda.
-    This function returns the final λ-return.
-    """
-    body_fun = partial(_lambda_return_t, t=trajectory, v=values)
-    return jax.lax.fori_loop(0, len(trajectory.rewards), body_fun, 0.0)
-
-
-def lambda_returns(trajectory: Trajectory, values: Value) -> jnp.ndarray:
-    body_fun = (
-        lambda g, i: (partial(_lambda_return_t, t=trajectory, v=values)(i, g),) * 2
-    )
-    # the recursive formulation of λ-returns requires to iterate backwards
-    return jnp.flip(fori_scan(0, len(trajectory.rewards), body_fun, 0.0, reversed=True))
-
-
-lambda_returns.__doc__ = (
-    lambda_return.__doc__
-    + """This function returns all the λ-returns for each time step.
-"""
-)
-
-generalised_advantage = lambda_returns
-generalised_advantage.__doc__ = (
-    """The Generalised Advantage Estimator (GAE) uses truncated lambda returns.
-    See: Schulman, J. et al., 2018, https://arxiv.org/abs/1506.02438\n"""
-    + lambda_returns.__doc__
-)
-
-
-def advantages(trajectory: Trajectory, v_t: Value, v_tk: Value):
-    """Advantage function used for actor-critic methods as in:
-    Mnih, V. et al., 2016, https://arxiv.org/abs/1602.01783
-    """
-    n = len(trajectory.rewards)
-    g_t = ewas(trajectory.rewards, trajectory.gammas)
-    return g_t + trajectory.gammas[-1] ** n * v_tk - v_t
-
-
-def ewa(series: Sequence, weights: Sequence) -> float:
-    """Exponentially weighted average"""
-    # makes sure weights and series have same rank
-    w = jnp.broadcast_to(weights, series)
-    body_fun = lambda x, i: x * w[i] ** i
-    return jax.lax.fori_loop(0, len(series), body_fun, (series[0], 0))
-
-
-def ewas(series: Sequence, weights: Sequence) -> jnp.ndarray:
-    """Exponentially weighted average that returns
-    all the intermediate values"""
-    w = jnp.broadcast_to(weights, series)
-    body_fun = lambda x, i: x * w[i] ** i
-    return fori_scan(0, len(series), body_fun, (series[0], 0))
diff --git a/helx/nn/rnn.py b/helx/stax.py
similarity index 92%
rename from helx/nn/rnn.py
rename to helx/stax.py
index 810b3b7..23c139d 100644
--- a/helx/nn/rnn.py
+++ b/helx/stax.py
@@ -1,3 +1,4 @@
+"""Common missing modules to the jax.stax library"""
 from typing import NamedTuple
 
 import jax
@@ -5,10 +6,7 @@
 from jax.nn import sigmoid
 from jax.nn.initializers import glorot_normal, normal, zeros
 
-from .module import module
 
-
-@module
 def Rnn(cell):
     """Layer construction function for an RNN unroll.
     Implements truncated backpropagation through time for the backward pass
@@ -44,7 +42,6 @@ class LSTMState(NamedTuple):
     c: jnp.ndarray
 
 
-@module
 def LSTMCell(
     hidden_size,
     W_init=glorot_normal(),
@@ -61,7 +58,7 @@ def initial_state(rng):
         return LSTMState(h_initial_state_fn(k1, shape), c_initial_state_fn(k2, shape))
 
     def init(rng, input_shape):
-        #  init params
+        #  init params
         in_dim, out_dim = input_shape[-1] + hidden_size, 4 * hidden_size
         output_shape = input_shape[:-1] + (hidden_size,)
         k1, k2, k3 = jax.random.split(rng, 3)
@@ -70,19 +67,18 @@ def init(rng, input_shape):
         return (output_shape, (output_shape, output_shape)), ((W, b), hidden_state)
 
     def apply(params, inputs, **kwargs):
-        prev_state = kwargs.get("prev_state")
+        prev_state = kwargs.get("prev_state")  # type: ignore
         W, b = params
-        xh = jnp.concatenate([inputs, prev_state.h], axis=-1)
+        xh = jnp.concatenate([inputs, prev_state.h], axis=-1)  # type: ignore
         gated = jnp.matmul(xh, W) + b
         i, f, o, g = jnp.split(gated, indices_or_sections=4, axis=-1)
-        c = sigmoid(f) * prev_state.c + sigmoid(i) * jnp.tanh(g)
+        c = sigmoid(f) * prev_state.c + sigmoid(i) * jnp.tanh(g)  # type: ignore
         h = sigmoid(o) * jnp.tanh(c)
         return h, LSTMState(h, c)
 
     return (init, apply)
 
 
-@module
 def LSTM(
     hidden_size,
     W_init=glorot_normal(),
@@ -101,7 +97,6 @@ def LSTM(
     )
 
 
-@module
 def GRUCell(
     hidden_size,
     W_init=glorot_normal(),
@@ -127,11 +122,11 @@ def init(rng, input_shape):
     def apply(params, inputs, **kwargs):
         prev_state = kwargs.get("prev_state")
         W_i, W_h, b = params
-        W_hz, W_ha = jnp.split(W_h, indices_or_sections=(2 * hidden_size,), axis=-1)
-        b_z, b_a = jnp.split(b, indices_or_sections=(2 * hidden_size,), axis=-1)
+        W_hz, W_ha = jnp.split(W_h, indices_or_sections=(2 * hidden_size,), axis=-1)  # type: ignore
+        b_z, b_a = jnp.split(b, indices_or_sections=(2 * hidden_size,), axis=-1)  # type: ignore
 
         gated = jnp.matmul(inputs, W_i)
-        zr_x, a_x = jnp.split(gated, indices_or_sections=[2 * hidden_size], axis=-1)
+        zr_x, a_x = jnp.split(gated, indices_or_sections=[2 * hidden_size], axis=-1)  # type: ignore
         zr_h = jnp.matmul(prev_state, W_hz)
         z, r = jnp.split(
             jax.nn.sigmoid(zr_x + zr_h + b_z), indices_or_sections=2, axis=-1
@@ -144,7 +139,6 @@ def apply(params, inputs, **kwargs):
     return (init, apply)
 
 
-@module
 def GRU(
     hidden_size,
     W_init=glorot_normal(),
diff --git a/helx/typing.py b/helx/typing.py
deleted file mode 100644
index 8ec5e1f..0000000
--- a/helx/typing.py
+++ /dev/null
@@ -1,48 +0,0 @@
-import functools
-from typing import Any, Callable, Dict, NamedTuple, Tuple, TypeVar, Union
-
-import jax.numpy as jnp
-
-Key = TypeVar("Key", bound=jnp.ndarray)
-Shape = TypeVar("Shape", bound=Tuple[int, ...])
-Size = TypeVar("Size", bound=Tuple[int, int])
-Params = TypeVar("Params", bound=Any)
-# Init = TypeVar("Init", Callable[[Key, Shape], Tuple[Shape, Params]])
-# TODO(ep) Error: TypeVar bound type cannot be generic. TODO(ep): add when generics are supported
-Init = Callable[[Key, Shape], Tuple[Shape, Params]]
-# Apply = TypeVar("Apply", Callable[[Params, jnp.ndarray, Dict], jnp.ndarray])
-# TODO(ep) Error: TypeVar bound type cannot be generic.
-Apply = Callable[[Params, jnp.ndarray, Dict], jnp.ndarray]
-
-HParams = TypeVar("HParams", bound=NamedTuple)
-Batch = Union
-
-State = TypeVar("State", bound=jnp.ndarray)
-Observation = TypeVar("Observation", bound=jnp.ndarray)
-Action = TypeVar("Action", bound=int)
-Reward = TypeVar("Reward", bound=float)
-Discount = TypeVar("Discount", bound=float)
-TraceDecay = TypeVar("TraceDecay", bound=float)
-
-Logits = TypeVar("Logits", bound=jnp.ndarray)
-Value = TypeVar("Value", bound=jnp.ndarray)
-Loss = TypeVar("Loss", bound=float)
-Return = TypeVar("Return", bound=float)
-
-
-def factory(cls_maker, T):
-    """Type factory decorator"""
-
-    @functools.wraps(cls_maker)
-    def fabricate(*args, **kwargs):
-        return T(*cls_maker(*args, **kwargs))
-
-    return fabricate
-
-
-def default(type):
-    return type()
-
-
-def nameof(type):
-    return type.__name__
diff --git a/playground.ipynb b/playground.ipynb
deleted file mode 100644
index 3e643b9..0000000
--- a/playground.ipynb
+++ /dev/null
@@ -1,577 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "id": "c163c0b0-8153-4d36-86f8-b7e7f35c4a7b",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%load_ext autoreload\n",
-    "%autoreload 2"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "id": "animal-promise",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "Finishing last run (ID:3rqkd57i) before initializing another..."
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<br/>Waiting for W&B process to finish, PID 19527<br/>Program ended successfully."
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "Find user logs for this run at: <code>/home/epignatelli/repos/helx/wandb/run-20210508_121802-3rqkd57i/logs/debug.log</code>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "Find internal logs for this run at: <code>/home/epignatelli/repos/helx/wandb/run-20210508_121802-3rqkd57i/logs/debug-internal.log</code>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "Synced 5 W&B file(s), 0 media file(s), 3 artifact file(s) and 1 other file(s)"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "\n",
-       "                    <br/>Synced <strong style=\"color:#cdcd00\">dry-pine-1</strong>: <a href=\"https://wandb.ai/epignatelli/helx/runs/3rqkd57i\" target=\"_blank\">https://wandb.ai/epignatelli/helx/runs/3rqkd57i</a><br/>\n",
-       "                "
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "...Successfully finished last run (ID:3rqkd57i). Initializing new run:<br/><br/>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\u001b[34m\u001b[1mwandb\u001b[0m: wandb version 0.10.30 is available!  To upgrade, please run:\n",
-      "\u001b[34m\u001b[1mwandb\u001b[0m:  $ pip install wandb --upgrade\n"
-     ]
-    },
-    {
-     "data": {
-      "text/html": [
-       "\n",
-       "                Tracking run with wandb version 0.10.27<br/>\n",
-       "                Syncing run <strong style=\"color:#cdcd00\">deep-sun-2</strong> to <a href=\"https://wandb.ai\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
-       "                Project page: <a href=\"https://wandb.ai/epignatelli/helx\" target=\"_blank\">https://wandb.ai/epignatelli/helx</a><br/>\n",
-       "                Run page: <a href=\"https://wandb.ai/epignatelli/helx/runs/21qflvyu\" target=\"_blank\">https://wandb.ai/epignatelli/helx/runs/21qflvyu</a><br/>\n",
-       "                Run data is saved locally in <code>/home/epignatelli/repos/helx/wandb/run-20210508_121816-21qflvyu</code><br/><br/>\n",
-       "            "
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "import helx\n",
-    "from helx.rl import baselines, environment\n",
-    "import wandb\n",
-    "\n",
-    "\n",
-    "if __name__ == \"__main__\":\n",
-    "\n",
-    "    env = environment.make_minigrid(\"MiniGrid-Empty-6x6-v0\")\n",
-    "    hparams = baselines.dqn.HParams(\n",
-    "        replay_memory_size=5000, replay_start=5000, batch_size=32\n",
-    "    )\n",
-    "    agent = baselines.dqn.Dqn((56, 56, 3), env.action_spec().num_values, hparams)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 18,
-   "id": "96e0fe85-2d99-4892-821a-99a5761180c5",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n",
-      "random policy\n",
-      "[0]\n"
-     ]
-    }
-   ],
-   "source": [
-    "import jax\n",
-    "timestep = env.reset()\n",
-    "while not timestep.last():\n",
-    "    #  apply policy\n",
-    "    action = agent.policy(timestep)\n",
-    "    ## get action\n",
-    "    self = agent\n",
-    "    if jax.random.uniform(self.rng, (1,)) < self.epsilon:\n",
-    "        print(\"random policy\")\n",
-    "        action = jax.random.randint(self.rng, (1,), 0, self.n_actions)\n",
-    "    else:\n",
-    "        print(\"greedy policy\")\n",
-    "        state = timestep.observation[None, ...]  # batching\n",
-    "        q_values = self.forward(self._online_params, state)\n",
-    "        action = jnp.argmax(q_values)    \n",
-    "    \n",
-    "    ## \n",
-    "    print(action)\n",
-    "    #  observe new state\n",
-    "    new_timestep = agent.observe(env, timestep, action)\n",
-    "    #  update policy\n",
-    "    loss = None\n",
-    "    if not eval:\n",
-    "        loss = agent.update(timestep, action, new_timestep)\n",
-    "    # prepare next iteration\n",
-    "    timestep = new_timestep"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 19,
-   "id": "d39c5bb9-56cf-4911-b5a0-1f15c736d808",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "DeviceArray([0], dtype=int32)"
-      ]
-     },
-     "execution_count": 19,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "jax.random.randint(self.rng, (1,), 0, self.n_actions)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "id": "5761879e-03a7-41bd-83d7-c5e1b3af2c28",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from helx.random import PRNGSequence"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "id": "4b40b57d-54e8-4ad7-8fef-ff8ca8fe9bbf",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "k = PRNGSequence(0)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "id": "3a8657b1-5a35-4ef4-a69e-28cbe71ff6e1",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "[4146024105  967050713]\n",
-      "[3746991216 1902760697]\n",
-      "[2415289872 3870638741]\n",
-      "[1908869515 1172436361]\n",
-      "[ 145021652 2649282573]\n",
-      "[1409544366 3423741852]\n",
-      "[2119404134 3000745408]\n",
-      "[2580787029  482974955]\n",
-      "[1297192877 2312490732]\n",
-      "[3489202107 1640625132]\n"
-     ]
-    }
-   ],
-   "source": [
-    "for i in range(10):\n",
-    "    print(next(k))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "3e151d4d-a73e-4c4b-9065-8deaa0696ebf",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.8.8"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}
diff --git a/setup.py b/setup.py
index 2a73d5c..aa6bc1c 100644
--- a/setup.py
+++ b/setup.py
@@ -1,7 +1,7 @@
 from setuptools import setup, find_packages
 
 
-__version__ = "0.0.7.0"
+__version__ = "0.1.0.0"
 
 
 def parse_requirements(filename):
@@ -18,14 +18,24 @@ def parse_requirements(filename):
     author_email="edu.pignatelli@gmail.com",
     url="https://github.com/epignatelli/helx",
     packages=find_packages(exclude="experiments"),
-    python_requires=">=3.7",
+    python_requires=">=3.9",
     install_requires=[
         "pytest",
         "bsuite",
-        "gym",
-        "gym-minigrid",
+        "gym[all]",
+        "minigrid",
         "wandb",
         "jupyterlab",
         "black",
+        "flake8",
+        "pytest",
+        "jax",
+        "chex",
+        "optax",
+        "rlax",
+        "flax",
+        "dm_env",
+        "wandb",
+        "absl-py",
     ],
 )


Test patch:
diff --git a/test/test_buffer.py b/test/test_buffer.py
deleted file mode 100644
index 948cb87..0000000
--- a/test/test_buffer.py
+++ /dev/null
@@ -1,33 +0,0 @@
-from helx.rl.memory import ReplayBuffer
-from gym_minigrid.wrappers import *
-from bsuite.utils.gym_wrapper import DMEnvFromGym
-from helx.rl import environment
-
-
-def test_buffer():
-    buffer = ReplayBuffer(10, 1)
-    env = environment.make("MiniGrid-Empty-8x8-v0")
-    num_episodes = 10
-
-    for episode in range(num_episodes):
-        print(
-            "Starting episode number {}/{}\t\t\t".format(episode, num_episodes - 1),
-            end="\r",
-        )
-        # initialise environment
-        timestep = env.reset()
-        while not timestep.last():
-            # policy
-            action = 1
-            # step environment
-            new_timestep = env.step(action)
-            # update
-            buffer.add(timestep, action, new_timestep)
-            # prepare next
-            timestep = new_timestep
-
-    transition_batched = buffer.sample(2)
-
-
-if __name__ == "__main__":
-    test_buffer()
diff --git a/test/test_dqn.py b/test/test_dqn.py
deleted file mode 100644
index b0df674..0000000
--- a/test/test_dqn.py
+++ /dev/null
@@ -1,20 +0,0 @@
-import gym
-import jax
-from bsuite.utils.gym_wrapper import DMEnvFromGym
-from gym_minigrid.wrappers import ImgObsWrapper, RGBImgPartialObsWrapper
-from helx.rl import baselines, environment
-
-
-def test_dqn():
-    hparams = baselines.dqn.HParams(
-        replay_memory_size=5000, replay_start=5000, batch_size=32
-    )
-    env = environment.make("MiniGrid-Empty-6x6-v0")
-    preprocess = jax.jit(lambda x: x / 255, backend="cpu")
-    dqn = baselines.dqn.Dqn(
-        (56, 56, 3), env.action_spec().num_values, hparams, preprocess=preprocess
-    )
-
-
-if __name__ == "__main__":
-    test_dqn()
diff --git a/test/test_environment.py b/test/test_environment.py
deleted file mode 100644
index d9b4135..0000000
--- a/test/test_environment.py
+++ /dev/null
@@ -1,38 +0,0 @@
-import logging
-import time
-import multiprocessing as mp
-from helx.rl.environment import make, MultiprocessEnv
-
-
-def test_multiprocess_env():
-    n = 32
-    env = make("MiniGrid-Empty-5x5-v0")
-    env = MultiprocessEnv(env, n)
-    env.reset()
-    logging.debug("resetted env")
-    for i in range(5):
-        env.step((0,) * n)
-        logging.debug("env stepped")
-
-
-def test_multiprocess_env_async():
-    n = 3
-    env = make("MiniGrid-Empty-5x5-v0")
-    env = MultiprocessEnv(env, n)
-    m = mp.Manager()
-    buffer = m.Queue(10)
-    logging.debug(type(buffer))
-    env.reset()
-    logging.debug("environment reset")
-    for i in range(5):
-        env.step_async([0] * n, buffer)
-    #  give the time to the env to compute the step
-    time.sleep(5)
-    for i in range(10):
-        buffer.get()
-    print("Queue size is:", buffer.qsize())
-
-
-if __name__ == "__main__":
-    test_multiprocess_env()
-    test_multiprocess_env_async()
diff --git a/test/test_jax.py b/test/test_jax.py
deleted file mode 100644
index f667de4..0000000
--- a/test/test_jax.py
+++ /dev/null
@@ -1,41 +0,0 @@
-from helx.jax import pure
-import os
-
-from helx.nn.module import module
-import jax
-import jax.numpy as jnp
-from jax.experimental.stax import Dense, serial
-import jax.experimental.stax as nn
-from jax.experimental.optimizers import adam
-from functools import partial
-
-jax.config.update("jax_log_compiles", True)
-
-
-def test_pmodule():
-    jax.config.update("jax_platform_name", "cpu")
-    os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=2"
-    n_devices = jax.local_device_count()
-    print(n_devices)
-
-    @module
-    def m(a, b):
-        return serial(Dense(a), Dense(b))
-
-    m(2, 3)
-
-
-def test_pure():
-    class Foo:
-        @pure
-        def impure(x):
-            return jnp.square(x)
-
-    foo = Foo()
-    foo.impure(jnp.ones((10, 10)))
-    jax.jit(jax.grad(foo.impure))(1.0)
-
-
-if __name__ == "__main__":
-    test_pmodule()
-    test_pure()
diff --git a/test/test_methods.py b/test/test_methods.py
deleted file mode 100644
index 0b7d79d..0000000
--- a/test/test_methods.py
+++ /dev/null
@@ -1,21 +0,0 @@
-import os
-import jax
-from jax.experimental.stax import serial, Dense
-import helx.nn
-
-
-def pmodule_test():
-    jax.config.update("jax_platform_name", "cpu")
-    os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=2"
-    n_devices = jax.local_device_count()
-    print(n_devices)
-
-    @helx.nn.module
-    def m(a, b):
-        return serial(Dense(a), Dense(b))
-
-    m(2, 3)
-
-
-if __name__ == "__main__":
-    pmodule_test()
\ No newline at end of file
diff --git a/test/test_nn.py b/test/test_nn.py
deleted file mode 100644
index 666b73a..0000000
--- a/test/test_nn.py
+++ /dev/null
@@ -1,63 +0,0 @@
-import jax
-from helx.nn import rnn
-
-
-def test_lstm_cell():
-    rng = jax.random.PRNGKey(0)
-    input_shape = (8,)
-    lstm = rnn.LSTMCell(16)
-    x = jax.random.normal(rng, input_shape)
-    out_shape, (params, prev_state) = lstm.init(rng, input_shape)
-    outputs, state = lstm.apply(params, x, prev_state=prev_state)
-    print(outputs.shape, state.h.shape, state.c.shape)
-    jax.jit(lstm.apply)(params, x, prev_state=prev_state)
-    jax.grad(lambda l: sum(lstm.apply(params, x, prev_state=prev_state)[0]))(1.0)
-
-
-def test_lstm():
-    rng = jax.random.PRNGKey(0)
-    SEQ_LEN = 5
-    INPUT_FEATURES = 8
-    HIDDEN_SIZE = 16
-    input_shape = (SEQ_LEN, INPUT_FEATURES)
-    lstm = rnn.LSTM(HIDDEN_SIZE)
-    x = jax.random.normal(rng, input_shape)
-    out_shape, (params, prev_state) = lstm.init(rng, input_shape)
-    outputs, prev_state = lstm.apply(params, x, prev_state=prev_state)
-    print(outputs.shape)
-    jax.jit(lstm.apply)(params, x, prev_state=prev_state)
-    jax.grad(lambda l: sum(sum(lstm.apply(params, x, prev_state=prev_state)[0])))(1.0)
-
-
-def test_gru_cell():
-    rng = jax.random.PRNGKey(0)
-    input_shape = (8,)
-    gru = rnn.LSTMCell(16)
-    x = jax.random.normal(rng, input_shape)
-    out_shape, (params, prev_state) = gru.init(rng, input_shape)
-    outputs, state = gru.apply(params, x, prev_state=prev_state)
-    print(outputs.shape, state.h.shape, state.c.shape)
-    jax.jit(gru.apply)(params, x, prev_state=prev_state)
-    jax.grad(lambda l: sum(gru.apply(params, x, prev_state=prev_state)[0]))(1.0)
-
-
-def test_gru():
-    rng = jax.random.PRNGKey(0)
-    SEQ_LEN = 5
-    INPUT_FEATURES = 8
-    HIDDEN_SIZE = 16
-    input_shape = (SEQ_LEN, INPUT_FEATURES)
-    gru = rnn.LSTM(HIDDEN_SIZE)
-    x = jax.random.normal(rng, input_shape)
-    out_shape, (params, prev_state) = gru.init(rng, input_shape)
-    outputs, hidden_state = gru.apply(params, x, prev_state=prev_state)
-    print(outputs.shape)
-    jax.jit(gru.apply)(params, x, prev_state=prev_state)
-    jax.grad(lambda l: sum(sum(gru.apply(params, x, prev_state=prev_state)[0])))(1.0)
-
-
-if __name__ == "__main__":
-    test_lstm_cell()
-    test_lstm()
-    test_gru_cell()
-    test_gru()
\ No newline at end of file
diff --git a/test/test_random.py b/test/test_random.py
deleted file mode 100644
index 838a2dd..0000000
--- a/test/test_random.py
+++ /dev/null
@@ -1,57 +0,0 @@
-import jax
-import jax.numpy as jnp
-
-import helx
-
-
-def test_shuffled_batch_indices():
-    rng = jax.random.PRNGKey(0)
-    indices = list(range(10))
-
-    batch_indices_no_drop = list(
-        map(
-            jnp.array,
-            ([2, 7, 9], [6, 0, 8], [1, 3, 4], [5]),
-        )
-    )
-    batch_indices_drop = list(
-        map(
-            jnp.array,
-            ([2, 7, 9], [6, 0, 8], [1, 3, 4]),
-        )
-    )
-    batch_indices_even_drop = list(
-        map(
-            jnp.array,
-            ([2, 7], [9, 6], [0, 8], [1, 3], [4, 5]),
-        )
-    )
-
-    # test spare, no drop
-    res_no_drop = helx.random.shuffled_batched_indices(rng, indices, 3, drop_last=False)
-    assert all(
-        [
-            jnp.array_equal(res_no_drop[i], batch_indices_no_drop[i])
-            for i in range(len(batch_indices_no_drop))
-        ]
-    ), "\n{} \nnot equal to \n{}".format(res_no_drop, batch_indices_no_drop)
-
-    # test spare drop
-    res_drop = helx.random.shuffled_batched_indices(rng, indices, 3, drop_last=True)
-    assert all(
-        [
-            jnp.array_equal(res_drop[i], batch_indices_drop[i])
-            for i in range(len(batch_indices_drop))
-        ]
-    ), "\n{} \nnot equal to \n{}".format(res_drop, batch_indices_drop)
-
-    # test even drop
-    res_even_drop = helx.random.shuffled_batched_indices(
-        rng, indices, 2, drop_last=True
-    )
-    assert all(
-        [
-            jnp.array_equal(res_even_drop[i], batch_indices_even_drop[i])
-            for i in range(len(batch_indices_no_drop))
-        ]
-    ), "\n{} \nnot equal to \n{}".format(res_even_drop, batch_indices_even_drop)