Repository: cpnota/autonomous-learning-library (language: python)

Issue:
Update Linting/Formatting
Some of the linting/formatting rules could stand to be updated. In particular, I want some stricter rules around unused imports.
Imports in scripts
@cpnota you may want to read through the imports in the files in the scripts folder. I'm pretty sure a decent number are unnecessary.


Discussion / hints:

Will address with #294 if still an issue.

Fix patch:
diff --git a/Makefile b/Makefile
index f2c5e5ef..cac991d1 100644
--- a/Makefile
+++ b/Makefile
@@ -11,10 +11,13 @@ integration-test:
 	python -m unittest discover -s integration -p "*test.py"
 
 lint:
-	flake8 --ignore "E501,E731,E74,E402,F401,W503,E128" all
+	black --check all benchmarks examples integration setup.py
+	isort --profile black --check all benchmarks examples integration setup.py
+	flake8 --select "F401" all benchmarks examples integration setup.py
 
 format:
-	autopep8 --in-place --aggressive --aggressive --ignore "E501,E731,E74,E402,F401,W503,E128" -r all
+	black all benchmarks examples integration setup.py
+	isort --profile black all benchmarks examples integration setup.py
 
 tensorboard:
 	tensorboard --logdir runs
diff --git a/all/__init__.py b/all/__init__.py
index 3bfd56fc..0432db2e 100644
--- a/all/__init__.py
+++ b/all/__init__.py
@@ -1,26 +1,16 @@
-import all.agents
-import all.approximation
-import all.core
-import all.environments
-import all.logging
-import all.memory
-import all.nn
-import all.optim
-import all.policies
-import all.presets
 from all.core import State, StateArray
 
 __all__ = [
-    'agents',
-    'approximation',
-    'core',
-    'environments',
-    'logging',
-    'memory',
-    'nn',
-    'optim',
-    'policies',
-    'presets',
-    'State',
-    'StateArray'
+    "agents",
+    "approximation",
+    "core",
+    "environments",
+    "logging",
+    "memory",
+    "nn",
+    "optim",
+    "policies",
+    "presets",
+    "State",
+    "StateArray",
 ]
diff --git a/all/agents/__init__.py b/all/agents/__init__.py
index 2811a056..7a62e6b0 100644
--- a/all/agents/__init__.py
+++ b/all/agents/__init__.py
@@ -15,7 +15,6 @@
 from .vqn import VQN, VQNTestAgent
 from .vsarsa import VSarsa, VSarsaTestAgent
 
-
 __all__ = [
     # Agent interfaces
     "Agent",
diff --git a/all/agents/_agent.py b/all/agents/_agent.py
index 45038f14..ed4cb357 100644
--- a/all/agents/_agent.py
+++ b/all/agents/_agent.py
@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+
 from all.optim import Schedulable
 
 
diff --git a/all/agents/_multiagent.py b/all/agents/_multiagent.py
index 11f4a18d..e98e0e40 100644
--- a/all/agents/_multiagent.py
+++ b/all/agents/_multiagent.py
@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+
 from all.optim import Schedulable
 
 
diff --git a/all/agents/_parallel_agent.py b/all/agents/_parallel_agent.py
index b98ae618..cba3b42d 100644
--- a/all/agents/_parallel_agent.py
+++ b/all/agents/_parallel_agent.py
@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+
 from all.optim import Schedulable
 
 
diff --git a/all/agents/a2c.py b/all/agents/a2c.py
index 082b9309..ccc97d48 100644
--- a/all/agents/a2c.py
+++ b/all/agents/a2c.py
@@ -1,7 +1,8 @@
-import torch
 from torch.nn.functional import mse_loss
+
 from all.logging import DummyLogger
 from all.memory import NStepAdvantageBuffer
+
 from ._agent import Agent
 from ._parallel_agent import ParallelAgent
 
@@ -28,15 +29,15 @@ class A2C(ParallelAgent):
     """
 
     def __init__(
-            self,
-            features,
-            v,
-            policy,
-            discount_factor=0.99,
-            entropy_loss_scaling=0.01,
-            n_envs=None,
-            n_steps=4,
-            logger=DummyLogger()
+        self,
+        features,
+        v,
+        policy,
+        discount_factor=0.99,
+        entropy_loss_scaling=0.01,
+        n_envs=None,
+        n_steps=4,
+        logger=DummyLogger(),
     ):
         if n_envs is None:
             raise RuntimeError("Must specify n_envs.")
@@ -80,7 +81,9 @@ def _train(self, next_states):
             value_loss = mse_loss(values, targets)
             policy_gradient_loss = -(distribution.log_prob(actions) * advantages).mean()
             entropy_loss = -distribution.entropy().mean()
-            policy_loss = policy_gradient_loss + self.entropy_loss_scaling * entropy_loss
+            policy_loss = (
+                policy_gradient_loss + self.entropy_loss_scaling * entropy_loss
+            )
             loss = value_loss + policy_loss
 
             # backward pass
@@ -90,8 +93,8 @@ def _train(self, next_states):
             self.features.step()
 
             # record metrics
-            self.logger.add_info('entropy', -entropy_loss)
-            self.logger.add_info('normalized_value_error', value_loss / targets.var())
+            self.logger.add_info("entropy", -entropy_loss)
+            self.logger.add_info("normalized_value_error", value_loss / targets.var())
 
     def _make_buffer(self):
         return NStepAdvantageBuffer(
@@ -99,7 +102,7 @@ def _make_buffer(self):
             self.features,
             self.n_steps,
             self.n_envs,
-            discount_factor=self.discount_factor
+            discount_factor=self.discount_factor,
         )
 
 
diff --git a/all/agents/c51.py b/all/agents/c51.py
index 8bdbcc96..3a1650b6 100644
--- a/all/agents/c51.py
+++ b/all/agents/c51.py
@@ -1,6 +1,8 @@
-import torch
 import numpy as np
+import torch
+
 from all.logging import DummyLogger
+
 from ._agent import Agent
 
 
@@ -26,16 +28,16 @@ class C51(Agent):
     """
 
     def __init__(
-            self,
-            q_dist,
-            replay_buffer,
-            discount_factor=0.99,
-            eps=1e-5,
-            exploration=0.02,
-            minibatch_size=32,
-            replay_start_size=5000,
-            update_frequency=1,
-            logger=DummyLogger(),
+        self,
+        q_dist,
+        replay_buffer,
+        discount_factor=0.99,
+        eps=1e-5,
+        exploration=0.02,
+        minibatch_size=32,
+        replay_start_size=5000,
+        update_frequency=1,
+        logger=DummyLogger(),
     ):
         # objects
         self.q_dist = q_dist
@@ -81,7 +83,9 @@ def _best_actions(self, probs):
     def _train(self):
         if self._should_train():
             # sample transitions from buffer
-            states, actions, rewards, next_states, weights = self.replay_buffer.sample(self.minibatch_size)
+            states, actions, rewards, next_states, weights = self.replay_buffer.sample(
+                self.minibatch_size
+            )
             # forward pass
             dist = self.q_dist(states, actions)
             # compute target distribution
@@ -100,14 +104,15 @@ def _train(self):
 
     def _should_train(self):
         self._frames_seen += 1
-        return self._frames_seen > self.replay_start_size and self._frames_seen % self.update_frequency == 0
+        return (
+            self._frames_seen > self.replay_start_size
+            and self._frames_seen % self.update_frequency == 0
+        )
 
     def _compute_target_dist(self, states, rewards):
         actions = self._best_actions(self.q_dist.no_grad(states))
         dist = self.q_dist.target(states, actions)
-        shifted_atoms = (
-            rewards.view((-1, 1)) + self.discount_factor * self.q_dist.atoms
-        )
+        shifted_atoms = rewards.view((-1, 1)) + self.discount_factor * self.q_dist.atoms
         return self.q_dist.project(dist, shifted_atoms)
 
     def _kl(self, dist, target_dist):
@@ -117,7 +122,7 @@ def _kl(self, dist, target_dist):
 
 
 class C51TestAgent(Agent):
-    def __init__(self, q_dist, n_actions, exploration=0.):
+    def __init__(self, q_dist, n_actions, exploration=0.0):
         self.q_dist = q_dist
         self.n_actions = n_actions
         self.exploration = exploration
diff --git a/all/agents/ddpg.py b/all/agents/ddpg.py
index 6a01f0ae..76bb5712 100644
--- a/all/agents/ddpg.py
+++ b/all/agents/ddpg.py
@@ -1,6 +1,7 @@
 import torch
 from torch.distributions.normal import Normal
 from torch.nn.functional import mse_loss
+
 from ._agent import Agent
 
 
@@ -28,17 +29,18 @@ class DDPG(Agent):
         update_frequency (int): Number of timesteps per training update.
     """
 
-    def __init__(self,
-                 q,
-                 policy,
-                 replay_buffer,
-                 action_space,
-                 discount_factor=0.99,
-                 minibatch_size=32,
-                 noise=0.1,
-                 replay_start_size=5000,
-                 update_frequency=1,
-                 ):
+    def __init__(
+        self,
+        q,
+        policy,
+        replay_buffer,
+        action_space,
+        discount_factor=0.99,
+        minibatch_size=32,
+        noise=0.1,
+        replay_start_size=5000,
+        update_frequency=1,
+    ):
         # objects
         self.q = q
         self.policy = policy
@@ -49,7 +51,13 @@ def __init__(self,
         self.minibatch_size = minibatch_size
         self.discount_factor = discount_factor
         # private
-        self._noise = Normal(0, noise * torch.tensor((action_space.high - action_space.low) / 2).to(policy.device))
+        self._noise = Normal(
+            0,
+            noise
+            * torch.tensor((action_space.high - action_space.low) / 2).to(
+                policy.device
+            ),
+        )
         self._low = torch.tensor(action_space.low, device=policy.device)
         self._high = torch.tensor(action_space.high, device=policy.device)
         self._state = None
@@ -76,11 +84,15 @@ def _choose_action(self, state):
     def _train(self):
         if self._should_train():
             # sample transitions from buffer
-            (states, actions, rewards, next_states, _) = self.replay_buffer.sample(self.minibatch_size)
+            (states, actions, rewards, next_states, _) = self.replay_buffer.sample(
+                self.minibatch_size
+            )
 
             # train q-network
             q_values = self.q(states, actions)
-            targets = rewards + self.discount_factor * self.q.target(next_states, self.policy.target(next_states))
+            targets = rewards + self.discount_factor * self.q.target(
+                next_states, self.policy.target(next_states)
+            )
             loss = mse_loss(q_values, targets)
             self.q.reinforce(loss)
 
@@ -92,7 +104,10 @@ def _train(self):
 
     def _should_train(self):
         self._frames_seen += 1
-        return self._frames_seen > self.replay_start_size and self._frames_seen % self.update_frequency == 0
+        return (
+            self._frames_seen > self.replay_start_size
+            and self._frames_seen % self.update_frequency == 0
+        )
 
 
 class DDPGTestAgent(Agent):
diff --git a/all/agents/ddqn.py b/all/agents/ddqn.py
index 0f9a1acc..f6bd9246 100644
--- a/all/agents/ddqn.py
+++ b/all/agents/ddqn.py
@@ -1,11 +1,13 @@
 import torch
+
 from all.nn import weighted_mse_loss
+
 from ._agent import Agent
 from .dqn import DQNTestAgent
 
 
 class DDQN(Agent):
-    '''
+    """
     Double Deep Q-Network (DDQN).
     DDQN is an enchancment to DQN that uses a "double Q-style" update,
     wherein the online network is used to select target actions
@@ -24,18 +26,19 @@ class DDQN(Agent):
         minibatch_size (int): The number of experiences to sample in each training update.
         replay_start_size (int): Number of experiences in replay buffer when training begins.
         update_frequency (int): Number of timesteps per training update.
-    '''
+    """
 
-    def __init__(self,
-                 q,
-                 policy,
-                 replay_buffer,
-                 discount_factor=0.99,
-                 loss=weighted_mse_loss,
-                 minibatch_size=32,
-                 replay_start_size=5000,
-                 update_frequency=1,
-                 ):
+    def __init__(
+        self,
+        q,
+        policy,
+        replay_buffer,
+        discount_factor=0.99,
+        loss=weighted_mse_loss,
+        minibatch_size=32,
+        replay_start_size=5000,
+        update_frequency=1,
+    ):
         # objects
         self.q = q
         self.policy = policy
@@ -64,12 +67,16 @@ def eval(self, state):
     def _train(self):
         if self._should_train():
             # sample transitions from buffer
-            (states, actions, rewards, next_states, weights) = self.replay_buffer.sample(self.minibatch_size)
+            (states, actions, rewards, next_states, weights) = (
+                self.replay_buffer.sample(self.minibatch_size)
+            )
             # forward pass
             values = self.q(states, actions)
             # compute targets
             next_actions = torch.argmax(self.q.no_grad(next_states), dim=1)
-            targets = rewards + self.discount_factor * self.q.target(next_states, next_actions)
+            targets = rewards + self.discount_factor * self.q.target(
+                next_states, next_actions
+            )
             # compute loss
             loss = self.loss(values, targets, weights)
             # backward pass
@@ -80,7 +87,10 @@ def _train(self):
 
     def _should_train(self):
         self._frames_seen += 1
-        return self._frames_seen > self.replay_start_size and self._frames_seen % self.update_frequency == 0
+        return (
+            self._frames_seen > self.replay_start_size
+            and self._frames_seen % self.update_frequency == 0
+        )
 
 
 DDQNTestAgent = DQNTestAgent
diff --git a/all/agents/dqn.py b/all/agents/dqn.py
index 0e39df3b..e0a40388 100644
--- a/all/agents/dqn.py
+++ b/all/agents/dqn.py
@@ -1,11 +1,11 @@
-import numpy as np
 import torch
 from torch.nn.functional import mse_loss
+
 from ._agent import Agent
 
 
 class DQN(Agent):
-    '''
+    """
     Deep Q-Network (DQN).
     DQN was one of the original deep reinforcement learning algorithms.
     It extends the ideas behind Q-learning to work well with modern convolution networks.
@@ -25,18 +25,19 @@ class DQN(Agent):
         n_actions (int): The number of available actions.
         replay_start_size (int): Number of experiences in replay buffer when training begins.
         update_frequency (int): Number of timesteps per training update.
-    '''
+    """
 
-    def __init__(self,
-                 q,
-                 policy,
-                 replay_buffer,
-                 discount_factor=0.99,
-                 loss=mse_loss,
-                 minibatch_size=32,
-                 replay_start_size=5000,
-                 update_frequency=1,
-                 ):
+    def __init__(
+        self,
+        q,
+        policy,
+        replay_buffer,
+        discount_factor=0.99,
+        loss=mse_loss,
+        minibatch_size=32,
+        replay_start_size=5000,
+        update_frequency=1,
+    ):
         # objects
         self.q = q
         self.policy = policy
@@ -65,11 +66,16 @@ def eval(self, state):
     def _train(self):
         if self._should_train():
             # sample transitions from buffer
-            (states, actions, rewards, next_states, _) = self.replay_buffer.sample(self.minibatch_size)
+            (states, actions, rewards, next_states, _) = self.replay_buffer.sample(
+                self.minibatch_size
+            )
             # forward pass
             values = self.q(states, actions)
             # compute targets
-            targets = rewards + self.discount_factor * torch.max(self.q.target(next_states), dim=1)[0]
+            targets = (
+                rewards
+                + self.discount_factor * torch.max(self.q.target(next_states), dim=1)[0]
+            )
             # compute loss
             loss = self.loss(values, targets)
             # backward pass
@@ -77,7 +83,10 @@ def _train(self):
 
     def _should_train(self):
         self._frames_seen += 1
-        return (self._frames_seen > self.replay_start_size and self._frames_seen % self.update_frequency == 0)
+        return (
+            self._frames_seen > self.replay_start_size
+            and self._frames_seen % self.update_frequency == 0
+        )
 
 
 class DQNTestAgent(Agent):
diff --git a/all/agents/independent.py b/all/agents/independent.py
index 32c3275f..886d398d 100644
--- a/all/agents/independent.py
+++ b/all/agents/independent.py
@@ -6,4 +6,4 @@ def __init__(self, agents):
         self.agents = agents
 
     def act(self, state):
-        return self.agents[state['agent']].act(state)
+        return self.agents[state["agent"]].act(state)
diff --git a/all/agents/ppo.py b/all/agents/ppo.py
index d13cde08..ce699279 100644
--- a/all/agents/ppo.py
+++ b/all/agents/ppo.py
@@ -1,8 +1,9 @@
 import torch
 from torch.nn.functional import mse_loss
+
 from all.logging import DummyLogger
 from all.memory import GeneralizedAdvantageBuffer
-from ._agent import Agent
+
 from ._parallel_agent import ParallelAgent
 from .a2c import A2CTestAgent
 
@@ -31,20 +32,20 @@ class PPO(ParallelAgent):
     """
 
     def __init__(
-            self,
-            features,
-            v,
-            policy,
-            discount_factor=0.99,
-            entropy_loss_scaling=0.01,
-            epochs=4,
-            epsilon=0.2,
-            lam=0.95,
-            minibatches=4,
-            compute_batch_size=256,
-            n_envs=None,
-            n_steps=4,
-            logger=DummyLogger()
+        self,
+        features,
+        v,
+        policy,
+        discount_factor=0.99,
+        entropy_loss_scaling=0.01,
+        epochs=4,
+        epsilon=0.2,
+        lam=0.95,
+        minibatches=4,
+        compute_batch_size=256,
+        n_envs=None,
+        n_steps=4,
+        logger=DummyLogger(),
     ):
         if n_envs is None:
             raise RuntimeError("Must specify n_envs.")
@@ -85,10 +86,18 @@ def _train(self, next_states):
             states, actions, advantages = self._buffer.advantages(next_states)
 
             # compute target values
-            features = states.batch_execute(self.compute_batch_size, self.features.no_grad)
-            features['actions'] = actions
-            pi_0 = features.batch_execute(self.compute_batch_size, lambda s: self.policy.no_grad(s).log_prob(s['actions']))
-            targets = features.batch_execute(self.compute_batch_size, self.v.no_grad) + advantages
+            features = states.batch_execute(
+                self.compute_batch_size, self.features.no_grad
+            )
+            features["actions"] = actions
+            pi_0 = features.batch_execute(
+                self.compute_batch_size,
+                lambda s: self.policy.no_grad(s).log_prob(s["actions"]),
+            )
+            targets = (
+                features.batch_execute(self.compute_batch_size, self.v.no_grad)
+                + advantages
+            )
 
             # train for several epochs
             for _ in range(self.epochs):
@@ -105,7 +114,9 @@ def _train_epoch(self, states, actions, advantages, targets, pi_0):
             i = indexes[first:last]
 
             # perform a single training step
-            self._train_minibatch(states[i], actions[i], pi_0[i], advantages[i], targets[i])
+            self._train_minibatch(
+                states[i], actions[i], pi_0[i], advantages[i], targets[i]
+            )
 
     def _train_minibatch(self, states, actions, pi_0, advantages, targets):
         # forward pass
@@ -116,7 +127,9 @@ def _train_minibatch(self, states, actions, pi_0, advantages, targets):
 
         # compute losses
         value_loss = mse_loss(values, targets)
-        policy_gradient_loss = self._clipped_policy_gradient_loss(pi_0, pi_i, advantages)
+        policy_gradient_loss = self._clipped_policy_gradient_loss(
+            pi_0, pi_i, advantages
+        )
         entropy_loss = -distribution.entropy().mean()
         policy_loss = policy_gradient_loss + self.entropy_loss_scaling * entropy_loss
         loss = value_loss + policy_loss
@@ -128,8 +141,8 @@ def _train_minibatch(self, states, actions, pi_0, advantages, targets):
         self.features.step()
 
         # debugging
-        self.logger.add_info('entropy', -entropy_loss)
-        self.logger.add_info('normalized_value_error', value_loss / targets.var())
+        self.logger.add_info("entropy", -entropy_loss)
+        self.logger.add_info("normalized_value_error", value_loss / targets.var())
 
     def _clipped_policy_gradient_loss(self, pi_0, pi_i, advantages):
         ratios = torch.exp(pi_i - pi_0)
@@ -146,7 +159,7 @@ def _make_buffer(self):
             self.n_envs,
             discount_factor=self.discount_factor,
             lam=self.lam,
-            compute_batch_size=self.compute_batch_size
+            compute_batch_size=self.compute_batch_size,
         )
 
 
diff --git a/all/agents/sac.py b/all/agents/sac.py
index c81e47ca..ef537c39 100644
--- a/all/agents/sac.py
+++ b/all/agents/sac.py
@@ -1,6 +1,8 @@
 import torch
 from torch.nn.functional import mse_loss
+
 from all.logging import DummyLogger
+
 from ._agent import Agent
 
 
@@ -29,21 +31,22 @@ class SAC(Agent):
         update_frequency (int): Number of timesteps per training update.
     """
 
-    def __init__(self,
-                 policy,
-                 q1,
-                 q2,
-                 replay_buffer,
-                 discount_factor=0.99,
-                 entropy_backups=True,
-                 entropy_target=-2.,
-                 lr_temperature=1e-4,
-                 minibatch_size=32,
-                 replay_start_size=5000,
-                 temperature_initial=0.1,
-                 update_frequency=1,
-                 logger=DummyLogger()
-                 ):
+    def __init__(
+        self,
+        policy,
+        q1,
+        q2,
+        replay_buffer,
+        discount_factor=0.99,
+        entropy_backups=True,
+        entropy_target=-2.0,
+        lr_temperature=1e-4,
+        minibatch_size=32,
+        replay_start_size=5000,
+        temperature_initial=0.1,
+        update_frequency=1,
+        logger=DummyLogger(),
+    ):
         # objects
         self.policy = policy
         self.q1 = q1
@@ -74,7 +77,9 @@ def act(self, state):
     def _train(self):
         if self._should_train():
             # sample from replay buffer
-            (states, actions, rewards, next_states, _) = self.replay_buffer.sample(self.minibatch_size)
+            (states, actions, rewards, next_states, _) = self.replay_buffer.sample(
+                self.minibatch_size
+            )
             discount_factor = self.discount_factor
 
             # compute targets for Q and V
@@ -100,21 +105,28 @@ def _train(self):
             self.q1.zero_grad()
 
             # adjust temperature
-            temperature_grad = (new_log_probs + self.entropy_target).mean() * self.temperature
-            self.temperature = max(0, self.temperature + self.lr_temperature * temperature_grad.detach())
+            temperature_grad = (
+                new_log_probs + self.entropy_target
+            ).mean() * self.temperature
+            self.temperature = max(
+                0, self.temperature + self.lr_temperature * temperature_grad.detach()
+            )
 
             # additional debugging info
-            self.logger.add_info('entropy', -new_log_probs.mean())
-            self.logger.add_info('q_values', q_values.mean())
-            self.logger.add_info('rewards', rewards.mean())
-            self.logger.add_info('normalized_q1_error', q1_loss / q_targets.var())
-            self.logger.add_info('normalized_q2_error', q2_loss / q_targets.var())
-            self.logger.add_info('temperature', self.temperature)
-            self.logger.add_info('temperature_grad', temperature_grad)
+            self.logger.add_info("entropy", -new_log_probs.mean())
+            self.logger.add_info("q_values", q_values.mean())
+            self.logger.add_info("rewards", rewards.mean())
+            self.logger.add_info("normalized_q1_error", q1_loss / q_targets.var())
+            self.logger.add_info("normalized_q2_error", q2_loss / q_targets.var())
+            self.logger.add_info("temperature", self.temperature)
+            self.logger.add_info("temperature_grad", temperature_grad)
 
     def _should_train(self):
         self._frames_seen += 1
-        return self._frames_seen > self.replay_start_size and self._frames_seen % self.update_frequency == 0
+        return (
+            self._frames_seen > self.replay_start_size
+            and self._frames_seen % self.update_frequency == 0
+        )
 
 
 class SACTestAgent(Agent):
diff --git a/all/agents/vac.py b/all/agents/vac.py
index b11e55ef..60655b1d 100644
--- a/all/agents/vac.py
+++ b/all/agents/vac.py
@@ -1,11 +1,11 @@
 from torch.nn.functional import mse_loss
-from ._agent import Agent
+
 from ._parallel_agent import ParallelAgent
 from .a2c import A2CTestAgent
 
 
 class VAC(ParallelAgent):
-    '''
+    """
     Vanilla Actor-Critic (VAC).
     VAC is an implementation of the actor-critic alogorithm found in the Sutton and Barto (2018) textbook.
     This implementation tweaks the algorithm slightly by using a shared feature layer.
@@ -20,7 +20,7 @@ class VAC(ParallelAgent):
         n_envs (int): Number of parallel actors/environments
         n_steps (int): Number of timesteps per rollout. Updates are performed once per rollout.
         logger (Logger): Used for logging.
-    '''
+    """
 
     def __init__(self, features, v, policy, discount_factor=1):
         self.features = features
@@ -47,12 +47,16 @@ def _train(self, state, reward):
             values = self.v(self._features)
 
             # compute targets
-            targets = reward + self.discount_factor * self.v.target(self.features.target(state))
+            targets = reward + self.discount_factor * self.v.target(
+                self.features.target(state)
+            )
             advantages = targets - values.detach()
 
             # compute losses
             value_loss = mse_loss(values, targets)
-            policy_loss = -(advantages * self._distribution.log_prob(self._action)).mean()
+            policy_loss = -(
+                advantages * self._distribution.log_prob(self._action)
+            ).mean()
             loss = value_loss + policy_loss
 
             # backward pass
diff --git a/all/agents/vpg.py b/all/agents/vpg.py
index e46e09ac..c0b54512 100644
--- a/all/agents/vpg.py
+++ b/all/agents/vpg.py
@@ -1,12 +1,14 @@
 import torch
 from torch.nn.functional import mse_loss
+
 from all.core import State
+
 from ._agent import Agent
 from .a2c import A2CTestAgent
 
 
 class VPG(Agent):
-    '''
+    """
     Vanilla Policy Gradient (VPG/REINFORCE).
     VPG (also known as REINFORCE) is the least biased implementation of the policy gradient theorem.
     It uses complete episode rollouts as unbiased estimates of the Q-function, rather than the n-step
@@ -25,16 +27,9 @@ class VPG(Agent):
         min_batch_size (int): Updates will occurs when an episode ends after at least
             this many state-action pairs are seen. Set this to a large value in order
             to train on multiple episodes at once.
-    '''
-
-    def __init__(
-            self,
-            features,
-            v,
-            policy,
-            discount_factor=0.99,
-            min_batch_size=1
-    ):
+    """
+
+    def __init__(self, features, v, policy, discount_factor=0.99, min_batch_size=1):
         self.features = features
         self.v = v
         self.policy = policy
@@ -92,21 +87,20 @@ def _terminal(self, state, reward):
 
     def _train(self):
         # forward pass
-        values = torch.cat([
-            self.v(features)
-            for (features, _, _)
-            in self._trajectories
-        ])
+        values = torch.cat(
+            [self.v(features) for (features, _, _) in self._trajectories]
+        )
 
         # forward passes for log_pis were stored during execution
         log_pis = torch.cat([log_pis for (_, _, log_pis) in self._trajectories])
 
         # compute targets
-        targets = torch.cat([
-            self._compute_discounted_returns(rewards)
-            for (_, rewards, _)
-            in self._trajectories
-        ])
+        targets = torch.cat(
+            [
+                self._compute_discounted_returns(rewards)
+                for (_, rewards, _) in self._trajectories
+            ]
+        )
         advantages = targets - values.detach()
 
         # compute losses
diff --git a/all/agents/vqn.py b/all/agents/vqn.py
index b4847b13..49ae9673 100644
--- a/all/agents/vqn.py
+++ b/all/agents/vqn.py
@@ -1,12 +1,12 @@
 import torch
 from torch.nn.functional import mse_loss
+
 from ._agent import Agent
 from ._parallel_agent import ParallelAgent
-from .dqn import DQNTestAgent
 
 
 class VQN(ParallelAgent):
-    '''
+    """
     Vanilla Q-Network (VQN).
     VQN is an implementation of the Q-learning algorithm found in the Sutton and Barto (2018) textbook.
     Q-learning algorithms attempt to learning the optimal policy while executing a (generally)
@@ -19,7 +19,7 @@ class VQN(ParallelAgent):
         q (QNetwork): An Approximation of the Q function.
         policy (GreedyPolicy): A policy derived from the Q-function.
         discount_factor (float): Discount factor for future rewards.
-    '''
+    """
 
     def __init__(self, q, policy, discount_factor=0.99):
         self.q = q
@@ -43,7 +43,10 @@ def _train(self, reward, next_state):
             # forward pass
             value = self.q(self._state, self._action)
             # compute target
-            target = reward + self.discount_factor * torch.max(self.q.target(next_state), dim=1)[0]
+            target = (
+                reward
+                + self.discount_factor * torch.max(self.q.target(next_state), dim=1)[0]
+            )
             # compute loss
             loss = mse_loss(value, target)
             # backward pass
diff --git a/all/agents/vsarsa.py b/all/agents/vsarsa.py
index d0b51646..a96e1b08 100644
--- a/all/agents/vsarsa.py
+++ b/all/agents/vsarsa.py
@@ -1,10 +1,11 @@
 from torch.nn.functional import mse_loss
+
 from ._parallel_agent import ParallelAgent
 from .vqn import VQNTestAgent
 
 
 class VSarsa(ParallelAgent):
-    '''
+    """
     Vanilla SARSA (VSarsa).
     SARSA (State-Action-Reward-State-Action) is an on-policy alternative to Q-learning. Unlike Q-learning,
     SARSA attempts to learn the Q-function for the current policy rather than the optimal policy. This
@@ -15,7 +16,7 @@ class VSarsa(ParallelAgent):
         q (QNetwork): An Approximation of the Q function.
         policy (GreedyPolicy): A policy derived from the Q-function.
         discount_factor (float): Discount factor for future rewards.
-    '''
+    """
 
     def __init__(self, q, policy, discount_factor=0.99):
         self.q = q
@@ -39,7 +40,9 @@ def _train(self, reward, next_state, next_action):
             # forward pass
             value = self.q(self._state, self._action)
             # compute target
-            target = reward + self.discount_factor * self.q.target(next_state, next_action)
+            target = reward + self.discount_factor * self.q.target(
+                next_state, next_action
+            )
             # compute loss
             loss = mse_loss(value, target)
             # backward pass
diff --git a/all/approximation/__init__.py b/all/approximation/__init__.py
index 2890f66c..89865f0e 100644
--- a/all/approximation/__init__.py
+++ b/all/approximation/__init__.py
@@ -1,13 +1,12 @@
 from .approximation import Approximation
+from .checkpointer import Checkpointer, DummyCheckpointer, PeriodicCheckpointer
+from .feature_network import FeatureNetwork
+from .identity import Identity
 from .q_continuous import QContinuous
 from .q_dist import QDist
 from .q_network import QNetwork
+from .target import FixedTarget, PolyakTarget, TargetNetwork, TrivialTarget
 from .v_network import VNetwork
-from .feature_network import FeatureNetwork
-from .identity import Identity
-from .target import TargetNetwork, FixedTarget, PolyakTarget, TrivialTarget
-from .checkpointer import Checkpointer, DummyCheckpointer, PeriodicCheckpointer
-
 
 __all__ = [
     "Approximation",
diff --git a/all/approximation/approximation.py b/all/approximation/approximation.py
index 3c346650..84a03966 100644
--- a/all/approximation/approximation.py
+++ b/all/approximation/approximation.py
@@ -1,15 +1,18 @@
 import os
+
 import torch
 from torch.nn import utils
+
 from all.logging import DummyLogger
-from .target import TrivialTarget
+
 from .checkpointer import DummyCheckpointer
+from .target import TrivialTarget
 
 DEFAULT_CHECKPOINT_FREQUENCY = 200
 
 
-class Approximation():
-    '''
+class Approximation:
+    """
     Base function approximation object.
 
     This defines a Pytorch-based function approximation object that
@@ -48,20 +51,20 @@ class Approximation():
             logger (all.logging.Logger:, optional): A Logger object used for logging.
                 The standard object logs to tensorboard, however, other types of Logger objects
                 may be implemented by the user.
-    '''
+    """
 
     def __init__(
-            self,
-            model,
-            optimizer=None,
-            checkpointer=None,
-            clip_grad=0,
-            device=None,
-            loss_scaling=1,
-            name='approximation',
-            scheduler=None,
-            target=None,
-            logger=DummyLogger(),
+        self,
+        model,
+        optimizer=None,
+        checkpointer=None,
+        clip_grad=0,
+        device=None,
+        loss_scaling=1,
+        name="approximation",
+        scheduler=None,
+        target=None,
+        logger=DummyLogger(),
     ):
         self.model = model
         self.device = device if device else next(model.parameters()).device
@@ -79,27 +82,24 @@ def __init__(
         if checkpointer is None:
             checkpointer = DummyCheckpointer()
         self._checkpointer = checkpointer
-        self._checkpointer.init(
-            self.model,
-            os.path.join(logger.log_dir, name + '.pt')
-        )
+        self._checkpointer.init(self.model, os.path.join(logger.log_dir, name + ".pt"))
 
     def __call__(self, *inputs):
-        '''
+        """
         Run a forward pass of the model.
-        '''
+        """
         return self.model(*inputs)
 
     def no_grad(self, *inputs):
-        '''Run a forward pass of the model in no_grad mode.'''
+        """Run a forward pass of the model in no_grad mode."""
         with torch.no_grad():
             return self.model(*inputs)
 
     def eval(self, *inputs):
-        '''
+        """
         Run a forward pass of the model in eval mode with no_grad.
         The model is returned to its previous mode afer the forward pass is made.
-        '''
+        """
         with torch.no_grad():
             # check current mode
             mode = self.model.training
@@ -112,11 +112,11 @@ def eval(self, *inputs):
             return result
 
     def target(self, *inputs):
-        '''Run a forward pass of the target network.'''
+        """Run a forward pass of the target network."""
         return self._target(*inputs)
 
     def reinforce(self, loss):
-        '''
+        """
         Backpropagate the loss through the model and make an update step.
         Internally, this will perform most of the activities associated with a control loop
         in standard machine learning environments, depending on the configuration of the object:
@@ -127,14 +127,14 @@ def reinforce(self, loss):
 
         Returns:
             self: The current Approximation object
-        '''
+        """
         loss = self._loss_scaling * loss
         loss.backward()
         self.step(loss=loss)
         return self
 
     def step(self, loss=None):
-        '''
+        """
         Given that a backward pass has been made, run an optimization step.
         Internally, this will perform most of the activities associated with a control loop
         in standard machine learning environments, depending on the configuration of the object:
@@ -145,7 +145,7 @@ def step(self, loss=None):
 
         Returns:
             self: The current Approximation object
-        '''
+        """
         if loss is not None:
             self._logger.add_loss(self._name, loss.detach())
         self._clip_grad_norm()
@@ -157,22 +157,26 @@ def step(self, loss=None):
         return self
 
     def zero_grad(self):
-        '''
+        """
         Clears the gradients of all optimized tensors
 
         Returns:
             self: The current Approximation object
-        '''
+        """
         self._optimizer.zero_grad()
         return self
 
     def _clip_grad_norm(self):
-        '''Clip the gradient norm if set. Raises RuntimeError if norm is non-finite.'''
+        """Clip the gradient norm if set. Raises RuntimeError if norm is non-finite."""
         if self._clip_grad != 0:
-            utils.clip_grad_norm_(self.model.parameters(), self._clip_grad, error_if_nonfinite=True)
+            utils.clip_grad_norm_(
+                self.model.parameters(), self._clip_grad, error_if_nonfinite=True
+            )
 
     def _step_lr_scheduler(self):
-        '''Step the . Raises RuntimeError if norm is non-finite.'''
+        """Step the . Raises RuntimeError if norm is non-finite."""
         if self._scheduler:
-            self._logger.add_schedule(self._name + '/lr', self._optimizer.param_groups[0]['lr'])
+            self._logger.add_schedule(
+                self._name + "/lr", self._optimizer.param_groups[0]["lr"]
+            )
             self._scheduler.step()
diff --git a/all/approximation/checkpointer/__init__.py b/all/approximation/checkpointer/__init__.py
index 36047f9d..495eca2e 100644
--- a/all/approximation/checkpointer/__init__.py
+++ b/all/approximation/checkpointer/__init__.py
@@ -1,5 +1,6 @@
 import warnings
-from abc import abstractmethod, ABC
+from abc import ABC, abstractmethod
+
 import torch
 
 
diff --git a/all/approximation/feature_network.py b/all/approximation/feature_network.py
index 559cbbe4..af3aea02 100644
--- a/all/approximation/feature_network.py
+++ b/all/approximation/feature_network.py
@@ -1,14 +1,15 @@
 import torch
+
 from .approximation import Approximation
 
 
 class FeatureNetwork(Approximation):
-    '''
+    """
     An Approximation that accepts a state updates the observation key
     based on the given model.
-    '''
+    """
 
-    def __init__(self, model, optimizer=None, name='feature', **kwargs):
+    def __init__(self, model, optimizer=None, name="feature", **kwargs):
         model = FeatureModule(model)
         super().__init__(model, optimizer, name=name, **kwargs)
 
@@ -19,5 +20,5 @@ def __init__(self, model):
         self.model = model
 
     def forward(self, states):
-        features = states.as_output(self.model(states.as_input('observation')))
-        return states.update('observation', features)
+        features = states.as_output(self.model(states.as_input("observation")))
+        return states.update("observation", features)
diff --git a/all/approximation/identity.py b/all/approximation/identity.py
index 77353b9b..db4bb87a 100644
--- a/all/approximation/identity.py
+++ b/all/approximation/identity.py
@@ -1,16 +1,16 @@
-import torch
 from torch import nn
+
 from .approximation import Approximation
 
 
 class Identity(Approximation):
-    '''
+    """
     An Approximation that represents the identity function.
 
     Because the model has no parameters, reinforce and step do nothing.
-    '''
+    """
 
-    def __init__(self, device, name='identity', **kwargs):
+    def __init__(self, device, name="identity", **kwargs):
         super().__init__(nn.Identity(), None, device=device, name=name, **kwargs)
 
     def reinforce(self):
diff --git a/all/approximation/q_continuous.py b/all/approximation/q_continuous.py
index 29dc754d..c53cee3c 100644
--- a/all/approximation/q_continuous.py
+++ b/all/approximation/q_continuous.py
@@ -1,23 +1,14 @@
 import torch
+
 from all.nn import RLNetwork
+
 from .approximation import Approximation
 
 
 class QContinuous(Approximation):
-    def __init__(
-            self,
-            model,
-            optimizer,
-            name='q',
-            **kwargs
-    ):
+    def __init__(self, model, optimizer, name="q", **kwargs):
         model = QContinuousModule(model)
-        super().__init__(
-            model,
-            optimizer,
-            name=name,
-            **kwargs
-        )
+        super().__init__(model, optimizer, name=name, **kwargs)
 
 
 class QContinuousModule(RLNetwork):
diff --git a/all/approximation/q_dist.py b/all/approximation/q_dist.py
index eb195169..c2e9a854 100644
--- a/all/approximation/q_dist.py
+++ b/all/approximation/q_dist.py
@@ -1,20 +1,22 @@
 import torch
 from torch.nn import functional as F
+
 from all import nn
+
 from .approximation import Approximation
 
 
 class QDist(Approximation):
     def __init__(
-            self,
-            model,
-            optimizer,
-            n_actions,
-            n_atoms,
-            v_min,
-            v_max,
-            name="q_dist",
-            **kwargs
+        self,
+        model,
+        optimizer,
+        n_actions,
+        n_atoms,
+        v_min,
+        v_max,
+        name="q_dist",
+        **kwargs
     ):
         device = next(model.parameters()).device
         self.n_actions = n_actions
diff --git a/all/approximation/q_network.py b/all/approximation/q_network.py
index 6cbe999c..1da6635d 100644
--- a/all/approximation/q_network.py
+++ b/all/approximation/q_network.py
@@ -1,23 +1,14 @@
 import torch
+
 from all.nn import RLNetwork
+
 from .approximation import Approximation
 
 
 class QNetwork(Approximation):
-    def __init__(
-            self,
-            model,
-            optimizer=None,
-            name='q',
-            **kwargs
-    ):
+    def __init__(self, model, optimizer=None, name="q", **kwargs):
         model = QModule(model)
-        super().__init__(
-            model,
-            optimizer,
-            name=name,
-            **kwargs
-        )
+        super().__init__(model, optimizer, name=name, **kwargs)
 
 
 class QModule(RLNetwork):
diff --git a/all/approximation/target/__init__.py b/all/approximation/target/__init__.py
index c81dadbd..451384e4 100644
--- a/all/approximation/target/__init__.py
+++ b/all/approximation/target/__init__.py
@@ -2,3 +2,5 @@
 from .fixed import FixedTarget
 from .polyak import PolyakTarget
 from .trivial import TrivialTarget
+
+__all__ = ["TargetNetwork", "FixedTarget", "PolyakTarget", "TrivialTarget"]
diff --git a/all/approximation/target/abstract.py b/all/approximation/target/abstract.py
index 31453816..411e2685 100644
--- a/all/approximation/target/abstract.py
+++ b/all/approximation/target/abstract.py
@@ -1,4 +1,4 @@
-from abc import abstractmethod, ABC
+from abc import ABC, abstractmethod
 
 
 class TargetNetwork(ABC):
diff --git a/all/approximation/target/fixed.py b/all/approximation/target/fixed.py
index 7f6c576d..756ee996 100644
--- a/all/approximation/target/fixed.py
+++ b/all/approximation/target/fixed.py
@@ -1,5 +1,7 @@
 import copy
+
 import torch
+
 from .abstract import TargetNetwork
 
 
diff --git a/all/approximation/target/polyak.py b/all/approximation/target/polyak.py
index 2d89c1e4..ae9271ff 100644
--- a/all/approximation/target/polyak.py
+++ b/all/approximation/target/polyak.py
@@ -1,10 +1,12 @@
 import copy
+
 import torch
+
 from .abstract import TargetNetwork
 
 
 class PolyakTarget(TargetNetwork):
-    '''TargetNetwork that updates using polyak averaging'''
+    """TargetNetwork that updates using polyak averaging"""
 
     def __init__(self, rate):
         self._source = None
@@ -20,5 +22,9 @@ def init(self, model):
         self._target = copy.deepcopy(model)
 
     def update(self):
-        for target_param, source_param in zip(self._target.parameters(), self._source.parameters()):
-            target_param.data.copy_(target_param.data * (1.0 - self._rate) + source_param.data * self._rate)
+        for target_param, source_param in zip(
+            self._target.parameters(), self._source.parameters()
+        ):
+            target_param.data.copy_(
+                target_param.data * (1.0 - self._rate) + source_param.data * self._rate
+            )
diff --git a/all/approximation/target/trivial.py b/all/approximation/target/trivial.py
index 2e95b8d7..7f8cbaec 100644
--- a/all/approximation/target/trivial.py
+++ b/all/approximation/target/trivial.py
@@ -1,4 +1,5 @@
 import torch
+
 from .abstract import TargetNetwork
 
 
diff --git a/all/approximation/v_network.py b/all/approximation/v_network.py
index 6df364c1..9947e6c6 100644
--- a/all/approximation/v_network.py
+++ b/all/approximation/v_network.py
@@ -1,22 +1,12 @@
 from all.nn import RLNetwork
+
 from .approximation import Approximation
 
 
 class VNetwork(Approximation):
-    def __init__(
-            self,
-            model,
-            optimizer,
-            name='v',
-            **kwargs
-    ):
+    def __init__(self, model, optimizer, name="v", **kwargs):
         model = VModule(model)
-        super().__init__(
-            model,
-            optimizer,
-            name=name,
-            **kwargs
-        )
+        super().__init__(model, optimizer, name=name, **kwargs)
 
 
 class VModule(RLNetwork):
diff --git a/all/bodies/__init__.py b/all/bodies/__init__.py
index 52fd4ab8..926f461d 100644
--- a/all/bodies/__init__.py
+++ b/all/bodies/__init__.py
@@ -4,10 +4,4 @@
 from .time import TimeFeature
 from .vision import FrameStack
 
-__all__ = [
-    "Body",
-    "ClipRewards",
-    "DeepmindAtariBody",
-    "FrameStack",
-    "TimeFeature"
-]
+__all__ = ["Body", "ClipRewards", "DeepmindAtariBody", "FrameStack", "TimeFeature"]
diff --git a/all/bodies/atari.py b/all/bodies/atari.py
index ecd8c7ca..0a560743 100644
--- a/all/bodies/atari.py
+++ b/all/bodies/atari.py
@@ -1,11 +1,19 @@
 import torch
+
 from ._body import Body
 from .rewards import ClipRewards
 from .vision import FrameStack
 
 
 class DeepmindAtariBody(Body):
-    def __init__(self, agent, lazy_frames=False, episodic_lives=True, frame_stack=4, clip_rewards=True):
+    def __init__(
+        self,
+        agent,
+        lazy_frames=False,
+        episodic_lives=True,
+        frame_stack=4,
+        clip_rewards=True,
+    ):
         if frame_stack > 1:
             agent = FrameStack(agent, lazy=lazy_frames, size=frame_stack)
         if clip_rewards:
@@ -17,19 +25,19 @@ def __init__(self, agent, lazy_frames=False, episodic_lives=True, frame_stack=4,
 
 class EpisodicLives(Body):
     def process_state(self, state):
-        if 'life_lost' not in state:
+        if "life_lost" not in state:
             return state
 
         if len(state.shape) == 0:
-            if state['life_lost']:
-                return state.update('mask', 0.)
+            if state["life_lost"]:
+                return state.update("mask", 0.0)
             return state
 
         masks = [None] * len(state)
-        life_lost = state['life_lost']
+        life_lost = state["life_lost"]
         for i, old_mask in enumerate(state.mask):
             if life_lost[i]:
-                masks[i] = 0.
+                masks[i] = 0.0
             else:
                 masks[i] = old_mask
-        return state.update('mask', torch.tensor(masks, device=state.device))
+        return state.update("mask", torch.tensor(masks, device=state.device))
diff --git a/all/bodies/rewards.py b/all/bodies/rewards.py
index 653bb508..8feededd 100644
--- a/all/bodies/rewards.py
+++ b/all/bodies/rewards.py
@@ -1,11 +1,12 @@
-import torch
 import numpy as np
+import torch
+
 from ._body import Body
 
 
 class ClipRewards(Body):
     def process_state(self, state):
-        return state.update('reward', self._clip(state.reward))
+        return state.update("reward", self._clip(state.reward))
 
     def _clip(self, reward):
         if torch.is_tensor(reward):
diff --git a/all/bodies/time.py b/all/bodies/time.py
index a11acc21..1e818908 100644
--- a/all/bodies/time.py
+++ b/all/bodies/time.py
@@ -1,5 +1,7 @@
 import torch
+
 from all.core import StateArray
+
 from ._body import Body
 
 
@@ -13,18 +15,23 @@ def process_state(self, state):
         if isinstance(state, StateArray):
             if self.timestep is None:
                 self.timestep = torch.zeros(state.shape, device=state.device)
-            observation = torch.cat((state.observation, self.scale * self.timestep.view(-1, 1)), dim=1)
-            state = state.update('observation', observation)
+            observation = torch.cat(
+                (state.observation, self.scale * self.timestep.view(-1, 1)), dim=1
+            )
+            state = state.update("observation", observation)
             self.timestep = state.mask.float() * (self.timestep + 1)
             return state
 
         if self.timestep is None:
             self.timestep = 0
-        state.update('timestep', self.timestep)
-        observation = torch.cat((
-            state.observation,
-            torch.tensor(self.scale * self.timestep, device=state.device).view(-1)
-        ), dim=0)
-        state = state.update('observation', observation)
+        state.update("timestep", self.timestep)
+        observation = torch.cat(
+            (
+                state.observation,
+                torch.tensor(self.scale * self.timestep, device=state.device).view(-1),
+            ),
+            dim=0,
+        )
+        state = state.update("observation", observation)
         self.timestep = state.mask * (self.timestep + 1)
         return state
diff --git a/all/bodies/vision.py b/all/bodies/vision.py
index 99ffd4db..82b7e696 100644
--- a/all/bodies/vision.py
+++ b/all/bodies/vision.py
@@ -1,5 +1,7 @@
 import torch
+
 from all.core import State, StateArray
+
 from ._body import Body
 
 
@@ -19,15 +21,15 @@ def process_state(self, state):
         if self._lazy:
             return LazyState.from_state(state, self._frames, self._to_cache)
         if isinstance(state, StateArray):
-            return state.update('observation', torch.cat(self._frames, dim=1))
-        return state.update('observation', torch.cat(self._frames, dim=0))
+            return state.update("observation", torch.cat(self._frames, dim=1))
+        return state.update("observation", torch.cat(self._frames, dim=0))
 
 
 class TensorDeviceCache:
-    '''
+    """
     To efficiently implement device trasfer of lazy states, this class
     caches the transfered tensor so that it is not copied multiple times.
-    '''
+    """
 
     def __init__(self, max_size=16):
         self.max_size = max_size
@@ -54,11 +56,11 @@ class LazyState(State):
     def from_state(cls, state, frames, to_cache):
         state = LazyState(state, device=frames[0].device)
         state.to_cache = to_cache
-        state['observation'] = frames
+        state["observation"] = frames
         return state
 
     def __getitem__(self, key):
-        if key == 'observation':
+        if key == "observation":
             v = dict.__getitem__(self, key)
             if torch.is_tensor(v):
                 return v
@@ -71,7 +73,7 @@ def update(self, key, value):
             if not k == key:
                 x[k] = dict.__getitem__(self, k)
         x[key] = value
-        state = LazyState.from_state(x, x['observation'], self.to_cache)
+        state = LazyState.from_state(x, x["observation"], self.to_cache)
         return state
 
     def to(self, device):
@@ -79,12 +81,12 @@ def to(self, device):
             return self
         x = {}
         for key, value in self.items():
-            if key == 'observation':
+            if key == "observation":
                 x[key] = [self.to_cache.convert(v, device) for v in value]
                 # x[key] = [v.to(device) for v in value]#torch.cat(value,axis=0).to(device)
             elif torch.is_tensor(value):
                 x[key] = value.to(device)
             else:
                 x[key] = value
-        state = LazyState.from_state(x, x['observation'], self.to_cache)
+        state = LazyState.from_state(x, x["observation"], self.to_cache)
         return state
diff --git a/all/core/__init__.py b/all/core/__init__.py
index 540c3800..9e173a13 100644
--- a/all/core/__init__.py
+++ b/all/core/__init__.py
@@ -1,3 +1,3 @@
-from .state import State, StateArray, MultiagentState
+from .state import MultiagentState, State, StateArray
 
-__all__ = ['State', 'StateArray', 'MultiagentState']
+__all__ = ["State", "StateArray", "MultiagentState"]
diff --git a/all/core/state.py b/all/core/state.py
index 29e785ae..12681441 100644
--- a/all/core/state.py
+++ b/all/core/state.py
@@ -1,6 +1,7 @@
+import warnings
+
 import numpy as np
 import torch
-import warnings
 
 
 class State(dict):
@@ -334,7 +335,7 @@ def update(self, key, value):
     def as_input(self, key):
         value = self[key]
         return value.view(
-            (np.prod(self.shape), *value.shape[len(self.shape):])
+            (np.prod(self.shape), *value.shape[len(self.shape) :])
         ).float()
 
     def as_output(self, tensor):
diff --git a/all/environments/__init__.py b/all/environments/__init__.py
index 20ba6fc1..9b1498fa 100644
--- a/all/environments/__init__.py
+++ b/all/environments/__init__.py
@@ -10,7 +10,6 @@
 from .pybullet import PybulletEnvironment
 from .vector_env import GymVectorEnvironment
 
-
 __all__ = [
     "AtariEnvironment",
     "DuplicateEnvironment",
@@ -22,4 +21,5 @@
     "MultiagentPettingZooEnv",
     "MujocoEnvironment",
     "PybulletEnvironment",
+    "VectorEnvironment",
 ]
diff --git a/all/environments/_multiagent_environment.py b/all/environments/_multiagent_environment.py
index a409168d..3ecce51b 100644
--- a/all/environments/_multiagent_environment.py
+++ b/all/environments/_multiagent_environment.py
@@ -2,7 +2,7 @@
 
 
 class MultiagentEnvironment(ABC):
-    '''
+    """
     A multiagent reinforcement learning Environment.
 
     The Multiagent variant of the Environment object.
@@ -10,20 +10,20 @@ class MultiagentEnvironment(ABC):
     the states, the actions, the transitions between states, and the rewards given to the agent.
     Environments are often used to benchmark reinforcement learning agents,
     or to define real problems that the user hopes to solve using reinforcement learning.
-    '''
+    """
 
     @abstractmethod
     def reset(self):
-        '''
+        """
         Reset the environment and return a new initial state for the first agent.
 
         Returns
             all.core.MultiagentState: The initial state for the next episode.
-        '''
+        """
 
     @abstractmethod
     def step(self, action):
-        '''
+        """
         Apply an action for the current agent and get the multiagent state for the next agent.
 
         Parameters:
@@ -31,37 +31,37 @@ def step(self, action):
 
         Returns:
             all.core.MultiagentState: The state for the next agent.
-        '''
+        """
 
     @abstractmethod
     def render(self, **kwargs):
-        '''Render the current environment state.'''
+        """Render the current environment state."""
 
     @abstractmethod
     def close(self):
-        '''Clean up any extraneous environment objects.'''
+        """Clean up any extraneous environment objects."""
 
     @abstractmethod
     def agent_iter(self):
-        '''
+        """
         Create an iterable which that the next element is always the name of the agent whose turn it is to act.
 
         Returns:
             An Iterable over Agent strings.
-        '''
+        """
 
     @abstractmethod
     def last(self):
-        '''
+        """
         Get the MultiagentState object for the current agent.
 
         Returns:
             The all.core.MultiagentState object for the current agent.
-        '''
+        """
 
     @abstractmethod
     def is_done(self, agent):
-        '''
+        """
         Determine whether a given agent is done.
 
         Args:
@@ -69,33 +69,33 @@ def is_done(self, agent):
 
         Returns:
             A boolean representing whether the given agent is done.
-        '''
+        """
 
     @property
     def state(self):
-        '''The State for the current agent.'''
+        """The State for the current agent."""
         return self.last()
 
     @property
     @abstractmethod
     def name(self):
-        '''str: The name of the environment.'''
+        """str: The name of the environment."""
 
     @abstractmethod
     def state_space(self, agent_id):
-        '''The state space for the given agent.'''
+        """The state space for the given agent."""
 
     def observation_space(self, agent_id):
-        '''Alias for MultiagentEnvironment.state_space(agent_id).'''
+        """Alias for MultiagentEnvironment.state_space(agent_id)."""
         return self.state_space(agent_id)
 
     @abstractmethod
     def action_space(self):
-        '''The action space for the given agent.'''
+        """The action space for the given agent."""
 
     @property
     @abstractmethod
     def device(self):
-        '''
+        """
         The torch device the environment lives on.
-        '''
+        """
diff --git a/all/environments/atari.py b/all/environments/atari.py
index e8a3a050..c889c943 100644
--- a/all/environments/atari.py
+++ b/all/environments/atari.py
@@ -1,19 +1,21 @@
 import gymnasium
 import torch
+
 from all.core import State
-from .duplicate_env import DuplicateEnvironment
+
+from ._environment import Environment
 from .atari_wrappers import (
-    NoopResetEnv,
-    MaxAndSkipEnv,
     FireResetEnv,
-    WarpFrame,
     LifeLostEnv,
+    MaxAndSkipEnv,
+    NoopResetEnv,
+    WarpFrame,
 )
-from ._environment import Environment
+from .duplicate_env import DuplicateEnvironment
 
 
 class AtariEnvironment(Environment):
-    def __init__(self, name, device='cpu', **gym_make_kwargs):
+    def __init__(self, name, device="cpu", **gym_make_kwargs):
 
         # construct the environment
         env = gymnasium.make(name + "NoFrameskip-v4", **gym_make_kwargs)
@@ -37,14 +39,18 @@ def __init__(self, name, device='cpu', **gym_make_kwargs):
         self._device = device
 
     def reset(self):
-        self._state = State.from_gym(self._env.reset(), dtype=self._env.observation_space.dtype, device=self._device)
+        self._state = State.from_gym(
+            self._env.reset(),
+            dtype=self._env.observation_space.dtype,
+            device=self._device,
+        )
         return self._state
 
     def step(self, action):
         self._state = State.from_gym(
             self._env.step(self._convert(action)),
             dtype=self._env.observation_space.dtype,
-            device=self._device
+            device=self._device,
         )
         return self._state
 
@@ -58,7 +64,9 @@ def seed(self, seed):
         self._env.seed(seed)
 
     def duplicate(self, n):
-        return DuplicateEnvironment([AtariEnvironment(self._name, device=self._device) for _ in range(n)])
+        return DuplicateEnvironment(
+            [AtariEnvironment(self._name, device=self._device) for _ in range(n)]
+        )
 
     @property
     def name(self):
diff --git a/all/environments/atari_wrappers.py b/all/environments/atari_wrappers.py
index df6013ef..7f513851 100644
--- a/all/environments/atari_wrappers.py
+++ b/all/environments/atari_wrappers.py
@@ -1,31 +1,34 @@
-'''
+"""
 A subset of Atari wrappers modified from:
 https://github.com/openai/baselines/blob/master/baselines/common/atari_wrappers.py
 Other behaviors were implemented as Bodies.
-'''
-import numpy as np
+"""
+
 import os
-os.environ.setdefault('PATH', '')
-from collections import deque
-import gymnasium
-from gymnasium import spaces
+
+import numpy as np
+
+os.environ.setdefault("PATH", "")
+
 import cv2
+import gymnasium
+
 cv2.ocl.setUseOpenCL(False)
 
 
 class NoopResetEnv(gymnasium.Wrapper):
     def __init__(self, env, noop_max=30):
-        '''Sample initial states by taking random number of no-ops on reset.
+        """Sample initial states by taking random number of no-ops on reset.
         No-op is assumed to be action 0.
-        '''
+        """
         gymnasium.Wrapper.__init__(self, env)
         self.noop_max = noop_max
         self.override_num_noops = None
         self.noop_action = 0
-        assert env.unwrapped.get_action_meanings()[0] == 'NOOP'
+        assert env.unwrapped.get_action_meanings()[0] == "NOOP"
 
     def reset(self, **kwargs):
-        ''' Do no-op action for a number of steps in [1, noop_max].'''
+        """Do no-op action for a number of steps in [1, noop_max]."""
         self.env.reset(**kwargs)
         if self.override_num_noops is not None:
             noops = self.override_num_noops
@@ -45,13 +48,13 @@ def step(self, ac):
 
 class FireResetEnv(gymnasium.Wrapper):
     def __init__(self, env):
-        '''
+        """
         Take action on reset for environments that are fixed until firing.
 
         Important: This was modified to also fire on lives lost.
-        '''
+        """
         gymnasium.Wrapper.__init__(self, env)
-        assert env.unwrapped.get_action_meanings()[1] == 'FIRE'
+        assert env.unwrapped.get_action_meanings()[1] == "FIRE"
         assert len(env.unwrapped.get_action_meanings()) >= 3
         self.lives = 0
         self.was_real_done = True
@@ -85,14 +88,14 @@ def lost_life(self):
 
 class MaxAndSkipEnv(gymnasium.Wrapper):
     def __init__(self, env, skip=4):
-        '''Return only every `skip`-th frame'''
+        """Return only every `skip`-th frame"""
         gymnasium.Wrapper.__init__(self, env)
         # most recent raw observations (for max pooling across time steps)
         self._obs_buffer = np.zeros((2,) + env.observation_space.shape, dtype=np.uint8)
         self._skip = skip
 
     def step(self, action):
-        '''Repeat action, sum reward, and max over last observations.'''
+        """Repeat action, sum reward, and max over last observations."""
         total_reward = 0.0
         for i in range(self._skip):
             obs, reward, terminated, truncated, info = self.env.step(action)
@@ -115,11 +118,11 @@ def reset(self, **kwargs):
 
 class WarpFrame(gymnasium.ObservationWrapper):
     def __init__(self, env, width=84, height=84, grayscale=True, dict_space_key=None):
-        '''
+        """
         Warp frames to 84x84 as done in the Nature paper and later work.
         If the environment uses dictionary observations, `dict_space_key` can be specified which indicates which
         observation should be warped.
-        '''
+        """
         super().__init__(env)
         self._width = width
         self._height = height
@@ -168,11 +171,11 @@ def observation(self, obs):
 
 class LifeLostEnv(gymnasium.Wrapper):
     def __init__(self, env):
-        '''
+        """
         Modified wrapper to add a "life_lost" key to info.
         This allows the agent Body to make the episode as done
         if it desires.
-        '''
+        """
         gymnasium.Wrapper.__init__(self, env)
         self.lives = 0
 
@@ -183,7 +186,7 @@ def reset(self):
     def step(self, action):
         obs, reward, terminated, truncated, _ = self.env.step(action)
         lives = self.env.unwrapped.ale.lives()
-        life_lost = (lives < self.lives and lives > 0)
+        life_lost = lives < self.lives and lives > 0
         self.lives = lives
-        info = {'life_lost': life_lost}
+        info = {"life_lost": life_lost}
         return obs, reward, terminated, truncated, info
diff --git a/all/environments/duplicate_env.py b/all/environments/duplicate_env.py
index 42fbd3b4..9b914255 100644
--- a/all/environments/duplicate_env.py
+++ b/all/environments/duplicate_env.py
@@ -1,12 +1,12 @@
-import gymnasium
 import torch
+
 from all.core import State
+
 from ._vector_environment import VectorEnvironment
-import numpy as np
 
 
 class DuplicateEnvironment(VectorEnvironment):
-    '''
+    """
     Turns a list of ALL Environment objects into a VectorEnvironment object
 
     This wrapper just takes the list of States the environments generate and outputs
@@ -16,9 +16,9 @@ class DuplicateEnvironment(VectorEnvironment):
     Args:
         envs: A list of ALL environments
         device (optional): the device on which tensors will be stored
-    '''
+    """
 
-    def __init__(self, envs, device=torch.device('cpu')):
+    def __init__(self, envs, device=torch.device("cpu")):
         self._name = envs[0].name
         self._envs = envs
         self._state = None
@@ -34,9 +34,16 @@ def name(self):
 
     def reset(self, seed=None, **kwargs):
         if seed is not None:
-            self._state = State.array([sub_env.reset(seed=(seed + i), **kwargs) for i, sub_env in enumerate(self._envs)])
+            self._state = State.array(
+                [
+                    sub_env.reset(seed=(seed + i), **kwargs)
+                    for i, sub_env in enumerate(self._envs)
+                ]
+            )
         else:
-            self._state = State.array([sub_env.reset(**kwargs) for sub_env in self._envs])
+            self._state = State.array(
+                [sub_env.reset(**kwargs) for sub_env in self._envs]
+            )
         return self._state
 
     def step(self, actions):
diff --git a/all/environments/gym.py b/all/environments/gym.py
index 7baa3c97..994fcfd7 100644
--- a/all/environments/gym.py
+++ b/all/environments/gym.py
@@ -1,13 +1,16 @@
 import gymnasium
 import torch
+
 from all.core import State
+
 from ._environment import Environment
 from .duplicate_env import DuplicateEnvironment
+
 gymnasium.logger.set_level(40)
 
 
 class GymEnvironment(Environment):
-    '''
+    """
     A wrapper for OpenAI Gym environments (see: https://gymnasium.openai.com).
 
     This wrapper converts the output of the gym environment to PyTorch tensors,
@@ -23,11 +26,19 @@ class GymEnvironment(Environment):
         device (str, optional): the device on which tensors will be stored
         legacy_gym (str, optional): If true, calls gym.make() instead of gymnasium.make()
         **gym_make_kwargs: kwargs passed to gymnasium.make(id, **gym_make_kwargs)
-    '''
-
-    def __init__(self, id, device=torch.device('cpu'), name=None, legacy_gym=False, **gym_make_kwargs):
+    """
+
+    def __init__(
+        self,
+        id,
+        device=torch.device("cpu"),
+        name=None,
+        legacy_gym=False,
+        **gym_make_kwargs
+    ):
         if legacy_gym:
             import gym
+
             self._gym = gym
         else:
             self._gym = gymnasium
@@ -46,14 +57,18 @@ def name(self):
         return self._name
 
     def reset(self, **kwargs):
-        self._state = State.from_gym(self._env.reset(**kwargs), dtype=self._env.observation_space.dtype, device=self._device)
+        self._state = State.from_gym(
+            self._env.reset(**kwargs),
+            dtype=self._env.observation_space.dtype,
+            device=self._device,
+        )
         return self._state
 
     def step(self, action):
         self._state = State.from_gym(
             self._env.step(self._convert(action)),
             dtype=self._env.observation_space.dtype,
-            device=self._device
+            device=self._device,
         )
         return self._state
 
@@ -67,7 +82,12 @@ def seed(self, seed):
         self._env.seed(seed)
 
     def duplicate(self, n):
-        return DuplicateEnvironment([GymEnvironment(self._id, device=self.device, name=self._name) for _ in range(n)])
+        return DuplicateEnvironment(
+            [
+                GymEnvironment(self._id, device=self.device, name=self._name)
+                for _ in range(n)
+            ]
+        )
 
     @property
     def state_space(self):
diff --git a/all/environments/multiagent_atari.py b/all/environments/multiagent_atari.py
index e9a10de1..c59437bb 100644
--- a/all/environments/multiagent_atari.py
+++ b/all/environments/multiagent_atari.py
@@ -1,14 +1,10 @@
 import importlib
-import numpy as np
-import torch
-import gymnasium
-from all.core import MultiagentState
-from ._multiagent_environment import MultiagentEnvironment
+
 from .multiagent_pettingzoo import MultiagentPettingZooEnv
 
 
 class MultiagentAtariEnv(MultiagentPettingZooEnv):
-    '''
+    """
     A wrapper for PettingZoo Atari environments (see: https://www.pettingzoo.ml/atari).
 
     This wrapper converts the output of the PettingZoo environment to PyTorch tensors,
@@ -17,16 +13,18 @@ class MultiagentAtariEnv(MultiagentPettingZooEnv):
     Args:
         env_name (string): A string representing the name of the environment (e.g. pong-v1)
         device (optional): the device on which tensors will be stored
-    '''
+    """
 
-    def __init__(self, env_name, device='cuda', **pettingzoo_params):
+    def __init__(self, env_name, device="cuda", **pettingzoo_params):
         env = self._load_env(env_name, pettingzoo_params)
         super().__init__(env, name=env_name, device=device)
 
     def _load_env(self, env_name, pettingzoo_params):
-        from pettingzoo import atari
-        from supersuit import resize_v1, frame_skip_v0, reshape_v0, max_observation_v0
-        env = importlib.import_module('pettingzoo.atari.{}'.format(env_name)).env(obs_type='grayscale_image', **pettingzoo_params)
+        from supersuit import frame_skip_v0, max_observation_v0, reshape_v0, resize_v1
+
+        env = importlib.import_module("pettingzoo.atari.{}".format(env_name)).env(
+            obs_type="grayscale_image", **pettingzoo_params
+        )
         env = max_observation_v0(env, 2)
         env = frame_skip_v0(env, 4)
         env = resize_v1(env, 84, 84)
diff --git a/all/environments/multiagent_pettingzoo.py b/all/environments/multiagent_pettingzoo.py
index 3bc96619..63f09b77 100644
--- a/all/environments/multiagent_pettingzoo.py
+++ b/all/environments/multiagent_pettingzoo.py
@@ -1,14 +1,14 @@
-import importlib
-import numpy as np
-import torch
 import cloudpickle
 import gymnasium
+import torch
+
 from all.core import MultiagentState
+
 from ._multiagent_environment import MultiagentEnvironment
 
 
 class MultiagentPettingZooEnv(MultiagentEnvironment):
-    '''
+    """
     A wrapper for generael PettingZoo environments (see: https://www.pettingzoo.ml/).
 
     This wrapper converts the output of the PettingZoo environment to PyTorch tensors,
@@ -17,9 +17,9 @@ class MultiagentPettingZooEnv(MultiagentEnvironment):
     Args:
         zoo_env (AECEnv): A PettingZoo AECEnv environment (e.g. pettingzoo.mpe.simple_push_v2)
         device (optional): the device on which tensors will be stored
-    '''
+    """
 
-    def __init__(self, zoo_env, name, device='cuda'):
+    def __init__(self, zoo_env, name, device="cuda"):
         env = zoo_env
         env.reset()
         self._env = env
@@ -27,22 +27,24 @@ def __init__(self, zoo_env, name, device='cuda'):
         self._device = device
         self.agents = self._env.agents
         self.subenvs = {
-            agent: SubEnv(agent, device, self.state_space(agent), self.action_space(agent))
+            agent: SubEnv(
+                agent, device, self.state_space(agent), self.action_space(agent)
+            )
             for agent in self.agents
         }
 
-    '''
+    """
     Reset the environment and return a new initial state.
 
     Returns:
         An initial MultiagentState object.
-    '''
+    """
 
     def reset(self, **kwargs):
         self._env.reset(**kwargs)
         return self.last()
 
-    '''
+    """
     Reset the environment and return a new initial state.
 
     Args:
@@ -50,7 +52,7 @@ def reset(self, **kwargs):
 
     Returns:
         The MultiagentState object for the next agent
-    '''
+    """
 
     def step(self, action):
         if action is None:
@@ -75,12 +77,24 @@ def is_done(self, agent):
         return self._env.terminations[agent]
 
     def duplicate(self, n):
-        return [MultiagentPettingZooEnv(cloudpickle.loads(cloudpickle.dumps(self._env)), self._name, device=self.device) for _ in range(n)]
+        return [
+            MultiagentPettingZooEnv(
+                cloudpickle.loads(cloudpickle.dumps(self._env)),
+                self._name,
+                device=self.device,
+            )
+            for _ in range(n)
+        ]
 
     def last(self):
         observation, reward, terminated, truncated, info = self._env.last()
         selected_obs_space = self._env.observation_space(self._env.agent_selection)
-        return MultiagentState.from_zoo(self._env.agent_selection, (observation, reward, terminated, truncated, info), device=self._device, dtype=selected_obs_space.dtype)
+        return MultiagentState.from_zoo(
+            self._env.agent_selection,
+            (observation, reward, terminated, truncated, info),
+            device=self._device,
+            dtype=selected_obs_space.dtype,
+        )
 
     @property
     def name(self):
@@ -112,7 +126,7 @@ def _convert(self, action):
         return action
 
 
-class SubEnv():
+class SubEnv:
     def __init__(self, name, device, state_space, action_space):
         self.name = name
         self.device = device
diff --git a/all/environments/pybullet.py b/all/environments/pybullet.py
index 70f379fb..48afe200 100644
--- a/all/environments/pybullet.py
+++ b/all/environments/pybullet.py
@@ -7,11 +7,13 @@ class PybulletEnvironment(GymEnvironment):
         "cheetah": "HalfCheetahBulletEnv-v0",
         "humanoid": "HumanoidBulletEnv-v0",
         "hopper": "HopperBulletEnv-v0",
-        "walker": "Walker2DBulletEnv-v0"
+        "walker": "Walker2DBulletEnv-v0",
     }
 
     def __init__(self, name, **kwargs):
-        import pybullet_envs
+        #
+        import pybullet_envs  # noqa: F401
+
         if name in self.short_names:
             name = self.short_names[name]
         super().__init__(name, legacy_gym=True, **kwargs)
diff --git a/all/environments/vector_env.py b/all/environments/vector_env.py
index d74a46af..cd5cf12b 100644
--- a/all/environments/vector_env.py
+++ b/all/environments/vector_env.py
@@ -1,12 +1,13 @@
-import gymnasium
+import numpy as np
 import torch
+
 from all.core import StateArray
+
 from ._vector_environment import VectorEnvironment
-import numpy as np
 
 
 class GymVectorEnvironment(VectorEnvironment):
-    '''
+    """
     A wrapper for Gym's vector environments
     (see: https://github.com/openai/gym/blob/master/gym/vector/vector_env.py).
 
@@ -19,9 +20,9 @@ class GymVectorEnvironment(VectorEnvironment):
     Args:
         vec_env: An OpenAI gym vector environment
         device (optional): the device on which tensors will be stored
-    '''
+    """
 
-    def __init__(self, vec_env, name, device=torch.device('cpu')):
+    def __init__(self, vec_env, name, device=torch.device("cpu")):
         self._name = name
         self._env = vec_env
         self._state = None
@@ -37,7 +38,13 @@ def name(self):
 
     def reset(self, **kwargs):
         obs, info = self._env.reset(**kwargs)
-        self._state = self._to_state(obs, np.zeros(self._env.num_envs), np.zeros(self._env.num_envs), np.zeros(self._env.num_envs), info)
+        self._state = self._to_state(
+            obs,
+            np.zeros(self._env.num_envs),
+            np.zeros(self._env.num_envs),
+            np.zeros(self._env.num_envs),
+            info,
+        )
         return self._state
 
     def _to_state(self, obs, rew, terminated, truncated, info):
@@ -45,12 +52,15 @@ def _to_state(self, obs, rew, terminated, truncated, info):
         rew = rew.astype("float32")
         done = (terminated + truncated).astype("bool")
         mask = (1 - terminated).astype("float32")
-        return StateArray({
-            "observation": torch.tensor(obs, device=self._device),
-            "reward": torch.tensor(rew, device=self._device),
-            "done": torch.tensor(done, device=self._device),
-            "mask": torch.tensor(mask, device=self._device)
-        }, shape=(self._env.num_envs,))
+        return StateArray(
+            {
+                "observation": torch.tensor(obs, device=self._device),
+                "reward": torch.tensor(rew, device=self._device),
+                "done": torch.tensor(done, device=self._device),
+                "mask": torch.tensor(mask, device=self._device),
+            },
+            shape=(self._env.num_envs,),
+        )
 
     def step(self, action):
         state_tuple = self._env.step(action.cpu().detach().numpy())
@@ -62,11 +72,17 @@ def close(self):
 
     @property
     def state_space(self):
-        return getattr(self._env, "single_observation_space", getattr(self._env, "observation_space"))
+        return getattr(
+            self._env,
+            "single_observation_space",
+            getattr(self._env, "observation_space"),
+        )
 
     @property
     def action_space(self):
-        return getattr(self._env, "single_action_space", getattr(self._env, "action_space"))
+        return getattr(
+            self._env, "single_action_space", getattr(self._env, "action_space")
+        )
 
     @property
     def state_array(self):
diff --git a/all/experiments/__init__.py b/all/experiments/__init__.py
index f0f7b530..c6d230f7 100644
--- a/all/experiments/__init__.py
+++ b/all/experiments/__init__.py
@@ -1,11 +1,11 @@
-from .run_experiment import run_experiment
 from .experiment import Experiment
-from .single_env_experiment import SingleEnvExperiment
-from .parallel_env_experiment import ParallelEnvExperiment
 from .multiagent_env_experiment import MultiagentEnvExperiment
+from .parallel_env_experiment import ParallelEnvExperiment
 from .plots import plot_returns_100
+from .run_experiment import run_experiment
+from .single_env_experiment import SingleEnvExperiment
 from .slurm import SlurmExperiment
-from .watch import watch, load_and_watch
+from .watch import load_and_watch, watch
 
 __all__ = [
     "run_experiment",
@@ -16,5 +16,5 @@
     "SlurmExperiment",
     "watch",
     "load_and_watch",
-    "plot_returns_100"
+    "plot_returns_100",
 ]
diff --git a/all/experiments/experiment.py b/all/experiments/experiment.py
index 0601da34..5611d7ae 100644
--- a/all/experiments/experiment.py
+++ b/all/experiments/experiment.py
@@ -1,16 +1,17 @@
 from abc import ABC, abstractmethod
+
 import numpy as np
 
 
 class Experiment(ABC):
-    '''
+    """
     An Experiment manages the basic train/test loop and logs results.
 
     Args:
             logger (:torch.logging.logger:): A Logger object used for logging.
             quiet (bool): If False, the Experiment will print information about
                 episode returns to standard out.
-    '''
+    """
 
     def __init__(self, logger, quiet):
         self._logger = logger
@@ -20,7 +21,7 @@ def __init__(self, logger, quiet):
 
     @abstractmethod
     def train(self, frames=np.inf, episodes=np.inf):
-        '''
+        """
         Train the agent for a certain number of frames or episodes.
         If both frames and episodes are specified, then the training loop will exit
         when either condition is satisfied.
@@ -28,11 +29,11 @@ def train(self, frames=np.inf, episodes=np.inf):
         Args:
                 frames (int): The maximum number of training frames.
                 episodes (bool): The maximum number of training episodes.
-        '''
+        """
 
     @abstractmethod
     def test(self, episodes=100):
-        '''
+        """
         Test the agent in eval mode for a certain number of episodes.
 
         Args:
@@ -40,58 +41,68 @@ def test(self, episodes=100):
 
         Returns:
             list(float): A list of all returns received during testing.
-        '''
+        """
 
     @property
     @abstractmethod
     def frame(self):
-        '''The index of the current training frame.'''
+        """The index of the current training frame."""
 
     @property
     @abstractmethod
     def episode(self):
-        '''The index of the current training episode'''
+        """The index of the current training episode"""
 
     def _log_training_episode(self, returns, episode_length, fps):
         if not self._quiet:
-            print('episode: {}, frame: {}, fps: {}, episode_length: {}, returns: {}'.format(
-                self.episode,
-                self.frame,
-                int(fps),
-                episode_length,
-                returns
-            ))
+            print(
+                "episode: {}, frame: {}, fps: {}, episode_length: {}, returns: {}".format(
+                    self.episode, self.frame, int(fps), episode_length, returns
+                )
+            )
         if returns > self._best_returns:
             self._best_returns = returns
         self._returns100.append(returns)
         if len(self._returns100) == 100:
             mean = np.mean(self._returns100)
             std = np.std(self._returns100)
-            self._logger.add_summary('returns100', mean, std, step="frame")
+            self._logger.add_summary("returns100", mean, std, step="frame")
             self._returns100 = []
-        self._logger.add_eval('returns/episode', returns, step="episode")
-        self._logger.add_eval('returns/frame', returns, step="frame")
+        self._logger.add_eval("returns/episode", returns, step="episode")
+        self._logger.add_eval("returns/frame", returns, step="frame")
         self._logger.add_eval("returns/max", self._best_returns, step="frame")
         self._logger.add_eval("episode_length", episode_length)
-        self._logger.add_eval('fps', fps, step="frame")
+        self._logger.add_eval("fps", fps, step="frame")
 
     def _log_test_episode(self, episode, returns, episode_length):
         if not self._quiet:
-            print('test episode: {}, episode_length: {}, returns: {}'.format(episode, episode_length, returns))
+            print(
+                "test episode: {}, episode_length: {}, returns: {}".format(
+                    episode, episode_length, returns
+                )
+            )
 
     def _log_test(self, returns, episode_lengths):
         if not self._quiet:
             returns_mean = np.mean(returns)
             returns_sem = np.std(returns) / np.sqrt(len(returns))
-            print('test returns (mean ± sem): {} ± {}'.format(returns_mean, returns_sem))
+            print(
+                "test returns (mean ± sem): {} ± {}".format(returns_mean, returns_sem)
+            )
             episode_length_mean = np.mean(episode_lengths)
             episode_length_sem = np.std(episode_lengths) / np.sqrt(len(episode_lengths))
-            print('test episode length (mean ± sem): {} ± {}'.format(episode_length_mean, episode_length_sem))
-        self._logger.add_summary('test_returns', np.mean(returns), np.std(returns))
-        self._logger.add_summary('test_episode_length', np.mean(episode_lengths), np.std(episode_lengths))
+            print(
+                "test episode length (mean ± sem): {} ± {}".format(
+                    episode_length_mean, episode_length_sem
+                )
+            )
+        self._logger.add_summary("test_returns", np.mean(returns), np.std(returns))
+        self._logger.add_summary(
+            "test_episode_length", np.mean(episode_lengths), np.std(episode_lengths)
+        )
 
     def save(self):
-        return self._preset.save('{}/preset.pt'.format(self._logger.log_dir))
+        return self._preset.save("{}/preset.pt".format(self._logger.log_dir))
 
     def close(self):
         self._logger.close()
diff --git a/all/experiments/multiagent_env_experiment.py b/all/experiments/multiagent_env_experiment.py
index 3314f401..e141e037 100644
--- a/all/experiments/multiagent_env_experiment.py
+++ b/all/experiments/multiagent_env_experiment.py
@@ -1,10 +1,12 @@
 from timeit import default_timer as timer
+
 import numpy as np
-from all.logging import ExperimentLogger, CometLogger
+
+from all.logging import CometLogger, ExperimentLogger
 
 
-class MultiagentEnvExperiment():
-    '''
+class MultiagentEnvExperiment:
+    """
     An Experiment object for training and testing Multiagents.
 
     Args:
@@ -17,20 +19,20 @@ class MultiagentEnvExperiment():
         save_freq (int, optional): How often to save the model to disk.
         train_steps (int, optional): The number of steps for which to train.
         verbose (bool, optional): Whether or not to log detailed information or only summaries.
-    '''
+    """
 
     def __init__(
-            self,
-            preset,
-            env,
-            logdir='runs',
-            name=None,
-            quiet=False,
-            render=False,
-            save_freq=100,
-            train_steps=float('inf'),
-            verbose=True,
-            logger="tensorboard"
+        self,
+        preset,
+        env,
+        logdir="runs",
+        name=None,
+        quiet=False,
+        render=False,
+        save_freq=100,
+        train_steps=float("inf"),
+        verbose=True,
+        logger="tensorboard",
     ):
         self._name = name if name is not None else preset.name
         self._logger = self._make_logger(logdir, self._name, env.name, verbose, logger)
@@ -47,7 +49,7 @@ def __init__(
         if render:
             self._env.render()
 
-    '''
+    """
     Train the Multiagent for a certain number of frames or episodes.
     If both frames and episodes are specified, then the training loop will exit
     when either condition is satisfied.
@@ -58,14 +60,14 @@ def __init__(
 
     Returns:
         MultiagentEnvExperiment: The experiment object.
-    '''
+    """
 
     def train(self, frames=np.inf, episodes=np.inf):
         while not self._done(frames, episodes):
             self._run_training_episode()
         return self
 
-    '''
+    """
     Test the agent in eval mode for a certain number of episodes.
 
     Args:
@@ -73,7 +75,7 @@ def train(self, frames=np.inf, episodes=np.inf):
 
     Returns:
         list(float): A list of all returns received during testing.
-    '''
+    """
 
     def test(self, episodes=100):
         test_agent = self._preset.test_agent()
@@ -90,17 +92,19 @@ def test(self, episodes=100):
         return returns
 
     def save(self):
-        return self._preset.save('{}/preset.pt'.format(self._logger.log_dir))
+        return self._preset.save("{}/preset.pt".format(self._logger.log_dir))
 
     def close(self):
         self._logger.close()
 
-    '''int: The number of completed training frames'''
+    """int: The number of completed training frames"""
+
     @property
     def frame(self):
         return self._frame
 
-    '''int: The number of completed training episodes'''
+    """int: The number of completed training episodes"""
+
     @property
     def episode(self):
         return self._episode
@@ -158,28 +162,38 @@ def _done(self, frames, episodes):
 
     def _log_training_episode(self, returns, fps):
         if not self._quiet:
-            print('returns: {}'.format(returns))
-            print('frames: {}, fps: {}'.format(self._frame, fps))
+            print("returns: {}".format(returns))
+            print("frames: {}, fps: {}".format(self._frame, fps))
         for agent in self._env.agents:
-            self._logger.add_eval('{}/returns/frame'.format(agent), returns[agent], step="frame")
+            self._logger.add_eval(
+                "{}/returns/frame".format(agent), returns[agent], step="frame"
+            )
 
     def _log_test_episode(self, episode, returns):
         if not self._quiet:
-            print('test episode: {}, returns: {}'.format(episode, returns))
+            print("test episode: {}, returns: {}".format(episode, returns))
 
     def _log_test(self, returns):
         for agent, agent_returns in returns.items():
             if not self._quiet:
                 mean = np.mean(agent_returns)
                 sem = np.std(agent_returns) / np.sqrt(len(agent_returns))
-                print('{} test returns (mean ± sem): {} ± {}'.format(agent, mean, sem))
-            self._logger.add_summary('{}/returns-test'.format(agent), np.mean(agent_returns), np.std(agent_returns))
+                print("{} test returns (mean ± sem): {} ± {}".format(agent, mean, sem))
+            self._logger.add_summary(
+                "{}/returns-test".format(agent),
+                np.mean(agent_returns),
+                np.std(agent_returns),
+            )
 
     def _save_model(self):
-        if self._save_freq != float('inf') and self._episode % self._save_freq == 0:
+        if self._save_freq != float("inf") and self._episode % self._save_freq == 0:
             self.save()
 
     def _make_logger(self, logdir, agent_name, env_name, verbose, logger):
         if logger == "comet":
-            return CometLogger(self, agent_name, env_name, verbose=verbose, logdir=logdir)
-        return ExperimentLogger(self, agent_name, env_name, verbose=verbose, logdir=logdir)
+            return CometLogger(
+                self, agent_name, env_name, verbose=verbose, logdir=logdir
+            )
+        return ExperimentLogger(
+            self, agent_name, env_name, verbose=verbose, logdir=logdir
+        )
diff --git a/all/experiments/parallel_env_experiment.py b/all/experiments/parallel_env_experiment.py
index 6cb2f926..cc462974 100644
--- a/all/experiments/parallel_env_experiment.py
+++ b/all/experiments/parallel_env_experiment.py
@@ -1,29 +1,33 @@
-
-import torch
 import time
+
 import numpy as np
-from all.logging import ExperimentLogger, CometLogger
-from .experiment import Experiment
+import torch
+
 from all.environments import VectorEnvironment
+from all.logging import CometLogger, ExperimentLogger
+
+from .experiment import Experiment
 
 
 class ParallelEnvExperiment(Experiment):
-    '''An Experiment object for training and testing agents that use parallel training environments.'''
+    """An Experiment object for training and testing agents that use parallel training environments."""
 
     def __init__(
-            self,
-            preset,
-            env,
-            name=None,
-            train_steps=float('inf'),
-            logdir='runs',
-            quiet=False,
-            render=False,
-            verbose=True,
-            logger="tensorboard"
+        self,
+        preset,
+        env,
+        name=None,
+        train_steps=float("inf"),
+        logdir="runs",
+        quiet=False,
+        render=False,
+        verbose=True,
+        logger="tensorboard",
     ):
         self._name = name if name is not None else preset.name
-        super().__init__(self._make_logger(logdir, self._name, env.name, verbose, logger), quiet)
+        super().__init__(
+            self._make_logger(logdir, self._name, env.name, verbose, logger), quiet
+        )
         self._n_envs = preset.n_envs
         if isinstance(env, VectorEnvironment):
             assert self._n_envs == env.num_envs
@@ -117,8 +121,10 @@ def test(self, episodes=100):
                         episode_length = episode_lengths[i]
                         test_returns.append(episode_return)
                         test_episode_lengths.append(episode_length)
-                        self._log_test_episode(len(test_returns), episode_return, episode_length)
-                    returns[i] = 0.
+                        self._log_test_episode(
+                            len(test_returns), episode_return, episode_length
+                        )
+                    returns[i] = 0.0
                     episode_lengths[i] = -1
                     episodes_started += 1
                     if episodes_started > episodes:
@@ -132,5 +138,9 @@ def _done(self, frames, episodes):
 
     def _make_logger(self, logdir, agent_name, env_name, verbose, logger):
         if logger == "comet":
-            return CometLogger(self, agent_name, env_name, verbose=verbose, logdir=logdir)
-        return ExperimentLogger(self, agent_name, env_name, verbose=verbose, logdir=logdir)
+            return CometLogger(
+                self, agent_name, env_name, verbose=verbose, logdir=logdir
+            )
+        return ExperimentLogger(
+            self, agent_name, env_name, verbose=verbose, logdir=logdir
+        )
diff --git a/all/experiments/plots.py b/all/experiments/plots.py
index a5f258fd..c3e69f37 100644
--- a/all/experiments/plots.py
+++ b/all/experiments/plots.py
@@ -1,6 +1,7 @@
 import os
-import numpy as np
+
 import matplotlib.pyplot as plt
+import numpy as np
 
 
 def plot_returns_100(runs_dir, timesteps=-1):
@@ -51,11 +52,11 @@ def subplot_returns_100(ax, env, data, lines, timesteps=-1):
         if agent in lines:
             ax.plot(x, mean, label=agent, color=lines[agent].get_color())
         else:
-            line, = ax.plot(x, mean, label=agent)
+            (line,) = ax.plot(x, mean, label=agent)
             lines[agent] = line
         ax.fill_between(
             x, mean + std, mean - std, alpha=0.2, color=lines[agent].get_color()
         )
         ax.set_title(env)
         ax.set_xlabel("timesteps")
-        ax.ticklabel_format(style='sci', axis='x', scilimits=(0, 5))
+        ax.ticklabel_format(style="sci", axis="x", scilimits=(0, 5))
diff --git a/all/experiments/run_experiment.py b/all/experiments/run_experiment.py
index 5b153519..3ed2021d 100644
--- a/all/experiments/run_experiment.py
+++ b/all/experiments/run_experiment.py
@@ -1,18 +1,19 @@
-from .single_env_experiment import SingleEnvExperiment
-from .parallel_env_experiment import ParallelEnvExperiment
 from all.presets import ParallelPreset
 
+from .parallel_env_experiment import ParallelEnvExperiment
+from .single_env_experiment import SingleEnvExperiment
+
 
 def run_experiment(
-        agents,
-        envs,
-        frames,
-        logdir='runs',
-        quiet=False,
-        render=False,
-        test_episodes=100,
-        verbose=True,
-        logger="tensorboard"
+    agents,
+    envs,
+    frames,
+    logdir="runs",
+    quiet=False,
+    render=False,
+    test_episodes=100,
+    verbose=True,
+    logger="tensorboard",
 ):
     if not isinstance(agents, list):
         agents = [agents]
@@ -32,7 +33,7 @@ def run_experiment(
                 quiet=quiet,
                 render=render,
                 verbose=verbose,
-                logger=logger
+                logger=logger,
             )
             experiment.save()
             experiment.train(frames=frames)
diff --git a/all/experiments/single_env_experiment.py b/all/experiments/single_env_experiment.py
index 49dd269a..d3ecb0ed 100644
--- a/all/experiments/single_env_experiment.py
+++ b/all/experiments/single_env_experiment.py
@@ -1,27 +1,31 @@
 from timeit import default_timer as timer
+
 import numpy as np
-from all.logging import ExperimentLogger, CometLogger
+
+from all.logging import CometLogger, ExperimentLogger
 
 from .experiment import Experiment
 
 
 class SingleEnvExperiment(Experiment):
-    '''An Experiment object for training and testing agents that interact with one environment at a time.'''
+    """An Experiment object for training and testing agents that interact with one environment at a time."""
 
     def __init__(
-            self,
-            preset,
-            env,
-            name=None,
-            train_steps=float('inf'),
-            logdir='runs',
-            quiet=False,
-            render=False,
-            verbose=True,
-            logger="tensorboard"
+        self,
+        preset,
+        env,
+        name=None,
+        train_steps=float("inf"),
+        logdir="runs",
+        quiet=False,
+        render=False,
+        verbose=True,
+        logger="tensorboard",
     ):
         self._name = name if name is not None else preset.name
-        super().__init__(self._make_logger(logdir, self._name, env.name, verbose, logger), quiet)
+        super().__init__(
+            self._make_logger(logdir, self._name, env.name, verbose, logger), quiet
+        )
         self._logdir = logdir
         self._preset = preset
         self._agent = self._preset.agent(logger=self._logger, train_steps=train_steps)
@@ -111,5 +115,9 @@ def _done(self, frames, episodes):
 
     def _make_logger(self, logdir, agent_name, env_name, verbose, logger):
         if logger == "comet":
-            return CometLogger(self, agent_name, env_name, verbose=verbose, logdir=logdir)
-        return ExperimentLogger(self, agent_name, env_name, verbose=verbose, logdir=logdir)
+            return CometLogger(
+                self, agent_name, env_name, verbose=verbose, logdir=logdir
+            )
+        return ExperimentLogger(
+            self, agent_name, env_name, verbose=verbose, logdir=logdir
+        )
diff --git a/all/experiments/slurm.py b/all/experiments/slurm.py
index f70396e0..7e4e0903 100644
--- a/all/experiments/slurm.py
+++ b/all/experiments/slurm.py
@@ -2,29 +2,27 @@
 import os
 import subprocess
 import sys
-from .run_experiment import run_experiment
 
+from .run_experiment import run_experiment
 
 # track the number of experiments created
 # in the current process
-ID = {
-    "value": 1
-}
+ID = {"value": 1}
 
 
 class SlurmExperiment:
     def __init__(
-            self,
-            agents,
-            envs,
-            frames,
-            test_episodes=100,
-            verbose=False,
-            job_name='autonomous-learning-library',
-            script_name='experiment.sh',
-            outdir='out',
-            logdir='runs',
-            sbatch_args=None,
+        self,
+        agents,
+        envs,
+        frames,
+        test_episodes=100,
+        verbose=False,
+        job_name="autonomous-learning-library",
+        script_name="experiment.sh",
+        outdir="out",
+        logdir="runs",
+        sbatch_args=None,
     ):
         if not isinstance(agents, list):
             agents = [agents]
@@ -59,12 +57,12 @@ def __init__(
             self.queue_jobs()
 
     def parse_args(self):
-        parser = argparse.ArgumentParser(description='Run an Atari benchmark.')
-        parser.add_argument('--experiment_id', type=int)
+        parser = argparse.ArgumentParser(description="Run an Atari benchmark.")
+        parser.add_argument("--experiment_id", type=int)
         self.args = parser.parse_args()
 
     def run_experiment(self):
-        task_id = int(os.environ['SLURM_ARRAY_TASK_ID'])
+        task_id = int(os.environ["SLURM_ARRAY_TASK_ID"])
         env = self.envs[int(task_id / len(self.agents))]
         agent = self.agents[task_id % len(self.agents)]
         run_experiment(
@@ -73,7 +71,7 @@ def run_experiment(self):
             self.frames,
             test_episodes=self.test_episodes,
             logdir=self.logdir,
-            verbose=self.verbose
+            verbose=self.verbose,
         )
 
     def queue_jobs(self):
@@ -82,41 +80,41 @@ def queue_jobs(self):
         self.run_sbatch_script()
 
     def create_sbatch_script(self):
-        script = open(self.script_name, 'w')
-        script.write('#!/bin/sh\n\n')
+        script = open(self.script_name, "w")
+        script.write("#!/bin/sh\n\n")
         num_experiments = len(self.envs) * len(self.agents)
 
         sbatch_args = {
-            'job-name': self.job_name,
-            'output': os.path.join(self.outdir, 'all_%A_%a.out'),
-            'error': os.path.join(self.outdir, 'all_%A_%a.err'),
-            'array': '0-' + str(num_experiments - 1),
-            'partition': '1080ti-short',
-            'ntasks': 1,
-            'mem-per-cpu': 4000,
-            'gres': 'gpu:1'
+            "job-name": self.job_name,
+            "output": os.path.join(self.outdir, "all_%A_%a.out"),
+            "error": os.path.join(self.outdir, "all_%A_%a.err"),
+            "array": "0-" + str(num_experiments - 1),
+            "partition": "1080ti-short",
+            "ntasks": 1,
+            "mem-per-cpu": 4000,
+            "gres": "gpu:1",
         }
         sbatch_args.update(self.sbatch_args)
 
         for key, value in sbatch_args.items():
-            script.write('#SBATCH --' + key + '=' + str(value) + '\n')
-        script.write('\n')
+            script.write("#SBATCH --" + key + "=" + str(value) + "\n")
+        script.write("\n")
 
-        script.write('python ' + sys.argv[0] + ' --experiment_id ' + str(self._id) + '\n')
+        script.write(
+            "python " + sys.argv[0] + " --experiment_id " + str(self._id) + "\n"
+        )
         script.close()
-        print('created sbatch script:', self.script_name)
+        print("created sbatch script:", self.script_name)
 
     def make_output_directory(self):
         try:
             os.mkdir(self.outdir)
-            print('Created output directory:', self.outdir)
+            print("Created output directory:", self.outdir)
         except FileExistsError:
-            print('Output directory already exists:', self.outdir)
+            print("Output directory already exists:", self.outdir)
 
     def run_sbatch_script(self):
         result = subprocess.run(
-            ['sbatch', self.script_name],
-            stdout=subprocess.PIPE,
-            check=True
+            ["sbatch", self.script_name], stdout=subprocess.PIPE, check=True
         )
-        print(result.stdout.decode('utf-8').rstrip())
+        print(result.stdout.decode("utf-8").rstrip())
diff --git a/all/experiments/watch.py b/all/experiments/watch.py
index b5922feb..0842448a 100644
--- a/all/experiments/watch.py
+++ b/all/experiments/watch.py
@@ -1,6 +1,7 @@
+import sys
 import time
+
 import torch
-import sys
 
 
 def watch(agent, env, fps=60, n_episodes=sys.maxsize):
@@ -15,7 +16,7 @@ def watch(agent, env, fps=60, n_episodes=sys.maxsize):
         returns += env.state.reward
 
         if env.state.done:
-            print('returns:', returns)
+            print("returns:", returns)
             env.reset()
             returns = 0
 
diff --git a/all/logging/__init__.py b/all/logging/__init__.py
index 8bd7deed..df81429c 100644
--- a/all/logging/__init__.py
+++ b/all/logging/__init__.py
@@ -1,6 +1,5 @@
 from ._logger import Logger
 from .dummy import DummyLogger
-from .experiment import ExperimentLogger, CometLogger
-
+from .experiment import CometLogger, ExperimentLogger
 
 __all__ = ["Logger", "DummyLogger", "ExperimentLogger", "CometLogger"]
diff --git a/all/logging/_logger.py b/all/logging/_logger.py
index 23d5ecde..728060b0 100644
--- a/all/logging/_logger.py
+++ b/all/logging/_logger.py
@@ -6,7 +6,7 @@ class Logger(ABC):
 
     @abstractmethod
     def add_summary(self, name, mean, std, step="frame"):
-        '''
+        """
         Log a summary statistic.
 
         Args:
@@ -14,54 +14,54 @@ def add_summary(self, name, mean, std, step="frame"):
             mean (float): The mean of the statistic at the current step
             std (float): The standard deviation of the statistic at the current step
             step (str, optional): Which step to use (e.g., "frame" or "episode")
-        '''
+        """
 
     @abstractmethod
     def add_loss(self, name, value, step="frame"):
-        '''
+        """
         Log the given loss metric at the current step.
 
         Args:
             name (str): The tag to associate with the loss
             value (number): The value of the loss at the current step
             step (str, optional): Which step to use (e.g., "frame" or "episode")
-        '''
+        """
 
     @abstractmethod
     def add_eval(self, name, value, step="frame"):
-        '''
+        """
         Log the given evaluation metric at the current step.
 
         Args:
             name (str): The tag to associate with the loss
             value (number): The evaluation metric at the current step
             step (str, optional): Which step to use (e.g., "frame" or "episode")
-        '''
+        """
 
     @abstractmethod
     def add_info(self, name, value, step="frame"):
-        '''
+        """
         Log the given informational metric at the current step.
 
         Args:
             name (str): The tag to associate with the loss
             value (number): The evaluation metric at the current step
             step (str, optional): Which step to use (e.g., "frame" or "episode")
-        '''
+        """
 
     @abstractmethod
     def add_schedule(self, name, value, step="frame"):
-        '''
+        """
         Log the current value of a hyperparameter according to some schedule.
 
         Args:
             name (str): The tag to associate with the hyperparameter schedule
             value (number): The value of the hyperparameter at the current step
             step (str, optional): Which step to use (e.g., "frame" or "episode")
-        '''
+        """
 
     @abstractmethod
     def close(self):
-        '''
+        """
         Close the logger and perform any necessary cleanup.
-        '''
+        """
diff --git a/all/logging/dummy.py b/all/logging/dummy.py
index 1e70d1d6..a7430105 100644
--- a/all/logging/dummy.py
+++ b/all/logging/dummy.py
@@ -2,7 +2,7 @@
 
 
 class DummyLogger(Logger):
-    '''A default Logger object that performs no logging and has no side effects.'''
+    """A default Logger object that performs no logging and has no side effects."""
 
     def add_summary(self, name, mean, std, step="frame"):
         pass
diff --git a/all/logging/experiment.py b/all/logging/experiment.py
index c2136f2d..d4fc5279 100644
--- a/all/logging/experiment.py
+++ b/all/logging/experiment.py
@@ -1,14 +1,15 @@
-
-import os
 import csv
+import os
 import subprocess
 from datetime import datetime
+
 from torch.utils.tensorboard import SummaryWriter
+
 from ._logger import Logger
 
 
 class ExperimentLogger(SummaryWriter, Logger):
-    '''
+    """
     The default Logger object used by all.experiments.Experiment.
     Writes logs using tensorboard into the current logdir directory ('runs' by default),
     tagging the run with a combination of the agent name, the commit hash of the
@@ -19,11 +20,11 @@ class ExperimentLogger(SummaryWriter, Logger):
         agent_name (str): The name of the Agent the Experiment is being performed on
         env_name (str): The name of the environment the Experiment is being performed in
         verbose (bool, optional): Whether or not to log all data or only summary metrics.
-    '''
+    """
 
-    def __init__(self, experiment, agent_name, env_name, verbose=True, logdir='runs'):
+    def __init__(self, experiment, agent_name, env_name, verbose=True, logdir="runs"):
         self.env_name = env_name
-        current_time = datetime.now().strftime('%Y-%m-%d_%H:%M:%S_%f')
+        current_time = datetime.now().strftime("%Y-%m-%d_%H:%M:%S_%f")
         dir_name = "%s_%s_%s" % (agent_name, COMMIT_HASH, current_time)
         os.makedirs(os.path.join(logdir, dir_name, env_name))
         self.log_dir = os.path.join(logdir, dir_name)
@@ -32,10 +33,16 @@ def __init__(self, experiment, agent_name, env_name, verbose=True, logdir='runs'
         super().__init__(log_dir=self.log_dir)
 
     def add_summary(self, name, mean, std, step="frame"):
-        super().add_scalar('{}/summary/{}/mean'.format(self.env_name, name), mean, self._get_step(step))
-        super().add_scalar('{}/summary/{}/std'.format(self.env_name, name), std, self._get_step(step))
+        super().add_scalar(
+            "{}/summary/{}/mean".format(self.env_name, name), mean, self._get_step(step)
+        )
+        super().add_scalar(
+            "{}/summary/{}/std".format(self.env_name, name), std, self._get_step(step)
+        )
 
-        with open(os.path.join(self.log_dir, self.env_name, name + ".csv"), "a") as csvfile:
+        with open(
+            os.path.join(self.log_dir, self.env_name, name + ".csv"), "a"
+        ) as csvfile:
             csv.writer(csvfile).writerow([self._get_step(step), mean, std])
 
     def add_loss(self, name, value, step="frame"):
@@ -66,7 +73,7 @@ def close(self):
 
 
 class CometLogger(Logger):
-    '''
+    """
     A Logger object to be used by all.experiments.Experiment.
     Writes logs using comet.ml Requires an API key to be stored in .comet.config or as an environment variable.
     Look at https://www.comet.ml/docs/python-sdk/advanced/#python-configuration for more info.
@@ -76,9 +83,9 @@ class CometLogger(Logger):
         env_name (str): The name of the environment the Experiment is being performed in
         loss (bool, optional): Whether or not to log loss/scheduling metrics, or only evaluation and summary metrics.
         logdir (str): The directory where run information is stored.
-    '''
+    """
 
-    def __init__(self, experiment, agent_name, env_name, verbose=True, logdir='runs'):
+    def __init__(self, experiment, agent_name, env_name, verbose=True, logdir="runs"):
         self.env_name = env_name
         self._experiment = experiment
         self._verbose = not verbose
@@ -86,23 +93,33 @@ def __init__(self, experiment, agent_name, env_name, verbose=True, logdir='runs'
         try:
             from comet_ml import Experiment
         except ImportError as e:
-            print("Failed to import comet_ml. CometLogger requires that comet_ml be installed")
+            print(
+                "Failed to import comet_ml. CometLogger requires that comet_ml be installed"
+            )
             raise e
         try:
             self._comet = Experiment(project_name=env_name)
         except ImportError as e:
-            print("See https://www.comet.ml/docs/python-sdk/warnings-errors/ for more info on this error.")
+            print(
+                "See https://www.comet.ml/docs/python-sdk/warnings-errors/ for more info on this error."
+            )
             raise e
         except ValueError as e:
-            print("See https://www.comet.ml/docs/python-sdk/advanced/#python-configuration for more info on this error.")
+            print(
+                "See https://www.comet.ml/docs/python-sdk/advanced/#python-configuration for more info on this error."
+            )
             raise e
 
         self._comet.set_name(agent_name)
         self.log_dir = logdir
 
     def add_summary(self, name, mean, std, step="frame"):
-        self._comet.log_metric('{}/summary/{}/mean'.format(self.env_name, name), mean, self._get_step(step))
-        self._comet.log_metric('{}/summary/{}/std'.format(self.env_name, name), std, self._get_step(step))
+        self._comet.log_metric(
+            "{}/summary/{}/mean".format(self.env_name, name), mean, self._get_step(step)
+        )
+        self._comet.log_metric(
+            "{}/summary/{}/std".format(self.env_name, name), std, self._get_step(step)
+        )
 
     def add_loss(self, name, value, step="frame"):
         self._add_scalar("loss/" + name, value, step)
@@ -137,11 +154,11 @@ def get_commit_hash():
             ["git", "rev-parse", "--short", "HEAD"],
             stdout=subprocess.PIPE,
             stderr=subprocess.DEVNULL,
-            check=False
+            check=False,
         )
         return result.stdout.decode("utf-8").rstrip()
     except Exception:
-        return ''
+        return ""
 
 
 COMMIT_HASH = get_commit_hash()
diff --git a/all/memory/__init__.py b/all/memory/__init__.py
index f9f70a9d..f6cfd6aa 100644
--- a/all/memory/__init__.py
+++ b/all/memory/__init__.py
@@ -1,11 +1,11 @@
+from .advantage import NStepAdvantageBuffer
+from .generalized_advantage import GeneralizedAdvantageBuffer
 from .replay_buffer import (
-    ReplayBuffer,
     ExperienceReplayBuffer,
-    PrioritizedReplayBuffer,
     NStepReplayBuffer,
+    PrioritizedReplayBuffer,
+    ReplayBuffer,
 )
-from .advantage import NStepAdvantageBuffer
-from .generalized_advantage import GeneralizedAdvantageBuffer
 
 __all__ = [
     "ReplayBuffer",
diff --git a/all/memory/advantage.py b/all/memory/advantage.py
index 7ab42b20..caacbada 100644
--- a/all/memory/advantage.py
+++ b/all/memory/advantage.py
@@ -1,4 +1,5 @@
 import torch
+
 from all.core import State
 
 
@@ -40,32 +41,22 @@ def advantages(self, states):
         advantages = self._compute_advantages(states, rewards, next_states, lengths)
         self._clear_buffers()
 
-        return (
-            states,
-            actions,
-            advantages
-        )
+        return (states, actions, advantages)
 
     def _compute_returns(self):
         sample_returns = torch.zeros(
-            (self.n_steps, self.n_envs),
-            device=self._rewards[0].device
+            (self.n_steps, self.n_envs), device=self._rewards[0].device
         )
         sample_lengths = torch.zeros(
-            (self.n_steps, self.n_envs),
-            device=self._rewards[0].device
+            (self.n_steps, self.n_envs), device=self._rewards[0].device
         )
         current_returns = self._rewards[0] * 0
         current_lengths = current_returns.clone()
         for i in range(self.n_steps):
             t = self.n_steps - 1 - i
             mask = self._states[t + 1].mask.float()
-            current_returns = (
-                self._rewards[t] + self.gamma * current_returns * mask
-            )
-            current_lengths = (
-                1 + current_lengths * mask
-            )
+            current_returns = self._rewards[t] + self.gamma * current_returns * mask
+            current_lengths = 1 + current_lengths * mask
             sample_returns[t] = current_returns
             sample_lengths[t] = current_lengths
 
@@ -95,7 +86,7 @@ def _summarize_transitions(self):
         return (
             State.array(sample_states),
             torch.stack(sample_actions),
-            State.array(sample_next_states)
+            State.array(sample_next_states),
         )
 
     def _compute_advantages(self, states, rewards, next_states, lengths):
diff --git a/all/memory/generalized_advantage.py b/all/memory/generalized_advantage.py
index ae5428f4..2bbbc8da 100644
--- a/all/memory/generalized_advantage.py
+++ b/all/memory/generalized_advantage.py
@@ -1,18 +1,19 @@
 import torch
+
 from all.core import State
 from all.optim import Schedulable
 
 
 class GeneralizedAdvantageBuffer(Schedulable):
     def __init__(
-            self,
-            v,
-            features,
-            n_steps,
-            n_envs,
-            discount_factor=1,
-            lam=1,
-            compute_batch_size=256,
+        self,
+        v,
+        features,
+        n_steps,
+        n_envs,
+        discount_factor=1,
+        lam=1,
+        compute_batch_size=256,
     ):
         self.v = v
         self.features = features
@@ -48,23 +49,26 @@ def advantages(self, next_states):
             raise Exception("Not enough states received!")
 
         self._states.append(next_states)
-        states = State.array(self._states[0:self.n_steps + 1])
-        actions = torch.cat(self._actions[:self.n_steps], dim=0)
-        rewards = torch.stack(self._rewards[:self.n_steps])
+        states = State.array(self._states[0 : self.n_steps + 1])
+        actions = torch.cat(self._actions[: self.n_steps], dim=0)
+        rewards = torch.stack(self._rewards[: self.n_steps])
 
-        _values = states.flatten().batch_execute(self.compute_batch_size, lambda s: self.v.target(self.features.target(s))).view(states.shape)
-        values = _values[0:self.n_steps]
+        _values = (
+            states.flatten()
+            .batch_execute(
+                self.compute_batch_size,
+                lambda s: self.v.target(self.features.target(s)),
+            )
+            .view(states.shape)
+        )
+        values = _values[0 : self.n_steps]
         next_values = _values[1:]
 
         td_errors = rewards + self.gamma * next_values - values
         advantages = self._compute_advantages(td_errors)
         self._clear_buffers()
 
-        return (
-            states[0:-1].flatten(),
-            actions,
-            advantages.view(-1)
-        )
+        return (states[0:-1].flatten(), actions, advantages.view(-1))
 
     def _compute_advantages(self, td_errors):
         advantages = td_errors.clone()
@@ -75,7 +79,9 @@ def _compute_advantages(self, td_errors):
         for i in range(self.n_steps):
             t = self.n_steps - 1 - i
             mask = self._states[t + 1].mask.float()
-            current_advantages = td_errors[t] + self.gamma * self.lam * current_advantages * mask
+            current_advantages = (
+                td_errors[t] + self.gamma * self.lam * current_advantages * mask
+            )
             advantages[t] = current_advantages
 
         return advantages
diff --git a/all/memory/replay_buffer.py b/all/memory/replay_buffer.py
index d949cace..a89c9c21 100644
--- a/all/memory/replay_buffer.py
+++ b/all/memory/replay_buffer.py
@@ -1,29 +1,32 @@
 from abc import ABC, abstractmethod
+
 import numpy as np
 import torch
+
 from all.core import State
 from all.optim import Schedulable
-from .segment_tree import SumSegmentTree, MinSegmentTree
+
+from .segment_tree import MinSegmentTree, SumSegmentTree
 
 
 class ReplayBuffer(ABC):
     @abstractmethod
     def store(self, state, action, reward, next_state):
-        '''Store the transition in the buffer'''
+        """Store the transition in the buffer"""
 
     @abstractmethod
     def sample(self, batch_size):
-        '''Sample from the stored transitions'''
+        """Sample from the stored transitions"""
 
     @abstractmethod
     def update_priorities(self, indexes, td_errors):
-        '''Update priorities based on the TD error'''
+        """Update priorities based on the TD error"""
 
 
 # Adapted from:
 # https://github.com/Shmuma/ptan/blob/master/ptan/experience.py
 class ExperienceReplayBuffer(ReplayBuffer):
-    def __init__(self, size, device='cpu', store_device=None):
+    def __init__(self, size, device="cpu", store_device=None):
         self.buffer = []
         self.capacity = int(size)
         self.pos = 0
@@ -58,7 +61,9 @@ def _reshape(self, minibatch, weights):
         if torch.is_tensor(minibatch[0][1]):
             actions = torch.stack([sample[1] for sample in minibatch]).to(self.device)
         else:
-            actions = torch.tensor([sample[1] for sample in minibatch], device=self.device)
+            actions = torch.tensor(
+                [sample[1] for sample in minibatch], device=self.device
+            )
         next_states = State.array([sample[2] for sample in minibatch]).to(self.device)
         return (states, actions, next_states.reward, next_states, weights)
 
@@ -71,13 +76,13 @@ def __iter__(self):
 
 class PrioritizedReplayBuffer(ExperienceReplayBuffer, Schedulable):
     def __init__(
-            self,
-            buffer_size,
-            alpha=0.6,
-            beta=0.4,
-            epsilon=1e-5,
-            device=torch.device('cpu'),
-            store_device=None
+        self,
+        buffer_size,
+        alpha=0.6,
+        beta=0.4,
+        epsilon=1e-5,
+        device=torch.device("cpu"),
+        store_device=None,
     ):
         super().__init__(buffer_size, device=device, store_device=store_device)
 
@@ -100,8 +105,8 @@ def store(self, state, action, next_state):
             return
         idx = self.pos
         super().store(state, action, next_state)
-        self._it_sum[idx] = self._max_priority ** self._alpha
-        self._it_min[idx] = self._max_priority ** self._alpha
+        self._it_sum[idx] = self._max_priority**self._alpha
+        self._it_min[idx] = self._max_priority**self._alpha
 
     def sample(self, batch_size):
         beta = self._beta
@@ -120,7 +125,7 @@ def sample(self, batch_size):
         try:
             samples = [self.buffer[idx] for idx in idxes]
         except IndexError as e:
-            print('index out of range: ', idxes)
+            print("index out of range: ", idxes)
             raise e
         self._cache = idxes
         return self._reshape(samples, torch.from_numpy(weights).to(self.device))
@@ -134,8 +139,8 @@ def update_priorities(self, priorities):
             assert priority > 0
             assert priority < np.inf
             assert 0 <= idx < len(self)
-            self._it_sum[idx] = priority ** self._alpha
-            self._it_min[idx] = priority ** self._alpha
+            self._it_sum[idx] = priority**self._alpha
+            self._it_min[idx] = priority**self._alpha
             self._max_priority = max(self._max_priority, priority)
 
     def _sample_proportional(self, batch_size):
@@ -150,13 +155,13 @@ def _sample_proportional(self, batch_size):
 
 
 class NStepReplayBuffer(ReplayBuffer):
-    '''Converts any ReplayBuffer into an NStepReplayBuffer'''
+    """Converts any ReplayBuffer into an NStepReplayBuffer"""
 
     def __init__(
-            self,
-            steps,
-            discount_factor,
-            buffer,
+        self,
+        steps,
+        discount_factor,
+        buffer,
     ):
         assert steps >= 1
         assert discount_factor >= 0
@@ -166,7 +171,7 @@ def __init__(
         self._states = []
         self._actions = []
         self._rewards = []
-        self._reward = 0.
+        self._reward = 0.0
 
     def store(self, state, action, next_state):
         if state is None or state.done:
@@ -175,7 +180,9 @@ def store(self, state, action, next_state):
         self._states.append(state)
         self._actions.append(action)
         self._rewards.append(next_state.reward)
-        self._reward += (self.discount_factor ** (len(self._states) - 1)) * next_state.reward
+        self._reward += (
+            self.discount_factor ** (len(self._states) - 1)
+        ) * next_state.reward
 
         if len(self._states) == self.steps:
             self._store_next(next_state)
@@ -183,12 +190,14 @@ def store(self, state, action, next_state):
         if next_state.done:
             while self._states:
                 self._store_next(next_state)
-            self._reward = 0.
+            self._reward = 0.0
 
     def _store_next(self, next_state):
-        self.buffer.store(self._states[0], self._actions[0], next_state.update('reward', self._reward))
+        self.buffer.store(
+            self._states[0], self._actions[0], next_state.update("reward", self._reward)
+        )
         self._reward = self._reward - self._rewards[0]
-        self._reward *= self.discount_factor ** -1
+        self._reward *= self.discount_factor**-1
         del self._states[0]
         del self._actions[0]
         del self._rewards[0]
diff --git a/all/memory/segment_tree.py b/all/memory/segment_tree.py
index a2313a5d..988ff907 100644
--- a/all/memory/segment_tree.py
+++ b/all/memory/segment_tree.py
@@ -27,7 +27,9 @@ def __init__(self, capacity, operation, neutral_element):
             neutral element for the operation above. eg. float('-inf')
             for max and 0 for sum.
         """
-        assert capacity > 0 and capacity & (capacity - 1) == 0, "capacity must be positive and a power of 2."
+        assert (
+            capacity > 0 and capacity & (capacity - 1) == 0
+        ), "capacity must be positive and a power of 2."
         self._capacity = capacity
         self._value = [neutral_element for _ in range(2 * capacity)]
         self._operation = operation
@@ -44,7 +46,7 @@ def _reduce_helper(self, start, end, node, node_start, node_end):
             else:
                 return self._operation(
                     self._reduce_helper(start, mid, 2 * node, node_start, mid),
-                    self._reduce_helper(mid + 1, end, 2 * node + 1, mid + 1, node_end)
+                    self._reduce_helper(mid + 1, end, 2 * node + 1, mid + 1, node_end),
                 )
 
     def reduce(self, start=0, end=None):
@@ -76,8 +78,7 @@ def __setitem__(self, idx, val):
         idx //= 2
         while idx >= 1:
             self._value[idx] = self._operation(
-                self._value[2 * idx],
-                self._value[2 * idx + 1]
+                self._value[2 * idx], self._value[2 * idx + 1]
             )
             idx //= 2
 
@@ -90,9 +91,7 @@ class SumSegmentTree(SegmentTree):
     # stolen from https://github.com/Shmuma/ptan/blob/master/ptan/common/utils.py
     def __init__(self, capacity):
         super(SumSegmentTree, self).__init__(
-            capacity=capacity,
-            operation=operator.add,
-            neutral_element=0.0
+            capacity=capacity, operation=operator.add, neutral_element=0.0
         )
 
     def sum(self, start=0, end=None):
@@ -129,9 +128,7 @@ class MinSegmentTree(SegmentTree):
     # stolen from https://github.com/Shmuma/ptan/blob/master/ptan/common/utils.py
     def __init__(self, capacity):
         super(MinSegmentTree, self).__init__(
-            capacity=capacity,
-            operation=min,
-            neutral_element=float('inf')
+            capacity=capacity, operation=min, neutral_element=float("inf")
         )
 
     def min(self, start=0, end=None):
diff --git a/all/nn/__init__.py b/all/nn/__init__.py
index c60c9a97..67dfa9db 100644
--- a/all/nn/__init__.py
+++ b/all/nn/__init__.py
@@ -1,10 +1,8 @@
+import numpy as np
 import torch
 from torch import nn
 from torch.nn import *  # noqa
 from torch.nn import functional as F
-import numpy as np
-from all.core import State
-
 
 """"A Pytorch Module"""
 Module = nn.Module
@@ -21,7 +19,7 @@ def __init__(self, model, _=None):
         self.device = next(model.parameters()).device
 
     def forward(self, state):
-        return state.apply(self.model, 'observation')
+        return state.apply(self.model, "observation")
 
 
 class Aggregation(nn.Module):
@@ -139,7 +137,9 @@ class NoisyFactorizedLinear(nn.Linear):
     N.B. nn.Linear already initializes weight and bias to
     """
 
-    def __init__(self, in_features, out_features, sigma_init=0.4, init_scale=3, bias=True):
+    def __init__(
+        self, in_features, out_features, sigma_init=0.4, init_scale=3, bias=True
+    ):
         self.init_scale = init_scale
         super().__init__(in_features, out_features, bias=bias)
         sigma_init = sigma_init / np.sqrt(in_features)
@@ -149,9 +149,7 @@ def __init__(self, in_features, out_features, sigma_init=0.4, init_scale=3, bias
         self.register_buffer("epsilon_input", torch.zeros(1, in_features))
         self.register_buffer("epsilon_output", torch.zeros(out_features, 1))
         if bias:
-            self.sigma_bias = nn.Parameter(
-                torch.Tensor(out_features).fill_(sigma_init)
-            )
+            self.sigma_bias = nn.Parameter(torch.Tensor(out_features).fill_(sigma_init))
 
     def reset_parameters(self):
         std = np.sqrt(self.init_scale / self.in_features)
@@ -218,13 +216,13 @@ def _loss(estimates, errors):
     return _loss
 
 
-def weighted_mse_loss(input, target, weight, reduction='mean'):
-    loss = (weight * ((target - input) ** 2))
-    return torch.mean(loss) if reduction == 'mean' else torch.sum(loss)
+def weighted_mse_loss(input, target, weight, reduction="mean"):
+    loss = weight * ((target - input) ** 2)
+    return torch.mean(loss) if reduction == "mean" else torch.sum(loss)
 
 
-def weighted_smooth_l1_loss(input, target, weight, reduction='mean'):
+def weighted_smooth_l1_loss(input, target, weight, reduction="mean"):
     t = torch.abs(input - target)
-    loss = torch.where(t < 1, 0.5 * t ** 2, t - 0.5)
+    loss = torch.where(t < 1, 0.5 * t**2, t - 0.5)
     loss = weight * loss
-    return torch.mean(loss) if reduction == 'mean' else torch.sum(loss)
+    return torch.mean(loss) if reduction == "mean" else torch.sum(loss)
diff --git a/all/optim/__init__.py b/all/optim/__init__.py
index c2a3058a..694126c7 100644
--- a/all/optim/__init__.py
+++ b/all/optim/__init__.py
@@ -1,3 +1,3 @@
 from .scheduler import LinearScheduler, Schedulable
 
-__all__ = ['Schedulable', 'LinearScheduler']
+__all__ = ["Schedulable", "LinearScheduler"]
diff --git a/all/optim/scheduler.py b/all/optim/scheduler.py
index d79c5d77..6fe5eae1 100644
--- a/all/optim/scheduler.py
+++ b/all/optim/scheduler.py
@@ -2,7 +2,7 @@
 
 
 class Schedulable:
-    '''Allow "instance" descriptors to implement parameter scheduling.'''
+    """Allow "instance" descriptors to implement parameter scheduling."""
 
     def __getattribute__(self, name):
         value = object.__getattribute__(self, name)
@@ -17,13 +17,13 @@ class Scheduler:
 
 class LinearScheduler(Scheduler):
     def __init__(
-            self,
-            initial_value,
-            final_value,
-            decay_start,
-            decay_end,
-            name='variable',
-            logger=DummyLogger(),
+        self,
+        initial_value,
+        final_value,
+        decay_start,
+        decay_end,
+        name="variable",
+        logger=DummyLogger(),
     ):
         self._initial_value = initial_value
         self._final_value = final_value
diff --git a/all/policies/__init__.py b/all/policies/__init__.py
index 1d264856..a92e7fac 100644
--- a/all/policies/__init__.py
+++ b/all/policies/__init__.py
@@ -1,8 +1,8 @@
+from .deterministic import DeterministicPolicy
 from .gaussian import GaussianPolicy
 from .greedy import GreedyPolicy, ParallelGreedyPolicy
-from .softmax import SoftmaxPolicy
-from .deterministic import DeterministicPolicy
 from .soft_deterministic import SoftDeterministicPolicy
+from .softmax import SoftmaxPolicy
 
 __all__ = [
     "GaussianPolicy",
@@ -10,5 +10,5 @@
     "ParallelGreedyPolicy",
     "SoftmaxPolicy",
     "DeterministicPolicy",
-    "SoftDeterministicPolicy"
+    "SoftDeterministicPolicy",
 ]
diff --git a/all/policies/deterministic.py b/all/policies/deterministic.py
index e69253ec..8a740ebe 100644
--- a/all/policies/deterministic.py
+++ b/all/policies/deterministic.py
@@ -1,10 +1,11 @@
 import torch
+
 from all.approximation import Approximation
 from all.nn import RLNetwork
 
 
 class DeterministicPolicy(Approximation):
-    '''
+    """
     A DDPG-style deterministic policy.
 
     Args:
@@ -15,23 +16,11 @@ class DeterministicPolicy(Approximation):
             model parameters, e.g. SGD, Adam, RMSprop, etc.
         action_space (gymnasium.spaces.Box): The Box representing the action space.
         kwargs (optional): Any other arguments accepted by all.approximation.Approximation
-    '''
-
-    def __init__(
-            self,
-            model,
-            optimizer=None,
-            space=None,
-            name='policy',
-            **kwargs
-    ):
+    """
+
+    def __init__(self, model, optimizer=None, space=None, name="policy", **kwargs):
         model = DeterministicPolicyNetwork(model, space)
-        super().__init__(
-            model,
-            optimizer,
-            name=name,
-            **kwargs
-        )
+        super().__init__(model, optimizer, name=name, **kwargs)
 
 
 class DeterministicPolicyNetwork(RLNetwork):
diff --git a/all/policies/gaussian.py b/all/policies/gaussian.py
index eaf76c67..85a68095 100644
--- a/all/policies/gaussian.py
+++ b/all/policies/gaussian.py
@@ -1,12 +1,13 @@
 import torch
 from torch.distributions.independent import Independent
 from torch.distributions.normal import Normal
+
 from all.approximation import Approximation
 from all.nn import RLNetwork
 
 
 class GaussianPolicy(Approximation):
-    '''
+    """
     A Gaussian stochastic policy.
 
     This policy will choose actions from a distribution represented by a spherical Gaussian.
@@ -24,21 +25,11 @@ class GaussianPolicy(Approximation):
             model parameters, e.g. SGD, Adam, RMSprop, etc.
         action_space (gymnasium.spaces.Box): The Box representing the action space.
         kwargs (optional): Any other arguments accepted by all.approximation.Approximation
-    '''
-
-    def __init__(
-            self,
-            model,
-            optimizer=None,
-            space=None,
-            name='policy',
-            **kwargs
-    ):
+    """
+
+    def __init__(self, model, optimizer=None, space=None, name="policy", **kwargs):
         super().__init__(
-            GaussianPolicyNetwork(model, space),
-            optimizer,
-            name=name,
-            **kwargs
+            GaussianPolicyNetwork(model, space), optimizer, name=name, **kwargs
         )
 
 
diff --git a/all/policies/greedy.py b/all/policies/greedy.py
index d24ed7c8..cc8e168a 100644
--- a/all/policies/greedy.py
+++ b/all/policies/greedy.py
@@ -1,10 +1,11 @@
 import numpy as np
 import torch
+
 from all.optim import Schedulable
 
 
 class GreedyPolicy(Schedulable):
-    '''
+    """
     An  "epsilon-greedy" action selection policy for discrete action spaces.
 
     This policy will usually choose the optimal action according to an approximation
@@ -16,13 +17,13 @@ class GreedyPolicy(Schedulable):
         q (all.approximation.QNetwork): The action-value or "q-function"
         num_actions (int): The number of available actions.
         epsilon (float, optional): The probability of selecting a random action.
-    '''
+    """
 
     def __init__(
-            self,
-            q,
-            num_actions,
-            epsilon=0.,
+        self,
+        q,
+        num_actions,
+        epsilon=0.0,
     ):
         self.q = q
         self.num_actions = num_actions
@@ -45,7 +46,7 @@ def eval(self, state):
 
 
 class ParallelGreedyPolicy(Schedulable):
-    '''
+    """
     A parallel version of the "epsilon-greedy" action selection policy for discrete action spaces.
 
     This policy will usually choose the optimal action according to an approximation
@@ -57,13 +58,13 @@ class ParallelGreedyPolicy(Schedulable):
         q (all.approximation.QNetwork): The action-value or "q-function"
         num_actions (int): The number of available actions.
         epsilon (float, optional): The probability of selecting a random action.
-    '''
+    """
 
     def __init__(
-            self,
-            q,
-            num_actions,
-            epsilon=0.,
+        self,
+        q,
+        num_actions,
+        epsilon=0.0,
     ):
         self.q = q
         self.num_actions = num_actions
@@ -80,6 +81,10 @@ def eval(self, state):
 
     def _choose_action(self, action_values):
         best_actions = torch.argmax(action_values, dim=-1)
-        random_actions = torch.randint(0, self.num_actions, best_actions.shape, device=best_actions.device)
-        choices = (torch.rand(best_actions.shape, device=best_actions.device) < self.epsilon).int()
+        random_actions = torch.randint(
+            0, self.num_actions, best_actions.shape, device=best_actions.device
+        )
+        choices = (
+            torch.rand(best_actions.shape, device=best_actions.device) < self.epsilon
+        ).int()
         return choices * random_actions + (1 - choices) * best_actions
diff --git a/all/policies/soft_deterministic.py b/all/policies/soft_deterministic.py
index ffb8a0d0..b3075e3e 100644
--- a/all/policies/soft_deterministic.py
+++ b/all/policies/soft_deterministic.py
@@ -1,10 +1,11 @@
 import torch
+
 from all.approximation import Approximation
 from all.nn import RLNetwork
 
 
 class SoftDeterministicPolicy(Approximation):
-    '''
+    """
     A "soft" deterministic policy compatible with soft actor-critic (SAC).
 
     Args:
@@ -17,16 +18,9 @@ class SoftDeterministicPolicy(Approximation):
             model parameters, e.g. SGD, Adam, RMSprop, etc.
         action_space (gymnasium.spaces.Box): The Box representing the action space.
         kwargs (optional): Any other arguments accepted by all.approximation.Approximation
-    '''
+    """
 
-    def __init__(
-            self,
-            model,
-            optimizer=None,
-            space=None,
-            name="policy",
-            **kwargs
-    ):
+    def __init__(self, model, optimizer=None, space=None, name="policy", **kwargs):
         model = SoftDeterministicPolicyNetwork(model, space)
         self._inner_model = model
         super().__init__(model, optimizer, name=name, **kwargs)
@@ -48,8 +42,8 @@ def forward(self, state):
         return self._squash(normal.loc)
 
     def _normal(self, outputs):
-        means = outputs[..., 0:self._action_dim]
-        logvars = outputs[..., self._action_dim:]
+        means = outputs[..., 0 : self._action_dim]
+        logvars = outputs[..., self._action_dim :]
         std = logvars.mul(0.5).exp_()
         return torch.distributions.normal.Normal(means, std)
 
@@ -59,7 +53,7 @@ def _sample(self, normal):
         return self._squash(raw), log_prob
 
     def _log_prob(self, normal, raw):
-        '''
+        """
         Compute the log probability of a raw action after the action is squashed.
         Both inputs act on the raw underlying distribution.
         Because tanh_mean does not affect the density, we can ignore it.
@@ -72,7 +66,7 @@ def _log_prob(self, normal, raw):
 
         Returns:
             torch.Tensor: The probability of the raw action, accounting for the affects of tanh.
-        '''
+        """
         log_prob = normal.log_prob(raw)
         log_prob -= torch.log(1 - torch.tanh(raw).pow(2) + 1e-6)
         log_prob -= torch.log(self._tanh_scale)
diff --git a/all/policies/softmax.py b/all/policies/softmax.py
index 431333fe..f011d5b9 100644
--- a/all/policies/softmax.py
+++ b/all/policies/softmax.py
@@ -1,11 +1,12 @@
 import torch
 from torch.nn import functional
-from all.nn import RLNetwork
+
 from all.approximation import Approximation
+from all.nn import RLNetwork
 
 
 class SoftmaxPolicy(Approximation):
-    '''
+    """
     A softmax (or Boltzmann) stochastic policy for discrete actions.
 
     Args:
@@ -15,15 +16,9 @@ class SoftmaxPolicy(Approximation):
         optimizer (torch.optim.Optimizer): A optimizer initialized with the
             model parameters, e.g. SGD, Adam, RMSprop, etc.
         kwargs (optional): Any other arguments accepted by all.approximation.Approximation
-    '''
-
-    def __init__(
-            self,
-            model,
-            optimizer=None,
-            name='policy',
-            **kwargs
-    ):
+    """
+
+    def __init__(self, model, optimizer=None, name="policy", **kwargs):
         model = SoftmaxPolicyNetwork(model)
         super().__init__(model, optimizer, name=name, **kwargs)
 
diff --git a/all/presets/__init__.py b/all/presets/__init__.py
index 82d875eb..be333fb1 100644
--- a/all/presets/__init__.py
+++ b/all/presets/__init__.py
@@ -1,9 +1,8 @@
-from all.presets import atari
-from all.presets import classic_control
-from all.presets import continuous
-from .preset import Preset, ParallelPreset
-from .builder import PresetBuilder, ParallelPresetBuilder
+from all.presets import atari, classic_control, continuous
+
+from .builder import ParallelPresetBuilder, PresetBuilder
 from .independent_multiagent import IndependentMultiagentPreset
+from .preset import ParallelPreset, Preset
 
 __all__ = [
     "Preset",
@@ -13,5 +12,5 @@
     "atari",
     "classic_control",
     "continuous",
-    "IndependentMultiagentPreset"
+    "IndependentMultiagentPreset",
 ]
diff --git a/all/presets/atari/__init__.py b/all/presets/atari/__init__.py
index 836b5657..f5b338de 100644
--- a/all/presets/atari/__init__.py
+++ b/all/presets/atari/__init__.py
@@ -1,14 +1,13 @@
-from .a2c import a2c, A2CAtariPreset
-from .c51 import c51, C51AtariPreset
-from .dqn import dqn, DQNAtariPreset
-from .ddqn import ddqn, DDQNAtariPreset
-from .ppo import ppo, PPOAtariPreset
-from .rainbow import rainbow, RainbowAtariPreset
-from .vac import vac, VACAtariPreset
-from .vpg import vpg, VPGAtariPreset
-from .vqn import vqn, VQNAtariPreset
-from .vsarsa import vsarsa, VSarsaAtariPreset
-
+from .a2c import a2c
+from .c51 import c51
+from .ddqn import ddqn
+from .dqn import dqn
+from .ppo import ppo
+from .rainbow import rainbow
+from .vac import vac
+from .vpg import vpg
+from .vqn import vqn
+from .vsarsa import vsarsa
 
 __all__ = [
     "a2c",
diff --git a/all/presets/atari/a2c.py b/all/presets/atari/a2c.py
index 030bff3f..a61c8e5d 100644
--- a/all/presets/atari/a2c.py
+++ b/all/presets/atari/a2c.py
@@ -1,16 +1,20 @@
 import copy
-import math
+
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
+
 from all.agents import A2C, A2CTestAgent
+from all.approximation import FeatureNetwork, VNetwork
 from all.bodies import DeepmindAtariBody
-from all.approximation import VNetwork, FeatureNetwork
 from all.logging import DummyLogger
 from all.policies import SoftmaxPolicy
+from all.presets.atari.models import (
+    nature_features,
+    nature_policy_head,
+    nature_value_head,
+)
 from all.presets.builder import ParallelPresetBuilder
 from all.presets.preset import ParallelPreset
-from all.presets.atari.models import nature_features, nature_value_head, nature_policy_head
-
 
 default_hyperparameters = {
     # Common settings
@@ -28,7 +32,7 @@
     # Model construction
     "feature_model_constructor": nature_features,
     "value_model_constructor": nature_value_head,
-    "policy_model_constructor": nature_policy_head
+    "policy_model_constructor": nature_policy_head,
 }
 
 
@@ -58,23 +62,37 @@ class A2CAtariPreset(ParallelPreset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.value_model = hyperparameters['value_model_constructor']().to(device)
-        self.policy_model = hyperparameters['policy_model_constructor'](env).to(device)
-        self.feature_model = hyperparameters['feature_model_constructor']().to(device)
+        self.value_model = hyperparameters["value_model_constructor"]().to(device)
+        self.policy_model = hyperparameters["policy_model_constructor"](env).to(device)
+        self.feature_model = hyperparameters["feature_model_constructor"]().to(device)
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        n_updates = train_steps / (self.hyperparameters['n_steps'] * self.hyperparameters['n_envs'])
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        n_updates = train_steps / (
+            self.hyperparameters["n_steps"] * self.hyperparameters["n_envs"]
+        )
 
-        feature_optimizer = Adam(self.feature_model.parameters(), lr=self.hyperparameters["lr"], eps=self.hyperparameters["eps"])
-        value_optimizer = Adam(self.value_model.parameters(), lr=self.hyperparameters["lr"], eps=self.hyperparameters["eps"])
-        policy_optimizer = Adam(self.policy_model.parameters(), lr=self.hyperparameters["lr"], eps=self.hyperparameters["eps"])
+        feature_optimizer = Adam(
+            self.feature_model.parameters(),
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
+        )
+        value_optimizer = Adam(
+            self.value_model.parameters(),
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
+        )
+        policy_optimizer = Adam(
+            self.policy_model.parameters(),
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
+        )
 
         features = FeatureNetwork(
             self.feature_model,
             feature_optimizer,
             scheduler=CosineAnnealingLR(feature_optimizer, n_updates),
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         v = VNetwork(
@@ -83,7 +101,7 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
             scheduler=CosineAnnealingLR(value_optimizer, n_updates),
             loss_scaling=self.hyperparameters["value_loss_scaling"],
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         policy = SoftmaxPolicy(
@@ -91,7 +109,7 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
             policy_optimizer,
             scheduler=CosineAnnealingLR(policy_optimizer, n_updates),
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         return DeepmindAtariBody(
@@ -103,7 +121,7 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
                 n_steps=self.hyperparameters["n_steps"],
                 discount_factor=self.hyperparameters["discount_factor"],
                 entropy_loss_scaling=self.hyperparameters["entropy_loss_scaling"],
-                logger=logger
+                logger=logger,
             ),
         )
 
@@ -116,4 +134,4 @@ def parallel_test_agent(self):
         return self.test_agent()
 
 
-a2c = ParallelPresetBuilder('a2c', default_hyperparameters, A2CAtariPreset)
+a2c = ParallelPresetBuilder("a2c", default_hyperparameters, A2CAtariPreset)
diff --git a/all/presets/atari/c51.py b/all/presets/atari/c51.py
index ce1ceeb4..6ebf56b7 100644
--- a/all/presets/atari/c51.py
+++ b/all/presets/atari/c51.py
@@ -1,16 +1,17 @@
 import copy
+
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
-from all.approximation import QDist, FixedTarget
+
 from all.agents import C51, C51TestAgent
+from all.approximation import FixedTarget, QDist
 from all.bodies import DeepmindAtariBody
 from all.logging import DummyLogger
 from all.memory import ExperienceReplayBuffer
 from all.optim import LinearScheduler
+from all.presets.atari.models import nature_c51
 from all.presets.builder import PresetBuilder
 from all.presets.preset import Preset
-from all.presets.atari.models import nature_c51
-
 
 default_hyperparameters = {
     "discount_factor": 0.99,
@@ -26,7 +27,7 @@
     "replay_buffer_size": 1000000,
     # Explicit exploration
     "initial_exploration": 0.02,
-    "final_exploration": 0.,
+    "final_exploration": 0.0,
     "final_exploration_step": 250000,
     "test_exploration": 0.001,
     # Distributional RL
@@ -34,7 +35,7 @@
     "v_min": -10,
     "v_max": 10,
     # Model construction
-    "model_constructor": nature_c51
+    "model_constructor": nature_c51,
 }
 
 
@@ -70,33 +71,36 @@ class C51AtariPreset(Preset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.model = hyperparameters['model_constructor'](env, atoms=hyperparameters['atoms']).to(device)
+        self.model = hyperparameters["model_constructor"](
+            env, atoms=hyperparameters["atoms"]
+        ).to(device)
         self.n_actions = env.action_space.n
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        n_updates = (train_steps - self.hyperparameters['replay_start_size']) / self.hyperparameters['update_frequency']
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        n_updates = (
+            train_steps - self.hyperparameters["replay_start_size"]
+        ) / self.hyperparameters["update_frequency"]
 
         optimizer = Adam(
             self.model.parameters(),
-            lr=self.hyperparameters['lr'],
-            eps=self.hyperparameters['eps']
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
         )
 
         q = QDist(
             self.model,
             optimizer,
             self.n_actions,
-            self.hyperparameters['atoms'],
-            v_min=self.hyperparameters['v_min'],
-            v_max=self.hyperparameters['v_max'],
-            target=FixedTarget(self.hyperparameters['target_update_frequency']),
+            self.hyperparameters["atoms"],
+            v_min=self.hyperparameters["v_min"],
+            v_max=self.hyperparameters["v_max"],
+            target=FixedTarget(self.hyperparameters["target_update_frequency"]),
             scheduler=CosineAnnealingLR(optimizer, n_updates),
             logger=logger,
         )
 
         replay_buffer = ExperienceReplayBuffer(
-            self.hyperparameters['replay_buffer_size'],
-            device=self.device
+            self.hyperparameters["replay_buffer_size"], device=self.device
         )
 
         return DeepmindAtariBody(
@@ -104,10 +108,11 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
                 q,
                 replay_buffer,
                 exploration=LinearScheduler(
-                    self.hyperparameters['initial_exploration'],
-                    self.hyperparameters['final_exploration'],
+                    self.hyperparameters["initial_exploration"],
+                    self.hyperparameters["final_exploration"],
                     0,
-                    self.hyperparameters["final_exploration_step"] - self.hyperparameters["replay_start_size"],
+                    self.hyperparameters["final_exploration_step"]
+                    - self.hyperparameters["replay_start_size"],
                     name="epsilon",
                     logger=logger,
                 ),
@@ -115,10 +120,10 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
                 minibatch_size=self.hyperparameters["minibatch_size"],
                 replay_start_size=self.hyperparameters["replay_start_size"],
                 update_frequency=self.hyperparameters["update_frequency"],
-                logger=logger
+                logger=logger,
             ),
             lazy_frames=True,
-            episodic_lives=True
+            episodic_lives=True,
         )
 
     def test_agent(self):
@@ -126,11 +131,15 @@ def test_agent(self):
             copy.deepcopy(self.model),
             None,
             self.n_actions,
-            self.hyperparameters['atoms'],
-            v_min=self.hyperparameters['v_min'],
-            v_max=self.hyperparameters['v_max'],
+            self.hyperparameters["atoms"],
+            v_min=self.hyperparameters["v_min"],
+            v_max=self.hyperparameters["v_max"],
+        )
+        return DeepmindAtariBody(
+            C51TestAgent(
+                q_dist, self.n_actions, self.hyperparameters["test_exploration"]
+            )
         )
-        return DeepmindAtariBody(C51TestAgent(q_dist, self.n_actions, self.hyperparameters["test_exploration"]))
 
 
-c51 = PresetBuilder('c51', default_hyperparameters, C51AtariPreset)
+c51 = PresetBuilder("c51", default_hyperparameters, C51AtariPreset)
diff --git a/all/presets/atari/ddqn.py b/all/presets/atari/ddqn.py
index 1d04f887..d5754534 100644
--- a/all/presets/atari/ddqn.py
+++ b/all/presets/atari/ddqn.py
@@ -1,18 +1,19 @@
 import copy
+
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
-from all.approximation import QNetwork, FixedTarget
+
 from all.agents import DDQN, DDQNTestAgent
+from all.approximation import FixedTarget, QNetwork
 from all.bodies import DeepmindAtariBody
 from all.logging import DummyLogger
 from all.memory import PrioritizedReplayBuffer
 from all.nn import weighted_smooth_l1_loss
 from all.optim import LinearScheduler
 from all.policies import GreedyPolicy
+from all.presets.atari.models import nature_ddqn
 from all.presets.builder import PresetBuilder
 from all.presets.preset import Preset
-from all.presets.atari.models import nature_ddqn
-
 
 default_hyperparameters = {
     # Common settings
@@ -30,12 +31,12 @@
     "alpha": 0.5,
     "beta": 0.5,
     # Explicit exploration
-    "initial_exploration": 1.,
+    "initial_exploration": 1.0,
     "final_exploration": 0.01,
     "final_exploration_step": 250000,
     "test_exploration": 0.001,
     # Model construction
-    "model_constructor": nature_ddqn
+    "model_constructor": nature_ddqn,
 }
 
 
@@ -71,65 +72,69 @@ class DDQNAtariPreset(Preset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.model = hyperparameters['model_constructor'](env).to(device)
+        self.model = hyperparameters["model_constructor"](env).to(device)
         self.n_actions = env.action_space.n
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        n_updates = (train_steps - self.hyperparameters['replay_start_size']) / self.hyperparameters['update_frequency']
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        n_updates = (
+            train_steps - self.hyperparameters["replay_start_size"]
+        ) / self.hyperparameters["update_frequency"]
 
         optimizer = Adam(
             self.model.parameters(),
-            lr=self.hyperparameters['lr'],
-            eps=self.hyperparameters['eps']
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
         )
 
         q = QNetwork(
             self.model,
             optimizer,
             scheduler=CosineAnnealingLR(optimizer, n_updates),
-            target=FixedTarget(self.hyperparameters['target_update_frequency']),
-            logger=logger
+            target=FixedTarget(self.hyperparameters["target_update_frequency"]),
+            logger=logger,
         )
 
         policy = GreedyPolicy(
             q,
             self.n_actions,
             epsilon=LinearScheduler(
-                self.hyperparameters['initial_exploration'],
-                self.hyperparameters['final_exploration'],
-                self.hyperparameters['replay_start_size'],
-                self.hyperparameters['final_exploration_step'] - self.hyperparameters['replay_start_size'],
+                self.hyperparameters["initial_exploration"],
+                self.hyperparameters["final_exploration"],
+                self.hyperparameters["replay_start_size"],
+                self.hyperparameters["final_exploration_step"]
+                - self.hyperparameters["replay_start_size"],
                 name="exploration",
-                logger=logger
-            )
+                logger=logger,
+            ),
         )
 
         replay_buffer = PrioritizedReplayBuffer(
-            self.hyperparameters['replay_buffer_size'],
-            alpha=self.hyperparameters['alpha'],
-            beta=self.hyperparameters['beta'],
-            device=self.device
+            self.hyperparameters["replay_buffer_size"],
+            alpha=self.hyperparameters["alpha"],
+            beta=self.hyperparameters["beta"],
+            device=self.device,
         )
 
         return DeepmindAtariBody(
-            DDQN(q, policy, replay_buffer,
-                 loss=weighted_smooth_l1_loss,
-                 discount_factor=self.hyperparameters["discount_factor"],
-                 minibatch_size=self.hyperparameters["minibatch_size"],
-                 replay_start_size=self.hyperparameters["replay_start_size"],
-                 update_frequency=self.hyperparameters["update_frequency"],
-                 ),
-            lazy_frames=True
+            DDQN(
+                q,
+                policy,
+                replay_buffer,
+                loss=weighted_smooth_l1_loss,
+                discount_factor=self.hyperparameters["discount_factor"],
+                minibatch_size=self.hyperparameters["minibatch_size"],
+                replay_start_size=self.hyperparameters["replay_start_size"],
+                update_frequency=self.hyperparameters["update_frequency"],
+            ),
+            lazy_frames=True,
         )
 
     def test_agent(self):
         q = QNetwork(copy.deepcopy(self.model))
         policy = GreedyPolicy(
-            q,
-            self.n_actions,
-            epsilon=self.hyperparameters['test_exploration']
+            q, self.n_actions, epsilon=self.hyperparameters["test_exploration"]
         )
         return DeepmindAtariBody(DDQNTestAgent(policy))
 
 
-ddqn = PresetBuilder('ddqn', default_hyperparameters, DDQNAtariPreset)
+ddqn = PresetBuilder("ddqn", default_hyperparameters, DDQNAtariPreset)
diff --git a/all/presets/atari/dqn.py b/all/presets/atari/dqn.py
index 2d5c4d3e..7d89a821 100644
--- a/all/presets/atari/dqn.py
+++ b/all/presets/atari/dqn.py
@@ -1,21 +1,19 @@
 import copy
-import torch
-import numpy as np
+
+from torch.nn.functional import smooth_l1_loss
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
-from torch.nn.functional import smooth_l1_loss
-from all import nn
-from all.approximation import QNetwork, FixedTarget
-from all.agents import Agent, DQN, DQNTestAgent
+
+from all.agents import DQN, DQNTestAgent
+from all.approximation import FixedTarget, QNetwork
 from all.bodies import DeepmindAtariBody
 from all.logging import DummyLogger
 from all.memory import ExperienceReplayBuffer
 from all.optim import LinearScheduler
 from all.policies import GreedyPolicy
+from all.presets.atari.models import nature_dqn
 from all.presets.builder import PresetBuilder
 from all.presets.preset import Preset
-from all.presets.atari.models import nature_dqn
-
 
 default_hyperparameters = {
     # Common settings
@@ -31,12 +29,12 @@
     "replay_start_size": 80000,
     "replay_buffer_size": 1000000,
     # Explicit exploration
-    "initial_exploration": 1.,
+    "initial_exploration": 1.0,
     "final_exploration": 0.01,
     "final_exploration_step": 250000,
     "test_exploration": 0.001,
     # Model construction
-    "model_constructor": nature_dqn
+    "model_constructor": nature_dqn,
 }
 
 
@@ -69,42 +67,44 @@ class DQNAtariPreset(Preset):
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
         hyperparameters = {**default_hyperparameters, **hyperparameters}
-        self.model = hyperparameters['model_constructor'](env).to(device)
+        self.model = hyperparameters["model_constructor"](env).to(device)
         self.n_actions = env.action_space.n
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        n_updates = (train_steps - self.hyperparameters['replay_start_size']) / self.hyperparameters['update_frequency']
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        n_updates = (
+            train_steps - self.hyperparameters["replay_start_size"]
+        ) / self.hyperparameters["update_frequency"]
 
         optimizer = Adam(
             self.model.parameters(),
-            lr=self.hyperparameters['lr'],
-            eps=self.hyperparameters['eps']
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
         )
 
         q = QNetwork(
             self.model,
             optimizer,
             scheduler=CosineAnnealingLR(optimizer, n_updates),
-            target=FixedTarget(self.hyperparameters['target_update_frequency']),
-            logger=logger
+            target=FixedTarget(self.hyperparameters["target_update_frequency"]),
+            logger=logger,
         )
 
         policy = GreedyPolicy(
             q,
             self.n_actions,
             epsilon=LinearScheduler(
-                self.hyperparameters['initial_exploration'],
-                self.hyperparameters['final_exploration'],
-                self.hyperparameters['replay_start_size'],
-                self.hyperparameters['final_exploration_step'] - self.hyperparameters['replay_start_size'],
+                self.hyperparameters["initial_exploration"],
+                self.hyperparameters["final_exploration"],
+                self.hyperparameters["replay_start_size"],
+                self.hyperparameters["final_exploration_step"]
+                - self.hyperparameters["replay_start_size"],
                 name="exploration",
-                logger=logger
-            )
+                logger=logger,
+            ),
         )
 
         replay_buffer = ExperienceReplayBuffer(
-            self.hyperparameters['replay_buffer_size'],
-            device=self.device
+            self.hyperparameters["replay_buffer_size"], device=self.device
         )
 
         return DeepmindAtariBody(
@@ -112,23 +112,21 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
                 q,
                 policy,
                 replay_buffer,
-                discount_factor=self.hyperparameters['discount_factor'],
+                discount_factor=self.hyperparameters["discount_factor"],
                 loss=smooth_l1_loss,
-                minibatch_size=self.hyperparameters['minibatch_size'],
-                replay_start_size=self.hyperparameters['replay_start_size'],
-                update_frequency=self.hyperparameters['update_frequency'],
+                minibatch_size=self.hyperparameters["minibatch_size"],
+                replay_start_size=self.hyperparameters["replay_start_size"],
+                update_frequency=self.hyperparameters["update_frequency"],
             ),
-            lazy_frames=True
+            lazy_frames=True,
         )
 
     def test_agent(self):
         q = QNetwork(copy.deepcopy(self.model))
         policy = GreedyPolicy(
-            q,
-            self.n_actions,
-            epsilon=self.hyperparameters['test_exploration']
+            q, self.n_actions, epsilon=self.hyperparameters["test_exploration"]
         )
         return DeepmindAtariBody(DQNTestAgent(policy))
 
 
-dqn = PresetBuilder('dqn', default_hyperparameters, DQNAtariPreset)
+dqn = PresetBuilder("dqn", default_hyperparameters, DQNAtariPreset)
diff --git a/all/presets/atari/models/__init__.py b/all/presets/atari/models/__init__.py
index c2627408..3d410de1 100644
--- a/all/presets/atari/models/__init__.py
+++ b/all/presets/atari/models/__init__.py
@@ -13,7 +13,7 @@ def nature_dqn(env, frames=4):
         nn.Flatten(),
         nn.Linear(3136, 512),
         nn.ReLU(),
-        nn.Linear0(512, env.action_space.n)
+        nn.Linear0(512, env.action_space.n),
     )
 
 
@@ -28,17 +28,11 @@ def nature_ddqn(env, frames=4):
         nn.ReLU(),
         nn.Flatten(),
         nn.Dueling(
+            nn.Sequential(nn.Linear(3136, 512), nn.ReLU(), nn.Linear0(512, 1)),
             nn.Sequential(
-                nn.Linear(3136, 512),
-                nn.ReLU(),
-                nn.Linear0(512, 1)
-            ),
-            nn.Sequential(
-                nn.Linear(3136, 512),
-                nn.ReLU(),
-                nn.Linear0(512, env.action_space.n)
+                nn.Linear(3136, 512), nn.ReLU(), nn.Linear0(512, env.action_space.n)
             ),
-        )
+        ),
     )
 
 
@@ -77,7 +71,7 @@ def nature_c51(env, frames=4, atoms=51):
         nn.Flatten(),
         nn.Linear(3136, 512),
         nn.ReLU(),
-        nn.Linear0(512, env.action_space.n * atoms)
+        nn.Linear0(512, env.action_space.n * atoms),
     )
 
 
@@ -95,22 +89,14 @@ def nature_rainbow(env, frames=4, hidden=512, atoms=51, sigma=0.5):
             nn.Sequential(
                 nn.NoisyFactorizedLinear(3136, hidden, sigma_init=sigma),
                 nn.ReLU(),
-                nn.NoisyFactorizedLinear(
-                    hidden,
-                    atoms,
-                    init_scale=0,
-                    sigma_init=sigma
-                )
+                nn.NoisyFactorizedLinear(hidden, atoms, init_scale=0, sigma_init=sigma),
             ),
             nn.Sequential(
                 nn.NoisyFactorizedLinear(3136, hidden, sigma_init=sigma),
                 nn.ReLU(),
                 nn.NoisyFactorizedLinear(
-                    hidden,
-                    env.action_space.n * atoms,
-                    init_scale=0,
-                    sigma_init=sigma
-                )
-            )
-        )
+                    hidden, env.action_space.n * atoms, init_scale=0, sigma_init=sigma
+                ),
+            ),
+        ),
     )
diff --git a/all/presets/atari/ppo.py b/all/presets/atari/ppo.py
index 47bccb90..9b50af93 100644
--- a/all/presets/atari/ppo.py
+++ b/all/presets/atari/ppo.py
@@ -1,16 +1,21 @@
 import copy
+
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
+
 from all.agents import PPO, PPOTestAgent
+from all.approximation import FeatureNetwork, VNetwork
 from all.bodies import DeepmindAtariBody
-from all.approximation import VNetwork, FeatureNetwork
 from all.logging import DummyLogger
 from all.optim import LinearScheduler
 from all.policies import SoftmaxPolicy
+from all.presets.atari.models import (
+    nature_features,
+    nature_policy_head,
+    nature_value_head,
+)
 from all.presets.builder import ParallelPresetBuilder
 from all.presets.preset import ParallelPreset
-from all.presets.atari.models import nature_features, nature_value_head, nature_policy_head
-
 
 default_hyperparameters = {
     # Common settings
@@ -34,7 +39,7 @@
     # Model construction
     "feature_model_constructor": nature_features,
     "value_model_constructor": nature_value_head,
-    "policy_model_constructor": nature_policy_head
+    "policy_model_constructor": nature_policy_head,
 }
 
 
@@ -69,23 +74,40 @@ class PPOAtariPreset(ParallelPreset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.value_model = hyperparameters['value_model_constructor']().to(device)
-        self.policy_model = hyperparameters['policy_model_constructor'](env).to(device)
-        self.feature_model = hyperparameters['feature_model_constructor']().to(device)
-
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        n_updates = train_steps * self.hyperparameters['epochs'] * self.hyperparameters['minibatches'] / (self.hyperparameters['n_steps'] * self.hyperparameters['n_envs'])
+        self.value_model = hyperparameters["value_model_constructor"]().to(device)
+        self.policy_model = hyperparameters["policy_model_constructor"](env).to(device)
+        self.feature_model = hyperparameters["feature_model_constructor"]().to(device)
+
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        n_updates = (
+            train_steps
+            * self.hyperparameters["epochs"]
+            * self.hyperparameters["minibatches"]
+            / (self.hyperparameters["n_steps"] * self.hyperparameters["n_envs"])
+        )
 
-        feature_optimizer = Adam(self.feature_model.parameters(), lr=self.hyperparameters["lr"], eps=self.hyperparameters["eps"])
-        value_optimizer = Adam(self.value_model.parameters(), lr=self.hyperparameters["lr"], eps=self.hyperparameters["eps"])
-        policy_optimizer = Adam(self.policy_model.parameters(), lr=self.hyperparameters["lr"], eps=self.hyperparameters["eps"])
+        feature_optimizer = Adam(
+            self.feature_model.parameters(),
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
+        )
+        value_optimizer = Adam(
+            self.value_model.parameters(),
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
+        )
+        policy_optimizer = Adam(
+            self.policy_model.parameters(),
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
+        )
 
         features = FeatureNetwork(
             self.feature_model,
             feature_optimizer,
             scheduler=CosineAnnealingLR(feature_optimizer, n_updates),
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         v = VNetwork(
@@ -94,7 +116,7 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
             scheduler=CosineAnnealingLR(value_optimizer, n_updates),
             loss_scaling=self.hyperparameters["value_loss_scaling"],
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         policy = SoftmaxPolicy(
@@ -102,7 +124,7 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
             policy_optimizer,
             scheduler=CosineAnnealingLR(policy_optimizer, n_updates),
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         return DeepmindAtariBody(
@@ -115,8 +137,8 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
                     self.hyperparameters["clip_final"],
                     0,
                     n_updates,
-                    name='clip',
-                    logger=logger
+                    name="clip",
+                    logger=logger,
                 ),
                 epochs=self.hyperparameters["epochs"],
                 minibatches=self.hyperparameters["minibatches"],
@@ -138,4 +160,4 @@ def parallel_test_agent(self):
         return self.test_agent()
 
 
-ppo = ParallelPresetBuilder('ppo', default_hyperparameters, PPOAtariPreset)
+ppo = ParallelPresetBuilder("ppo", default_hyperparameters, PPOAtariPreset)
diff --git a/all/presets/atari/rainbow.py b/all/presets/atari/rainbow.py
index 1e99e2a5..e2465ad8 100644
--- a/all/presets/atari/rainbow.py
+++ b/all/presets/atari/rainbow.py
@@ -1,16 +1,17 @@
 import copy
+
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
-from all.approximation import QDist, FixedTarget
+
 from all.agents import Rainbow, RainbowTestAgent
+from all.approximation import FixedTarget, QDist
 from all.bodies import DeepmindAtariBody
 from all.logging import DummyLogger
-from all.memory import PrioritizedReplayBuffer, NStepReplayBuffer
+from all.memory import NStepReplayBuffer, PrioritizedReplayBuffer
 from all.optim import LinearScheduler
+from all.presets.atari.models import nature_rainbow
 from all.presets.builder import PresetBuilder
 from all.presets.preset import Preset
-from all.presets.atari.models import nature_rainbow
-
 
 default_hyperparameters = {
     "discount_factor": 0.99,
@@ -25,7 +26,7 @@
     "replay_buffer_size": 1000000,
     # Explicit exploration
     "initial_exploration": 0.02,
-    "final_exploration": 0.,
+    "final_exploration": 0.0,
     "test_exploration": 0.001,
     # Prioritized replay settings
     "alpha": 0.5,
@@ -39,7 +40,7 @@
     # Noisy Nets
     "sigma": 0.5,
     # Model construction
-    "model_constructor": nature_rainbow
+    "model_constructor": nature_rainbow,
 }
 
 
@@ -79,39 +80,43 @@ class RainbowAtariPreset(Preset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.model = hyperparameters['model_constructor'](env, atoms=hyperparameters["atoms"], sigma=hyperparameters["sigma"]).to(device)
+        self.model = hyperparameters["model_constructor"](
+            env, atoms=hyperparameters["atoms"], sigma=hyperparameters["sigma"]
+        ).to(device)
         self.n_actions = env.action_space.n
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        n_updates = (train_steps - self.hyperparameters['replay_start_size']) / self.hyperparameters['update_frequency']
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        n_updates = (
+            train_steps - self.hyperparameters["replay_start_size"]
+        ) / self.hyperparameters["update_frequency"]
 
         optimizer = Adam(
             self.model.parameters(),
-            lr=self.hyperparameters['lr'],
-            eps=self.hyperparameters['eps']
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
         )
 
         q_dist = QDist(
             self.model,
             optimizer,
             self.n_actions,
-            self.hyperparameters['atoms'],
+            self.hyperparameters["atoms"],
             scheduler=CosineAnnealingLR(optimizer, n_updates),
-            v_min=self.hyperparameters['v_min'],
-            v_max=self.hyperparameters['v_max'],
-            target=FixedTarget(self.hyperparameters['target_update_frequency']),
+            v_min=self.hyperparameters["v_min"],
+            v_max=self.hyperparameters["v_max"],
+            target=FixedTarget(self.hyperparameters["target_update_frequency"]),
             logger=logger,
         )
 
         replay_buffer = NStepReplayBuffer(
-            self.hyperparameters['n_steps'],
-            self.hyperparameters['discount_factor'],
+            self.hyperparameters["n_steps"],
+            self.hyperparameters["discount_factor"],
             PrioritizedReplayBuffer(
-                self.hyperparameters['replay_buffer_size'],
-                alpha=self.hyperparameters['alpha'],
-                beta=self.hyperparameters['beta'],
-                device=self.device
-            )
+                self.hyperparameters["replay_buffer_size"],
+                alpha=self.hyperparameters["alpha"],
+                beta=self.hyperparameters["beta"],
+                device=self.device,
+            ),
         )
 
         return DeepmindAtariBody(
@@ -119,21 +124,22 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
                 q_dist,
                 replay_buffer,
                 exploration=LinearScheduler(
-                    self.hyperparameters['initial_exploration'],
-                    self.hyperparameters['final_exploration'],
+                    self.hyperparameters["initial_exploration"],
+                    self.hyperparameters["final_exploration"],
                     0,
-                    train_steps - self.hyperparameters['replay_start_size'],
+                    train_steps - self.hyperparameters["replay_start_size"],
                     name="exploration",
-                    logger=logger
+                    logger=logger,
                 ),
-                discount_factor=self.hyperparameters['discount_factor'] ** self.hyperparameters["n_steps"],
-                minibatch_size=self.hyperparameters['minibatch_size'],
-                replay_start_size=self.hyperparameters['replay_start_size'],
-                update_frequency=self.hyperparameters['update_frequency'],
+                discount_factor=self.hyperparameters["discount_factor"]
+                ** self.hyperparameters["n_steps"],
+                minibatch_size=self.hyperparameters["minibatch_size"],
+                replay_start_size=self.hyperparameters["replay_start_size"],
+                update_frequency=self.hyperparameters["update_frequency"],
                 logger=logger,
             ),
             lazy_frames=True,
-            episodic_lives=True
+            episodic_lives=True,
         )
 
     def test_agent(self):
@@ -141,11 +147,15 @@ def test_agent(self):
             copy.deepcopy(self.model),
             None,
             self.n_actions,
-            self.hyperparameters['atoms'],
-            v_min=self.hyperparameters['v_min'],
-            v_max=self.hyperparameters['v_max'],
+            self.hyperparameters["atoms"],
+            v_min=self.hyperparameters["v_min"],
+            v_max=self.hyperparameters["v_max"],
+        )
+        return DeepmindAtariBody(
+            RainbowTestAgent(
+                q_dist, self.n_actions, self.hyperparameters["test_exploration"]
+            )
         )
-        return DeepmindAtariBody(RainbowTestAgent(q_dist, self.n_actions, self.hyperparameters["test_exploration"]))
 
 
-rainbow = PresetBuilder('rainbow', default_hyperparameters, RainbowAtariPreset)
+rainbow = PresetBuilder("rainbow", default_hyperparameters, RainbowAtariPreset)
diff --git a/all/presets/atari/vac.py b/all/presets/atari/vac.py
index 81d38ab5..3cb2638e 100644
--- a/all/presets/atari/vac.py
+++ b/all/presets/atari/vac.py
@@ -1,15 +1,20 @@
 import copy
+
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
+
 from all.agents import VAC, VACTestAgent
-from all.approximation import VNetwork, FeatureNetwork
+from all.approximation import FeatureNetwork, VNetwork
 from all.bodies import DeepmindAtariBody
 from all.logging import DummyLogger
 from all.policies import SoftmaxPolicy
+from all.presets.atari.models import (
+    nature_features,
+    nature_policy_head,
+    nature_value_head,
+)
 from all.presets.builder import ParallelPresetBuilder
 from all.presets.preset import ParallelPreset
-from all.presets.atari.models import nature_features, nature_value_head, nature_policy_head
-
 
 default_hyperparameters = {
     # Common settings
@@ -26,7 +31,7 @@
     # Model construction
     "feature_model_constructor": nature_features,
     "value_model_constructor": nature_value_head,
-    "policy_model_constructor": nature_policy_head
+    "policy_model_constructor": nature_policy_head,
 }
 
 
@@ -55,23 +60,35 @@ class VACAtariPreset(ParallelPreset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.value_model = hyperparameters['value_model_constructor']().to(device)
-        self.policy_model = hyperparameters['policy_model_constructor'](env).to(device)
-        self.feature_model = hyperparameters['feature_model_constructor']().to(device)
+        self.value_model = hyperparameters["value_model_constructor"]().to(device)
+        self.policy_model = hyperparameters["policy_model_constructor"](env).to(device)
+        self.feature_model = hyperparameters["feature_model_constructor"]().to(device)
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
         n_updates = train_steps / self.hyperparameters["n_envs"]
 
-        feature_optimizer = Adam(self.feature_model.parameters(), lr=self.hyperparameters["lr_pi"], eps=self.hyperparameters["eps"])
-        value_optimizer = Adam(self.value_model.parameters(), lr=self.hyperparameters["lr_v"], eps=self.hyperparameters["eps"])
-        policy_optimizer = Adam(self.policy_model.parameters(), lr=self.hyperparameters["lr_pi"], eps=self.hyperparameters["eps"])
+        feature_optimizer = Adam(
+            self.feature_model.parameters(),
+            lr=self.hyperparameters["lr_pi"],
+            eps=self.hyperparameters["eps"],
+        )
+        value_optimizer = Adam(
+            self.value_model.parameters(),
+            lr=self.hyperparameters["lr_v"],
+            eps=self.hyperparameters["eps"],
+        )
+        policy_optimizer = Adam(
+            self.policy_model.parameters(),
+            lr=self.hyperparameters["lr_pi"],
+            eps=self.hyperparameters["eps"],
+        )
 
         features = FeatureNetwork(
             self.feature_model,
             feature_optimizer,
             scheduler=CosineAnnealingLR(feature_optimizer, n_updates),
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         v = VNetwork(
@@ -80,7 +97,7 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
             scheduler=CosineAnnealingLR(value_optimizer, n_updates),
             loss_scaling=self.hyperparameters["value_loss_scaling"],
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         policy = SoftmaxPolicy(
@@ -88,11 +105,16 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
             policy_optimizer,
             scheduler=CosineAnnealingLR(policy_optimizer, n_updates),
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         return DeepmindAtariBody(
-            VAC(features, v, policy, discount_factor=self.hyperparameters["discount_factor"]),
+            VAC(
+                features,
+                v,
+                policy,
+                discount_factor=self.hyperparameters["discount_factor"],
+            ),
         )
 
     def test_agent(self):
@@ -104,4 +126,4 @@ def parallel_test_agent(self):
         return self.test_agent()
 
 
-vac = ParallelPresetBuilder('vac', default_hyperparameters, VACAtariPreset)
+vac = ParallelPresetBuilder("vac", default_hyperparameters, VACAtariPreset)
diff --git a/all/presets/atari/vpg.py b/all/presets/atari/vpg.py
index 0b6fb68e..b51a486d 100644
--- a/all/presets/atari/vpg.py
+++ b/all/presets/atari/vpg.py
@@ -1,15 +1,20 @@
 import copy
+
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
+
 from all.agents import VPG, VPGTestAgent
-from all.approximation import VNetwork, FeatureNetwork
+from all.approximation import FeatureNetwork, VNetwork
 from all.bodies import DeepmindAtariBody
 from all.logging import DummyLogger
 from all.policies import SoftmaxPolicy
+from all.presets.atari.models import (
+    nature_features,
+    nature_policy_head,
+    nature_value_head,
+)
 from all.presets.builder import PresetBuilder
 from all.presets.preset import Preset
-from all.presets.atari.models import nature_features, nature_value_head, nature_policy_head
-
 
 default_hyperparameters = {
     # Common settings
@@ -25,7 +30,7 @@
     # Model construction
     "feature_model_constructor": nature_features,
     "value_model_constructor": nature_value_head,
-    "policy_model_constructor": nature_policy_head
+    "policy_model_constructor": nature_policy_head,
 }
 
 
@@ -54,23 +59,35 @@ class VPGAtariPreset(Preset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.value_model = hyperparameters['value_model_constructor']().to(device)
-        self.policy_model = hyperparameters['policy_model_constructor'](env).to(device)
-        self.feature_model = hyperparameters['feature_model_constructor']().to(device)
+        self.value_model = hyperparameters["value_model_constructor"]().to(device)
+        self.policy_model = hyperparameters["policy_model_constructor"](env).to(device)
+        self.feature_model = hyperparameters["feature_model_constructor"]().to(device)
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
         n_updates = train_steps / self.hyperparameters["min_batch_size"]
 
-        feature_optimizer = Adam(self.feature_model.parameters(), lr=self.hyperparameters["lr_pi"], eps=self.hyperparameters["eps"])
-        value_optimizer = Adam(self.value_model.parameters(), lr=self.hyperparameters["lr_v"], eps=self.hyperparameters["eps"])
-        policy_optimizer = Adam(self.policy_model.parameters(), lr=self.hyperparameters["lr_pi"], eps=self.hyperparameters["eps"])
+        feature_optimizer = Adam(
+            self.feature_model.parameters(),
+            lr=self.hyperparameters["lr_pi"],
+            eps=self.hyperparameters["eps"],
+        )
+        value_optimizer = Adam(
+            self.value_model.parameters(),
+            lr=self.hyperparameters["lr_v"],
+            eps=self.hyperparameters["eps"],
+        )
+        policy_optimizer = Adam(
+            self.policy_model.parameters(),
+            lr=self.hyperparameters["lr_pi"],
+            eps=self.hyperparameters["eps"],
+        )
 
         features = FeatureNetwork(
             self.feature_model,
             feature_optimizer,
             scheduler=CosineAnnealingLR(feature_optimizer, n_updates),
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         v = VNetwork(
@@ -79,7 +96,7 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
             scheduler=CosineAnnealingLR(value_optimizer, n_updates),
             loss_scaling=self.hyperparameters["value_loss_scaling"],
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         policy = SoftmaxPolicy(
@@ -87,11 +104,17 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
             policy_optimizer,
             scheduler=CosineAnnealingLR(policy_optimizer, n_updates),
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         return DeepmindAtariBody(
-            VPG(features, v, policy, discount_factor=self.hyperparameters["discount_factor"], min_batch_size=self.hyperparameters["min_batch_size"]),
+            VPG(
+                features,
+                v,
+                policy,
+                discount_factor=self.hyperparameters["discount_factor"],
+                min_batch_size=self.hyperparameters["min_batch_size"],
+            ),
         )
 
     def test_agent(self):
@@ -103,4 +126,4 @@ def parallel_test_agent(self):
         return self.test_agent()
 
 
-vpg = PresetBuilder('vpg', default_hyperparameters, VPGAtariPreset)
+vpg = PresetBuilder("vpg", default_hyperparameters, VPGAtariPreset)
diff --git a/all/presets/atari/vqn.py b/all/presets/atari/vqn.py
index 4818bdae..99a2adc5 100644
--- a/all/presets/atari/vqn.py
+++ b/all/presets/atari/vqn.py
@@ -1,16 +1,17 @@
 import copy
+
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
-from all.approximation import QNetwork
+
 from all.agents import VQN, VQNTestAgent
+from all.approximation import QNetwork
 from all.bodies import DeepmindAtariBody
 from all.logging import DummyLogger
 from all.optim import LinearScheduler
 from all.policies import GreedyPolicy, ParallelGreedyPolicy
+from all.presets.atari.models import nature_ddqn
 from all.presets.builder import ParallelPresetBuilder
 from all.presets.preset import ParallelPreset
-from all.presets.atari.models import nature_ddqn
-
 
 default_hyperparameters = {
     # Common settings
@@ -19,14 +20,14 @@
     "lr": 1e-3,
     "eps": 1.5e-4,
     # Explicit exploration
-    "initial_exploration": 1.,
+    "initial_exploration": 1.0,
     "final_exploration": 0.01,
     "final_exploration_step": 250000,
     "test_exploration": 0.001,
     # Parallel actors
     "n_envs": 64,
     # Model construction
-    "model_constructor": nature_ddqn
+    "model_constructor": nature_ddqn,
 }
 
 
@@ -54,51 +55,56 @@ class VQNAtariPreset(ParallelPreset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.model = hyperparameters['model_constructor'](env).to(device)
+        self.model = hyperparameters["model_constructor"](env).to(device)
         self.n_actions = env.action_space.n
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        n_updates = train_steps / self.hyperparameters['n_envs']
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        n_updates = train_steps / self.hyperparameters["n_envs"]
 
         optimizer = Adam(
             self.model.parameters(),
-            lr=self.hyperparameters['lr'],
-            eps=self.hyperparameters['eps']
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
         )
 
         q = QNetwork(
             self.model,
             optimizer,
             scheduler=CosineAnnealingLR(optimizer, n_updates),
-            logger=logger
+            logger=logger,
         )
 
         policy = ParallelGreedyPolicy(
             q,
             self.n_actions,
             epsilon=LinearScheduler(
-                self.hyperparameters['initial_exploration'],
-                self.hyperparameters['final_exploration'],
+                self.hyperparameters["initial_exploration"],
+                self.hyperparameters["final_exploration"],
                 0,
-                self.hyperparameters["final_exploration_step"] / self.hyperparameters["n_envs"],
+                self.hyperparameters["final_exploration_step"]
+                / self.hyperparameters["n_envs"],
                 name="exploration",
-                logger=logger
-            )
+                logger=logger,
+            ),
         )
 
         return DeepmindAtariBody(
-            VQN(q, policy, discount_factor=self.hyperparameters['discount_factor']),
+            VQN(q, policy, discount_factor=self.hyperparameters["discount_factor"]),
         )
 
     def test_agent(self):
         q = QNetwork(copy.deepcopy(self.model))
-        policy = GreedyPolicy(q, self.n_actions, epsilon=self.hyperparameters['test_exploration'])
+        policy = GreedyPolicy(
+            q, self.n_actions, epsilon=self.hyperparameters["test_exploration"]
+        )
         return DeepmindAtariBody(VQNTestAgent(policy))
 
     def parallel_test_agent(self):
         q = QNetwork(copy.deepcopy(self.model))
-        policy = ParallelGreedyPolicy(q, self.n_actions, epsilon=self.hyperparameters['test_exploration'])
+        policy = ParallelGreedyPolicy(
+            q, self.n_actions, epsilon=self.hyperparameters["test_exploration"]
+        )
         return DeepmindAtariBody(VQNTestAgent(policy))
 
 
-vqn = ParallelPresetBuilder('vqn', default_hyperparameters, VQNAtariPreset)
+vqn = ParallelPresetBuilder("vqn", default_hyperparameters, VQNAtariPreset)
diff --git a/all/presets/atari/vsarsa.py b/all/presets/atari/vsarsa.py
index d41369e3..4aba9971 100644
--- a/all/presets/atari/vsarsa.py
+++ b/all/presets/atari/vsarsa.py
@@ -1,16 +1,17 @@
 import copy
+
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
-from all.approximation import QNetwork
+
 from all.agents import VSarsa, VSarsaTestAgent
+from all.approximation import QNetwork
 from all.bodies import DeepmindAtariBody
 from all.logging import DummyLogger
 from all.optim import LinearScheduler
 from all.policies import GreedyPolicy, ParallelGreedyPolicy
+from all.presets.atari.models import nature_ddqn
 from all.presets.builder import ParallelPresetBuilder
 from all.presets.preset import ParallelPreset
-from all.presets.atari.models import nature_ddqn
-
 
 default_hyperparameters = {
     # Common settings
@@ -19,14 +20,14 @@
     "lr": 1e-3,
     "eps": 1.5e-4,
     # Explicit exploration
-    "initial_exploration": 1.,
+    "initial_exploration": 1.0,
     "final_exploration": 0.01,
     "final_exploration_step": 250000,
     "test_exploration": 0.001,
     # Parallel actors
     "n_envs": 64,
     # Model construction
-    "model_constructor": nature_ddqn
+    "model_constructor": nature_ddqn,
 }
 
 
@@ -54,51 +55,56 @@ class VSarsaAtariPreset(ParallelPreset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.model = hyperparameters['model_constructor'](env).to(device)
+        self.model = hyperparameters["model_constructor"](env).to(device)
         self.n_actions = env.action_space.n
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        n_updates = train_steps / self.hyperparameters['n_envs']
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        n_updates = train_steps / self.hyperparameters["n_envs"]
 
         optimizer = Adam(
             self.model.parameters(),
-            lr=self.hyperparameters['lr'],
-            eps=self.hyperparameters['eps']
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
         )
 
         q = QNetwork(
             self.model,
             optimizer,
             scheduler=CosineAnnealingLR(optimizer, n_updates),
-            logger=logger
+            logger=logger,
         )
 
         policy = ParallelGreedyPolicy(
             q,
             self.n_actions,
             epsilon=LinearScheduler(
-                self.hyperparameters['initial_exploration'],
-                self.hyperparameters['final_exploration'],
+                self.hyperparameters["initial_exploration"],
+                self.hyperparameters["final_exploration"],
                 0,
-                self.hyperparameters["final_exploration_step"] / self.hyperparameters["n_envs"],
+                self.hyperparameters["final_exploration_step"]
+                / self.hyperparameters["n_envs"],
                 name="exploration",
-                logger=logger
-            )
+                logger=logger,
+            ),
         )
 
         return DeepmindAtariBody(
-            VSarsa(q, policy, discount_factor=self.hyperparameters['discount_factor']),
+            VSarsa(q, policy, discount_factor=self.hyperparameters["discount_factor"]),
         )
 
     def test_agent(self):
         q = QNetwork(copy.deepcopy(self.model))
-        policy = GreedyPolicy(q, self.n_actions, epsilon=self.hyperparameters['test_exploration'])
+        policy = GreedyPolicy(
+            q, self.n_actions, epsilon=self.hyperparameters["test_exploration"]
+        )
         return DeepmindAtariBody(VSarsaTestAgent(policy))
 
     def parallel_test_agent(self):
         q = QNetwork(copy.deepcopy(self.model))
-        policy = ParallelGreedyPolicy(q, self.n_actions, epsilon=self.hyperparameters['test_exploration'])
+        policy = ParallelGreedyPolicy(
+            q, self.n_actions, epsilon=self.hyperparameters["test_exploration"]
+        )
         return DeepmindAtariBody(VSarsaTestAgent(policy))
 
 
-vsarsa = ParallelPresetBuilder('vsarsa', default_hyperparameters, VSarsaAtariPreset)
+vsarsa = ParallelPresetBuilder("vsarsa", default_hyperparameters, VSarsaAtariPreset)
diff --git a/all/presets/builder.py b/all/presets/builder.py
index 768c2610..dd51e933 100644
--- a/all/presets/builder.py
+++ b/all/presets/builder.py
@@ -1,7 +1,4 @@
-from abc import ABC, abstractmethod
-
-
-class PresetBuilder():
+class PresetBuilder:
     def __init__(
         self,
         default_name,
@@ -17,7 +14,9 @@ def __init__(
         self.constructor = constructor
         self._device = device
         self._env = env
-        self._hyperparameters = self._merge_hyperparameters(default_hyperparameters, hyperparameters)
+        self._hyperparameters = self._merge_hyperparameters(
+            default_hyperparameters, hyperparameters
+        )
         self._name = name or default_name
 
     def __call__(self, **kwargs):
@@ -30,20 +29,21 @@ def env(self, env):
         return self._preset_builder(env=env)
 
     def hyperparameters(self, **hyperparameters):
-        return self._preset_builder(hyperparameters=self._merge_hyperparameters(self._hyperparameters, hyperparameters))
+        return self._preset_builder(
+            hyperparameters=self._merge_hyperparameters(
+                self._hyperparameters, hyperparameters
+            )
+        )
 
     def name(self, name):
         return self._preset_builder(name=name)
 
     def build(self):
         if not self._env:
-            raise Exception('Env is required')
+            raise Exception("Env is required")
 
         return self.constructor(
-            self._env,
-            device=self._device,
-            name=self._name,
-            **self._hyperparameters
+            self._env, device=self._device, name=self._name, **self._hyperparameters
         )
 
     def _merge_hyperparameters(self, h1, h2):
@@ -56,12 +56,17 @@ def _merge_hyperparameters(self, h1, h2):
 
     def _preset_builder(self, **kwargs):
         old_kwargs = {
-            'device': self._device,
-            'env': self._env,
-            'hyperparameters': self._hyperparameters,
-            'name': self._name,
+            "device": self._device,
+            "env": self._env,
+            "hyperparameters": self._hyperparameters,
+            "name": self._name,
         }
-        return PresetBuilder(self.default_name, self.default_hyperparameters, self.constructor, **{**old_kwargs, **kwargs})
+        return PresetBuilder(
+            self.default_name,
+            self.default_hyperparameters,
+            self.constructor,
+            **{**old_kwargs, **kwargs}
+        )
 
 
 class ParallelPresetBuilder(PresetBuilder):
@@ -75,8 +80,8 @@ def __init__(
         hyperparameters=None,
         name=None,
     ):
-        if 'n_envs' not in default_hyperparameters:
-            raise Exception('ParallelPreset hyperparameters must include n_envs')
+        if "n_envs" not in default_hyperparameters:
+            raise Exception("ParallelPreset hyperparameters must include n_envs")
         super().__init__(
             default_name,
             default_hyperparameters,
@@ -84,7 +89,7 @@ def __init__(
             device=device,
             env=env,
             hyperparameters=hyperparameters,
-            name=name
+            name=name,
         )
 
     def build(self):
diff --git a/all/presets/classic_control/__init__.py b/all/presets/classic_control/__init__.py
index 8eb72a2c..f5b338de 100644
--- a/all/presets/classic_control/__init__.py
+++ b/all/presets/classic_control/__init__.py
@@ -1,13 +1,13 @@
-from .a2c import a2c, A2CClassicControlPreset
-from .c51 import c51, C51ClassicControlPreset
-from .ddqn import ddqn, DDQNClassicControlPreset
-from .dqn import dqn, DQNClassicControlPreset
-from .ppo import ppo, PPOClassicControlPreset
-from .rainbow import rainbow, RainbowClassicControlPreset
-from .vac import vac, VACClassicControlPreset
-from .vpg import vpg, VPGClassicControlPreset
-from .vqn import vqn, VQNClassicControlPreset
-from .vsarsa import vsarsa, VSarsaClassicControlPreset
+from .a2c import a2c
+from .c51 import c51
+from .ddqn import ddqn
+from .dqn import dqn
+from .ppo import ppo
+from .rainbow import rainbow
+from .vac import vac
+from .vpg import vpg
+from .vqn import vqn
+from .vsarsa import vsarsa
 
 __all__ = [
     "a2c",
diff --git a/all/presets/classic_control/a2c.py b/all/presets/classic_control/a2c.py
index 2726e684..2551fba9 100644
--- a/all/presets/classic_control/a2c.py
+++ b/all/presets/classic_control/a2c.py
@@ -1,13 +1,18 @@
 import copy
+
 from torch.optim import Adam
+
 from all.agents import A2C, A2CTestAgent
-from all.approximation import VNetwork, FeatureNetwork
+from all.approximation import FeatureNetwork, VNetwork
 from all.logging import DummyLogger
 from all.policies import SoftmaxPolicy
 from all.presets.builder import ParallelPresetBuilder
+from all.presets.classic_control.models import (
+    fc_policy_head,
+    fc_relu_features,
+    fc_value_head,
+)
 from all.presets.preset import ParallelPreset
-from all.presets.classic_control.models import fc_relu_features, fc_policy_head, fc_value_head
-
 
 default_hyperparameters = {
     # Common settings
@@ -24,7 +29,7 @@
     # Model construction
     "feature_model_constructor": fc_relu_features,
     "value_model_constructor": fc_value_head,
-    "policy_model_constructor": fc_policy_head
+    "policy_model_constructor": fc_policy_head,
 }
 
 
@@ -54,33 +59,41 @@ class A2CClassicControlPreset(ParallelPreset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.value_model = hyperparameters['value_model_constructor']().to(device)
-        self.policy_model = hyperparameters['policy_model_constructor'](env).to(device)
-        self.feature_model = hyperparameters['feature_model_constructor'](env).to(device)
+        self.value_model = hyperparameters["value_model_constructor"]().to(device)
+        self.policy_model = hyperparameters["policy_model_constructor"](env).to(device)
+        self.feature_model = hyperparameters["feature_model_constructor"](env).to(
+            device
+        )
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        feature_optimizer = Adam(self.feature_model.parameters(), lr=self.hyperparameters["lr"])
-        value_optimizer = Adam(self.value_model.parameters(), lr=self.hyperparameters["lr"])
-        policy_optimizer = Adam(self.policy_model.parameters(), lr=self.hyperparameters["lr"])
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        feature_optimizer = Adam(
+            self.feature_model.parameters(), lr=self.hyperparameters["lr"]
+        )
+        value_optimizer = Adam(
+            self.value_model.parameters(), lr=self.hyperparameters["lr"]
+        )
+        policy_optimizer = Adam(
+            self.policy_model.parameters(), lr=self.hyperparameters["lr"]
+        )
 
         features = FeatureNetwork(
             self.feature_model,
             feature_optimizer,
-            clip_grad=self.hyperparameters["clip_grad"]
+            clip_grad=self.hyperparameters["clip_grad"],
         )
 
         v = VNetwork(
             self.value_model,
             value_optimizer,
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         policy = SoftmaxPolicy(
             self.policy_model,
             policy_optimizer,
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         return A2C(
@@ -91,7 +104,7 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
             n_steps=self.hyperparameters["n_steps"],
             discount_factor=self.hyperparameters["discount_factor"],
             entropy_loss_scaling=self.hyperparameters["entropy_loss_scaling"],
-            logger=logger
+            logger=logger,
         )
 
     def test_agent(self):
@@ -103,4 +116,4 @@ def parallel_test_agent(self):
         return self.test_agent()
 
 
-a2c = ParallelPresetBuilder('a2c', default_hyperparameters, A2CClassicControlPreset)
+a2c = ParallelPresetBuilder("a2c", default_hyperparameters, A2CClassicControlPreset)
diff --git a/all/presets/classic_control/c51.py b/all/presets/classic_control/c51.py
index da54d871..adfb3983 100644
--- a/all/presets/classic_control/c51.py
+++ b/all/presets/classic_control/c51.py
@@ -1,14 +1,15 @@
 import copy
+
 from torch.optim import Adam
+
 from all.agents import C51, C51TestAgent
-from all.approximation import QDist, FixedTarget
+from all.approximation import FixedTarget, QDist
 from all.logging import DummyLogger
 from all.memory import ExperienceReplayBuffer
 from all.optim import LinearScheduler
 from all.presets.builder import PresetBuilder
-from all.presets.preset import Preset
 from all.presets.classic_control.models import fc_relu_dist_q
-
+from all.presets.preset import Preset
 
 default_hyperparameters = {
     "discount_factor": 0.99,
@@ -31,7 +32,7 @@
     "v_min": -100,
     "v_max": 100,
     # Model construction
-    "model_constructor": fc_relu_dist_q
+    "model_constructor": fc_relu_dist_q,
 }
 
 
@@ -66,36 +67,38 @@ class C51ClassicControlPreset(Preset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.model = hyperparameters['model_constructor'](env, atoms=hyperparameters['atoms']).to(device)
+        self.model = hyperparameters["model_constructor"](
+            env, atoms=hyperparameters["atoms"]
+        ).to(device)
         self.n_actions = env.action_space.n
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        optimizer = Adam(self.model.parameters(), lr=self.hyperparameters['lr'])
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        optimizer = Adam(self.model.parameters(), lr=self.hyperparameters["lr"])
 
         q = QDist(
             self.model,
             optimizer,
             self.n_actions,
-            self.hyperparameters['atoms'],
-            v_min=self.hyperparameters['v_min'],
-            v_max=self.hyperparameters['v_max'],
-            target=FixedTarget(self.hyperparameters['target_update_frequency']),
+            self.hyperparameters["atoms"],
+            v_min=self.hyperparameters["v_min"],
+            v_max=self.hyperparameters["v_max"],
+            target=FixedTarget(self.hyperparameters["target_update_frequency"]),
             logger=logger,
         )
 
         replay_buffer = ExperienceReplayBuffer(
-            self.hyperparameters['replay_buffer_size'],
-            device=self.device
+            self.hyperparameters["replay_buffer_size"], device=self.device
         )
 
         return C51(
             q,
             replay_buffer,
             exploration=LinearScheduler(
-                self.hyperparameters['initial_exploration'],
-                self.hyperparameters['final_exploration'],
+                self.hyperparameters["initial_exploration"],
+                self.hyperparameters["final_exploration"],
                 0,
-                self.hyperparameters["final_exploration_step"] - self.hyperparameters["replay_start_size"],
+                self.hyperparameters["final_exploration_step"]
+                - self.hyperparameters["replay_start_size"],
                 name="epsilon",
                 logger=logger,
             ),
@@ -103,7 +106,7 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
             minibatch_size=self.hyperparameters["minibatch_size"],
             replay_start_size=self.hyperparameters["replay_start_size"],
             update_frequency=self.hyperparameters["update_frequency"],
-            logger=logger
+            logger=logger,
         )
 
     def test_agent(self):
@@ -111,11 +114,13 @@ def test_agent(self):
             copy.deepcopy(self.model),
             None,
             self.n_actions,
-            self.hyperparameters['atoms'],
-            v_min=self.hyperparameters['v_min'],
-            v_max=self.hyperparameters['v_max'],
+            self.hyperparameters["atoms"],
+            v_min=self.hyperparameters["v_min"],
+            v_max=self.hyperparameters["v_max"],
+        )
+        return C51TestAgent(
+            q_dist, self.n_actions, self.hyperparameters["test_exploration"]
         )
-        return C51TestAgent(q_dist, self.n_actions, self.hyperparameters["test_exploration"])
 
 
-c51 = PresetBuilder('c51', default_hyperparameters, C51ClassicControlPreset)
+c51 = PresetBuilder("c51", default_hyperparameters, C51ClassicControlPreset)
diff --git a/all/presets/classic_control/ddqn.py b/all/presets/classic_control/ddqn.py
index e8e70c45..f7203c7e 100644
--- a/all/presets/classic_control/ddqn.py
+++ b/all/presets/classic_control/ddqn.py
@@ -1,15 +1,16 @@
 import copy
+
 from torch.optim import Adam
+
 from all.agents import DDQN, DDQNTestAgent
-from all.approximation import QNetwork, FixedTarget
+from all.approximation import FixedTarget, QNetwork
 from all.logging import DummyLogger
 from all.memory import PrioritizedReplayBuffer
 from all.optim import LinearScheduler
 from all.policies import GreedyPolicy
 from all.presets.builder import PresetBuilder
-from all.presets.preset import Preset
 from all.presets.classic_control.models import dueling_fc_relu_q
-
+from all.presets.preset import Preset
 
 default_hyperparameters = {
     "discount_factor": 0.99,
@@ -23,15 +24,15 @@
     "replay_start_size": 1000,
     "replay_buffer_size": 10000,
     # Exploration settings
-    "initial_exploration": 1.,
-    "final_exploration": 0.,
+    "initial_exploration": 1.0,
+    "final_exploration": 0.0,
     "final_exploration_step": 10000,
     "test_exploration": 0.001,
     # Prioritized replay settings
     "alpha": 0.2,
     "beta": 0.6,
     # Model construction
-    "model_constructor": dueling_fc_relu_q
+    "model_constructor": dueling_fc_relu_q,
 }
 
 
@@ -66,37 +67,38 @@ class DDQNClassicControlPreset(Preset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.model = hyperparameters['model_constructor'](env).to(device)
+        self.model = hyperparameters["model_constructor"](env).to(device)
         self.n_actions = env.action_space.n
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        optimizer = Adam(self.model.parameters(), lr=self.hyperparameters['lr'])
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        optimizer = Adam(self.model.parameters(), lr=self.hyperparameters["lr"])
 
         q = QNetwork(
             self.model,
             optimizer,
-            target=FixedTarget(self.hyperparameters['target_update_frequency']),
-            logger=logger
+            target=FixedTarget(self.hyperparameters["target_update_frequency"]),
+            logger=logger,
         )
 
         policy = GreedyPolicy(
             q,
             self.n_actions,
             epsilon=LinearScheduler(
-                self.hyperparameters['initial_exploration'],
-                self.hyperparameters['final_exploration'],
-                self.hyperparameters['replay_start_size'],
-                self.hyperparameters['final_exploration_step'] - self.hyperparameters['replay_start_size'],
+                self.hyperparameters["initial_exploration"],
+                self.hyperparameters["final_exploration"],
+                self.hyperparameters["replay_start_size"],
+                self.hyperparameters["final_exploration_step"]
+                - self.hyperparameters["replay_start_size"],
                 name="exploration",
-                logger=logger
-            )
+                logger=logger,
+            ),
         )
 
         replay_buffer = PrioritizedReplayBuffer(
-            self.hyperparameters['replay_buffer_size'],
-            alpha=self.hyperparameters['alpha'],
-            beta=self.hyperparameters['beta'],
-            device=self.device
+            self.hyperparameters["replay_buffer_size"],
+            alpha=self.hyperparameters["alpha"],
+            beta=self.hyperparameters["beta"],
+            device=self.device,
         )
 
         return DDQN(
@@ -111,8 +113,10 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
 
     def test_agent(self):
         q = QNetwork(copy.deepcopy(self.model))
-        policy = GreedyPolicy(q, self.n_actions, epsilon=self.hyperparameters['test_exploration'])
+        policy = GreedyPolicy(
+            q, self.n_actions, epsilon=self.hyperparameters["test_exploration"]
+        )
         return DDQNTestAgent(policy)
 
 
-ddqn = PresetBuilder('ddqn', default_hyperparameters, DDQNClassicControlPreset)
+ddqn = PresetBuilder("ddqn", default_hyperparameters, DDQNClassicControlPreset)
diff --git a/all/presets/classic_control/dqn.py b/all/presets/classic_control/dqn.py
index 6fb84bc1..fdc94e4e 100644
--- a/all/presets/classic_control/dqn.py
+++ b/all/presets/classic_control/dqn.py
@@ -1,15 +1,16 @@
 import copy
+
 from torch.optim import Adam
+
 from all.agents import DQN, DQNTestAgent
-from all.approximation import QNetwork, FixedTarget
+from all.approximation import FixedTarget, QNetwork
 from all.logging import DummyLogger
 from all.memory import ExperienceReplayBuffer
 from all.optim import LinearScheduler
 from all.policies import GreedyPolicy
 from all.presets.builder import PresetBuilder
-from all.presets.preset import Preset
 from all.presets.classic_control.models import fc_relu_q
-
+from all.presets.preset import Preset
 
 default_hyperparameters = {
     # Common settings
@@ -24,12 +25,12 @@
     "replay_start_size": 1000,
     "replay_buffer_size": 10000,
     # Explicit exploration
-    "initial_exploration": 1.,
-    "final_exploration": 0.,
+    "initial_exploration": 1.0,
+    "final_exploration": 0.0,
     "final_exploration_step": 10000,
     "test_exploration": 0.001,
     # Model construction
-    "model_constructor": fc_relu_q
+    "model_constructor": fc_relu_q,
 }
 
 
@@ -60,51 +61,53 @@ class DQNClassicControlPreset(Preset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.model = hyperparameters['model_constructor'](env).to(device)
+        self.model = hyperparameters["model_constructor"](env).to(device)
         self.n_actions = env.action_space.n
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        optimizer = Adam(self.model.parameters(), lr=self.hyperparameters['lr'])
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        optimizer = Adam(self.model.parameters(), lr=self.hyperparameters["lr"])
 
         q = QNetwork(
             self.model,
             optimizer,
-            target=FixedTarget(self.hyperparameters['target_update_frequency']),
-            logger=logger
+            target=FixedTarget(self.hyperparameters["target_update_frequency"]),
+            logger=logger,
         )
 
         policy = GreedyPolicy(
             q,
             self.n_actions,
             epsilon=LinearScheduler(
-                self.hyperparameters['initial_exploration'],
-                self.hyperparameters['final_exploration'],
-                self.hyperparameters['replay_start_size'],
-                self.hyperparameters['final_exploration_step'] - self.hyperparameters['replay_start_size'],
+                self.hyperparameters["initial_exploration"],
+                self.hyperparameters["final_exploration"],
+                self.hyperparameters["replay_start_size"],
+                self.hyperparameters["final_exploration_step"]
+                - self.hyperparameters["replay_start_size"],
                 name="exploration",
-                logger=logger
-            )
+                logger=logger,
+            ),
         )
 
         replay_buffer = ExperienceReplayBuffer(
-            self.hyperparameters['replay_buffer_size'],
-            device=self.device
+            self.hyperparameters["replay_buffer_size"], device=self.device
         )
 
         return DQN(
             q,
             policy,
             replay_buffer,
-            discount_factor=self.hyperparameters['discount_factor'],
-            minibatch_size=self.hyperparameters['minibatch_size'],
-            replay_start_size=self.hyperparameters['replay_start_size'],
-            update_frequency=self.hyperparameters['update_frequency'],
+            discount_factor=self.hyperparameters["discount_factor"],
+            minibatch_size=self.hyperparameters["minibatch_size"],
+            replay_start_size=self.hyperparameters["replay_start_size"],
+            update_frequency=self.hyperparameters["update_frequency"],
         )
 
     def test_agent(self):
         q = QNetwork(copy.deepcopy(self.model))
-        policy = GreedyPolicy(q, self.n_actions, epsilon=self.hyperparameters['test_exploration'])
+        policy = GreedyPolicy(
+            q, self.n_actions, epsilon=self.hyperparameters["test_exploration"]
+        )
         return DQNTestAgent(policy)
 
 
-dqn = PresetBuilder('dqn', default_hyperparameters, DQNClassicControlPreset)
+dqn = PresetBuilder("dqn", default_hyperparameters, DQNClassicControlPreset)
diff --git a/all/presets/classic_control/ppo.py b/all/presets/classic_control/ppo.py
index 961405e7..cf27709f 100644
--- a/all/presets/classic_control/ppo.py
+++ b/all/presets/classic_control/ppo.py
@@ -1,16 +1,19 @@
 import copy
+
 from torch.optim import Adam
-from torch.optim.lr_scheduler import CosineAnnealingLR
+
 from all.agents import PPO, PPOTestAgent
-from all.bodies import DeepmindAtariBody
-from all.approximation import VNetwork, FeatureNetwork
+from all.approximation import FeatureNetwork, VNetwork
 from all.logging import DummyLogger
 from all.optim import LinearScheduler
 from all.policies import SoftmaxPolicy
 from all.presets.builder import ParallelPresetBuilder
+from all.presets.classic_control.models import (
+    fc_policy_head,
+    fc_relu_features,
+    fc_value_head,
+)
 from all.presets.preset import ParallelPreset
-from all.presets.classic_control.models import fc_relu_features, fc_policy_head, fc_value_head
-
 
 default_hyperparameters = {
     # Common settings
@@ -34,7 +37,7 @@
     # Model construction
     "feature_model_constructor": fc_relu_features,
     "value_model_constructor": fc_value_head,
-    "policy_model_constructor": fc_policy_head
+    "policy_model_constructor": fc_policy_head,
 }
 
 
@@ -69,22 +72,41 @@ class PPOClassicControlPreset(ParallelPreset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.value_model = hyperparameters['value_model_constructor']().to(device)
-        self.policy_model = hyperparameters['policy_model_constructor'](env).to(device)
-        self.feature_model = hyperparameters['feature_model_constructor'](env).to(device)
+        self.value_model = hyperparameters["value_model_constructor"]().to(device)
+        self.policy_model = hyperparameters["policy_model_constructor"](env).to(device)
+        self.feature_model = hyperparameters["feature_model_constructor"](env).to(
+            device
+        )
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        n_updates = train_steps * self.hyperparameters['epochs'] * self.hyperparameters['minibatches'] / (self.hyperparameters['n_steps'] * self.hyperparameters['n_envs'])
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        n_updates = (
+            train_steps
+            * self.hyperparameters["epochs"]
+            * self.hyperparameters["minibatches"]
+            / (self.hyperparameters["n_steps"] * self.hyperparameters["n_envs"])
+        )
 
-        feature_optimizer = Adam(self.feature_model.parameters(), lr=self.hyperparameters["lr"], eps=self.hyperparameters["eps"])
-        value_optimizer = Adam(self.value_model.parameters(), lr=self.hyperparameters["lr"], eps=self.hyperparameters["eps"])
-        policy_optimizer = Adam(self.policy_model.parameters(), lr=self.hyperparameters["lr"], eps=self.hyperparameters["eps"])
+        feature_optimizer = Adam(
+            self.feature_model.parameters(),
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
+        )
+        value_optimizer = Adam(
+            self.value_model.parameters(),
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
+        )
+        policy_optimizer = Adam(
+            self.policy_model.parameters(),
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
+        )
 
         features = FeatureNetwork(
             self.feature_model,
             feature_optimizer,
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         v = VNetwork(
@@ -92,14 +114,14 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
             value_optimizer,
             loss_scaling=self.hyperparameters["value_loss_scaling"],
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         policy = SoftmaxPolicy(
             self.policy_model,
             policy_optimizer,
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         return PPO(
@@ -111,8 +133,8 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
                 self.hyperparameters["clip_final"],
                 0,
                 n_updates,
-                name='clip',
-                logger=logger
+                name="clip",
+                logger=logger,
             ),
             epochs=self.hyperparameters["epochs"],
             minibatches=self.hyperparameters["minibatches"],
@@ -133,4 +155,4 @@ def parallel_test_agent(self):
         return self.test_agent()
 
 
-ppo = ParallelPresetBuilder('ppo', default_hyperparameters, PPOClassicControlPreset)
+ppo = ParallelPresetBuilder("ppo", default_hyperparameters, PPOClassicControlPreset)
diff --git a/all/presets/classic_control/rainbow.py b/all/presets/classic_control/rainbow.py
index c280e742..f60e1186 100644
--- a/all/presets/classic_control/rainbow.py
+++ b/all/presets/classic_control/rainbow.py
@@ -1,16 +1,15 @@
 import copy
+
 from torch.optim import Adam
-from torch.optim.lr_scheduler import CosineAnnealingLR
-from all.approximation import QDist, FixedTarget
+
 from all.agents import Rainbow, RainbowTestAgent
-from all.bodies import DeepmindAtariBody
+from all.approximation import FixedTarget, QDist
 from all.logging import DummyLogger
-from all.memory import PrioritizedReplayBuffer, NStepReplayBuffer
+from all.memory import NStepReplayBuffer, PrioritizedReplayBuffer
 from all.optim import LinearScheduler
 from all.presets.builder import PresetBuilder
-from all.presets.preset import Preset
 from all.presets.classic_control.models import fc_relu_rainbow
-
+from all.presets.preset import Preset
 
 default_hyperparameters = {
     "discount_factor": 0.99,
@@ -25,7 +24,7 @@
     "replay_buffer_size": 20000,
     # Explicit exploration
     "initial_exploration": 0.02,
-    "final_exploration": 0.,
+    "final_exploration": 0.0,
     "test_exploration": 0.001,
     # Prioritized replay settings
     "alpha": 0.5,
@@ -39,7 +38,7 @@
     # Noisy Nets
     "sigma": 0.5,
     # Model construction
-    "model_constructor": fc_relu_rainbow
+    "model_constructor": fc_relu_rainbow,
 }
 
 
@@ -79,53 +78,56 @@ class RainbowClassicControlPreset(Preset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.model = hyperparameters['model_constructor'](env, atoms=hyperparameters["atoms"], sigma=hyperparameters["sigma"]).to(device)
+        self.model = hyperparameters["model_constructor"](
+            env, atoms=hyperparameters["atoms"], sigma=hyperparameters["sigma"]
+        ).to(device)
         self.n_actions = env.action_space.n
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
         optimizer = Adam(
             self.model.parameters(),
-            lr=self.hyperparameters['lr'],
-            eps=self.hyperparameters['eps']
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
         )
 
         q_dist = QDist(
             self.model,
             optimizer,
             self.n_actions,
-            self.hyperparameters['atoms'],
-            v_min=self.hyperparameters['v_min'],
-            v_max=self.hyperparameters['v_max'],
-            target=FixedTarget(self.hyperparameters['target_update_frequency']),
+            self.hyperparameters["atoms"],
+            v_min=self.hyperparameters["v_min"],
+            v_max=self.hyperparameters["v_max"],
+            target=FixedTarget(self.hyperparameters["target_update_frequency"]),
             logger=logger,
         )
 
         replay_buffer = NStepReplayBuffer(
-            self.hyperparameters['n_steps'],
-            self.hyperparameters['discount_factor'],
+            self.hyperparameters["n_steps"],
+            self.hyperparameters["discount_factor"],
             PrioritizedReplayBuffer(
-                self.hyperparameters['replay_buffer_size'],
-                alpha=self.hyperparameters['alpha'],
-                beta=self.hyperparameters['beta'],
-                device=self.device
-            )
+                self.hyperparameters["replay_buffer_size"],
+                alpha=self.hyperparameters["alpha"],
+                beta=self.hyperparameters["beta"],
+                device=self.device,
+            ),
         )
 
         return Rainbow(
             q_dist,
             replay_buffer,
             exploration=LinearScheduler(
-                self.hyperparameters['initial_exploration'],
-                self.hyperparameters['final_exploration'],
+                self.hyperparameters["initial_exploration"],
+                self.hyperparameters["final_exploration"],
                 0,
-                train_steps - self.hyperparameters['replay_start_size'],
+                train_steps - self.hyperparameters["replay_start_size"],
                 name="exploration",
-                logger=logger
+                logger=logger,
             ),
-            discount_factor=self.hyperparameters['discount_factor'] ** self.hyperparameters["n_steps"],
-            minibatch_size=self.hyperparameters['minibatch_size'],
-            replay_start_size=self.hyperparameters['replay_start_size'],
-            update_frequency=self.hyperparameters['update_frequency'],
+            discount_factor=self.hyperparameters["discount_factor"]
+            ** self.hyperparameters["n_steps"],
+            minibatch_size=self.hyperparameters["minibatch_size"],
+            replay_start_size=self.hyperparameters["replay_start_size"],
+            update_frequency=self.hyperparameters["update_frequency"],
             logger=logger,
         )
 
@@ -134,11 +136,13 @@ def test_agent(self):
             copy.deepcopy(self.model),
             None,
             self.n_actions,
-            self.hyperparameters['atoms'],
-            v_min=self.hyperparameters['v_min'],
-            v_max=self.hyperparameters['v_max'],
+            self.hyperparameters["atoms"],
+            v_min=self.hyperparameters["v_min"],
+            v_max=self.hyperparameters["v_max"],
+        )
+        return RainbowTestAgent(
+            q_dist, self.n_actions, self.hyperparameters["test_exploration"]
         )
-        return RainbowTestAgent(q_dist, self.n_actions, self.hyperparameters["test_exploration"])
 
 
-rainbow = PresetBuilder('rainbow', default_hyperparameters, RainbowClassicControlPreset)
+rainbow = PresetBuilder("rainbow", default_hyperparameters, RainbowClassicControlPreset)
diff --git a/all/presets/classic_control/vac.py b/all/presets/classic_control/vac.py
index c607e250..52fcbaae 100644
--- a/all/presets/classic_control/vac.py
+++ b/all/presets/classic_control/vac.py
@@ -1,15 +1,18 @@
 import copy
+
 from torch.optim import Adam
-from torch.optim.lr_scheduler import CosineAnnealingLR
+
 from all.agents import VAC, VACTestAgent
-from all.approximation import VNetwork, FeatureNetwork
-from all.bodies import DeepmindAtariBody
+from all.approximation import FeatureNetwork, VNetwork
 from all.logging import DummyLogger
 from all.policies import SoftmaxPolicy
 from all.presets.builder import ParallelPresetBuilder
+from all.presets.classic_control.models import (
+    fc_policy_head,
+    fc_relu_features,
+    fc_value_head,
+)
 from all.presets.preset import ParallelPreset
-from all.presets.classic_control.models import fc_relu_features, fc_policy_head, fc_value_head
-
 
 default_hyperparameters = {
     # Common settings
@@ -26,7 +29,7 @@
     # Model construction
     "feature_model_constructor": fc_relu_features,
     "value_model_constructor": fc_value_head,
-    "policy_model_constructor": fc_policy_head
+    "policy_model_constructor": fc_policy_head,
 }
 
 
@@ -55,20 +58,34 @@ class VACClassicControlPreset(ParallelPreset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.value_model = hyperparameters['value_model_constructor']().to(device)
-        self.policy_model = hyperparameters['policy_model_constructor'](env).to(device)
-        self.feature_model = hyperparameters['feature_model_constructor'](env).to(device)
+        self.value_model = hyperparameters["value_model_constructor"]().to(device)
+        self.policy_model = hyperparameters["policy_model_constructor"](env).to(device)
+        self.feature_model = hyperparameters["feature_model_constructor"](env).to(
+            device
+        )
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        feature_optimizer = Adam(self.feature_model.parameters(), lr=self.hyperparameters["lr_pi"], eps=self.hyperparameters["eps"])
-        value_optimizer = Adam(self.value_model.parameters(), lr=self.hyperparameters["lr_v"], eps=self.hyperparameters["eps"])
-        policy_optimizer = Adam(self.policy_model.parameters(), lr=self.hyperparameters["lr_pi"], eps=self.hyperparameters["eps"])
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        feature_optimizer = Adam(
+            self.feature_model.parameters(),
+            lr=self.hyperparameters["lr_pi"],
+            eps=self.hyperparameters["eps"],
+        )
+        value_optimizer = Adam(
+            self.value_model.parameters(),
+            lr=self.hyperparameters["lr_v"],
+            eps=self.hyperparameters["eps"],
+        )
+        policy_optimizer = Adam(
+            self.policy_model.parameters(),
+            lr=self.hyperparameters["lr_pi"],
+            eps=self.hyperparameters["eps"],
+        )
 
         features = FeatureNetwork(
             self.feature_model,
             feature_optimizer,
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         v = VNetwork(
@@ -76,17 +93,19 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
             value_optimizer,
             loss_scaling=self.hyperparameters["value_loss_scaling"],
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         policy = SoftmaxPolicy(
             self.policy_model,
             policy_optimizer,
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
-        return VAC(features, v, policy, discount_factor=self.hyperparameters["discount_factor"])
+        return VAC(
+            features, v, policy, discount_factor=self.hyperparameters["discount_factor"]
+        )
 
     def test_agent(self):
         features = FeatureNetwork(copy.deepcopy(self.feature_model))
@@ -97,4 +116,4 @@ def parallel_test_agent(self):
         return self.test_agent()
 
 
-vac = ParallelPresetBuilder('vac', default_hyperparameters, VACClassicControlPreset)
+vac = ParallelPresetBuilder("vac", default_hyperparameters, VACClassicControlPreset)
diff --git a/all/presets/classic_control/vpg.py b/all/presets/classic_control/vpg.py
index 4392d80b..b1b39d2b 100644
--- a/all/presets/classic_control/vpg.py
+++ b/all/presets/classic_control/vpg.py
@@ -1,15 +1,18 @@
 import copy
+
 from torch.optim import Adam
-from torch.optim.lr_scheduler import CosineAnnealingLR
+
 from all.agents import VPG, VPGTestAgent
-from all.approximation import VNetwork, FeatureNetwork
-from all.bodies import DeepmindAtariBody
+from all.approximation import FeatureNetwork, VNetwork
 from all.logging import DummyLogger
 from all.policies import SoftmaxPolicy
 from all.presets.builder import PresetBuilder
+from all.presets.classic_control.models import (
+    fc_policy_head,
+    fc_relu_features,
+    fc_value_head,
+)
 from all.presets.preset import Preset
-from all.presets.classic_control.models import fc_relu_features, fc_policy_head, fc_value_head
-
 
 default_hyperparameters = {
     # Common settings
@@ -25,7 +28,7 @@
     # Model construction
     "feature_model_constructor": fc_relu_features,
     "value_model_constructor": fc_value_head,
-    "policy_model_constructor": fc_policy_head
+    "policy_model_constructor": fc_policy_head,
 }
 
 
@@ -54,20 +57,34 @@ class VPGClassicControlPreset(Preset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.value_model = hyperparameters['value_model_constructor']().to(device)
-        self.policy_model = hyperparameters['policy_model_constructor'](env).to(device)
-        self.feature_model = hyperparameters['feature_model_constructor'](env).to(device)
+        self.value_model = hyperparameters["value_model_constructor"]().to(device)
+        self.policy_model = hyperparameters["policy_model_constructor"](env).to(device)
+        self.feature_model = hyperparameters["feature_model_constructor"](env).to(
+            device
+        )
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        feature_optimizer = Adam(self.feature_model.parameters(), lr=self.hyperparameters["lr_pi"], eps=self.hyperparameters["eps"])
-        value_optimizer = Adam(self.value_model.parameters(), lr=self.hyperparameters["lr_v"], eps=self.hyperparameters["eps"])
-        policy_optimizer = Adam(self.policy_model.parameters(), lr=self.hyperparameters["lr_pi"], eps=self.hyperparameters["eps"])
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        feature_optimizer = Adam(
+            self.feature_model.parameters(),
+            lr=self.hyperparameters["lr_pi"],
+            eps=self.hyperparameters["eps"],
+        )
+        value_optimizer = Adam(
+            self.value_model.parameters(),
+            lr=self.hyperparameters["lr_v"],
+            eps=self.hyperparameters["eps"],
+        )
+        policy_optimizer = Adam(
+            self.policy_model.parameters(),
+            lr=self.hyperparameters["lr_pi"],
+            eps=self.hyperparameters["eps"],
+        )
 
         features = FeatureNetwork(
             self.feature_model,
             feature_optimizer,
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         v = VNetwork(
@@ -75,17 +92,23 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
             value_optimizer,
             loss_scaling=self.hyperparameters["value_loss_scaling"],
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
         policy = SoftmaxPolicy(
             self.policy_model,
             policy_optimizer,
             clip_grad=self.hyperparameters["clip_grad"],
-            logger=logger
+            logger=logger,
         )
 
-        return VPG(features, v, policy, discount_factor=self.hyperparameters["discount_factor"], min_batch_size=self.hyperparameters["min_batch_size"])
+        return VPG(
+            features,
+            v,
+            policy,
+            discount_factor=self.hyperparameters["discount_factor"],
+            min_batch_size=self.hyperparameters["min_batch_size"],
+        )
 
     def test_agent(self):
         features = FeatureNetwork(copy.deepcopy(self.feature_model))
@@ -96,4 +119,4 @@ def parallel_test_agent(self):
         return self.test_agent()
 
 
-vpg = PresetBuilder('vpg', default_hyperparameters, VPGClassicControlPreset)
+vpg = PresetBuilder("vpg", default_hyperparameters, VPGClassicControlPreset)
diff --git a/all/presets/classic_control/vqn.py b/all/presets/classic_control/vqn.py
index 76258676..557db597 100644
--- a/all/presets/classic_control/vqn.py
+++ b/all/presets/classic_control/vqn.py
@@ -1,16 +1,16 @@
 import copy
+
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
-from all.approximation import QNetwork
+
 from all.agents import VQN, VQNTestAgent
-from all.bodies import DeepmindAtariBody
+from all.approximation import QNetwork
 from all.logging import DummyLogger
 from all.optim import LinearScheduler
 from all.policies import GreedyPolicy, ParallelGreedyPolicy
 from all.presets.builder import ParallelPresetBuilder
-from all.presets.preset import ParallelPreset
 from all.presets.classic_control.models import dueling_fc_relu_q
-
+from all.presets.preset import ParallelPreset
 
 default_hyperparameters = {
     # Common settings
@@ -19,14 +19,14 @@
     "lr": 1e-2,
     "eps": 1.5e-4,
     # Explicit exploration
-    "initial_exploration": 1.,
-    "final_exploration": 0.,
+    "initial_exploration": 1.0,
+    "final_exploration": 0.0,
     "final_exploration_step": 10000,
     "test_exploration": 0.001,
     # Parallel actors
     "n_envs": 8,
     # Model construction
-    "model_constructor": dueling_fc_relu_q
+    "model_constructor": dueling_fc_relu_q,
 }
 
 
@@ -54,49 +54,54 @@ class VQNClassicControlPreset(ParallelPreset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.model = hyperparameters['model_constructor'](env).to(device)
+        self.model = hyperparameters["model_constructor"](env).to(device)
         self.n_actions = env.action_space.n
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        n_updates = train_steps / self.hyperparameters['n_envs']
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        n_updates = train_steps / self.hyperparameters["n_envs"]
 
         optimizer = Adam(
             self.model.parameters(),
-            lr=self.hyperparameters['lr'],
-            eps=self.hyperparameters['eps']
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
         )
 
         q = QNetwork(
             self.model,
             optimizer,
             scheduler=CosineAnnealingLR(optimizer, n_updates),
-            logger=logger
+            logger=logger,
         )
 
         policy = ParallelGreedyPolicy(
             q,
             self.n_actions,
             epsilon=LinearScheduler(
-                self.hyperparameters['initial_exploration'],
-                self.hyperparameters['final_exploration'],
+                self.hyperparameters["initial_exploration"],
+                self.hyperparameters["final_exploration"],
                 0,
-                self.hyperparameters["final_exploration_step"] / self.hyperparameters["n_envs"],
+                self.hyperparameters["final_exploration_step"]
+                / self.hyperparameters["n_envs"],
                 name="exploration",
-                logger=logger
-            )
+                logger=logger,
+            ),
         )
 
-        return VQN(q, policy, discount_factor=self.hyperparameters['discount_factor'])
+        return VQN(q, policy, discount_factor=self.hyperparameters["discount_factor"])
 
     def test_agent(self):
         q = QNetwork(copy.deepcopy(self.model))
-        policy = GreedyPolicy(q, self.n_actions, epsilon=self.hyperparameters["test_exploration"])
+        policy = GreedyPolicy(
+            q, self.n_actions, epsilon=self.hyperparameters["test_exploration"]
+        )
         return VQNTestAgent(policy)
 
     def parallel_test_agent(self):
         q = QNetwork(copy.deepcopy(self.model))
-        policy = ParallelGreedyPolicy(q, self.n_actions, epsilon=self.hyperparameters["test_exploration"])
+        policy = ParallelGreedyPolicy(
+            q, self.n_actions, epsilon=self.hyperparameters["test_exploration"]
+        )
         return VQNTestAgent(policy)
 
 
-vqn = ParallelPresetBuilder('vqn', default_hyperparameters, VQNClassicControlPreset)
+vqn = ParallelPresetBuilder("vqn", default_hyperparameters, VQNClassicControlPreset)
diff --git a/all/presets/classic_control/vsarsa.py b/all/presets/classic_control/vsarsa.py
index 50db3372..e1c963c1 100644
--- a/all/presets/classic_control/vsarsa.py
+++ b/all/presets/classic_control/vsarsa.py
@@ -1,16 +1,16 @@
 import copy
+
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
-from all.approximation import QNetwork
+
 from all.agents import VSarsa, VSarsaTestAgent
-from all.bodies import DeepmindAtariBody
+from all.approximation import QNetwork
 from all.logging import DummyLogger
 from all.optim import LinearScheduler
 from all.policies import GreedyPolicy, ParallelGreedyPolicy
 from all.presets.builder import ParallelPresetBuilder
-from all.presets.preset import ParallelPreset
 from all.presets.classic_control.models import dueling_fc_relu_q
-
+from all.presets.preset import ParallelPreset
 
 default_hyperparameters = {
     # Common settings
@@ -19,14 +19,14 @@
     "lr": 1e-2,
     "eps": 1.5e-4,
     # Explicit exploration
-    "initial_exploration": 1.,
-    "final_exploration": 0.,
+    "initial_exploration": 1.0,
+    "final_exploration": 0.0,
     "final_exploration_step": 10000,
     "test_exploration": 0.001,
     # Parallel actors
     "n_envs": 8,
     # Model construction
-    "model_constructor": dueling_fc_relu_q
+    "model_constructor": dueling_fc_relu_q,
 }
 
 
@@ -54,49 +54,58 @@ class VSarsaClassicControlPreset(ParallelPreset):
 
     def __init__(self, env, name, device, **hyperparameters):
         super().__init__(name, device, hyperparameters)
-        self.model = hyperparameters['model_constructor'](env).to(device)
+        self.model = hyperparameters["model_constructor"](env).to(device)
         self.n_actions = env.action_space.n
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        n_updates = train_steps / self.hyperparameters['n_envs']
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        n_updates = train_steps / self.hyperparameters["n_envs"]
 
         optimizer = Adam(
             self.model.parameters(),
-            lr=self.hyperparameters['lr'],
-            eps=self.hyperparameters['eps']
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
         )
 
         q = QNetwork(
             self.model,
             optimizer,
             scheduler=CosineAnnealingLR(optimizer, n_updates),
-            logger=logger
+            logger=logger,
         )
 
         policy = ParallelGreedyPolicy(
             q,
             self.n_actions,
             epsilon=LinearScheduler(
-                self.hyperparameters['initial_exploration'],
-                self.hyperparameters['final_exploration'],
+                self.hyperparameters["initial_exploration"],
+                self.hyperparameters["final_exploration"],
                 0,
-                self.hyperparameters["final_exploration_step"] / self.hyperparameters["n_envs"],
+                self.hyperparameters["final_exploration_step"]
+                / self.hyperparameters["n_envs"],
                 name="exploration",
-                logger=logger
-            )
+                logger=logger,
+            ),
         )
 
-        return VSarsa(q, policy, discount_factor=self.hyperparameters['discount_factor'])
+        return VSarsa(
+            q, policy, discount_factor=self.hyperparameters["discount_factor"]
+        )
 
     def test_agent(self):
         q = QNetwork(copy.deepcopy(self.model))
-        policy = GreedyPolicy(q, self.n_actions, epsilon=self.hyperparameters["test_exploration"])
+        policy = GreedyPolicy(
+            q, self.n_actions, epsilon=self.hyperparameters["test_exploration"]
+        )
         return VSarsaTestAgent(policy)
 
     def parallel_test_agent(self):
         q = QNetwork(copy.deepcopy(self.model))
-        policy = ParallelGreedyPolicy(q, self.n_actions, epsilon=self.hyperparameters["test_exploration"])
+        policy = ParallelGreedyPolicy(
+            q, self.n_actions, epsilon=self.hyperparameters["test_exploration"]
+        )
         return VSarsaTestAgent(policy)
 
 
-vsarsa = ParallelPresetBuilder('vsarsa', default_hyperparameters, VSarsaClassicControlPreset)
+vsarsa = ParallelPresetBuilder(
+    "vsarsa", default_hyperparameters, VSarsaClassicControlPreset
+)
diff --git a/all/presets/continuous/__init__.py b/all/presets/continuous/__init__.py
index 32a98557..7b0a9fc4 100644
--- a/all/presets/continuous/__init__.py
+++ b/all/presets/continuous/__init__.py
@@ -1,10 +1,10 @@
 # from .actor_critic import actor_critic
-from .ddpg import ddpg, DDPGContinuousPreset
-from .ppo import ppo, PPOContinuousPreset
+from .ddpg import ddpg
+from .ppo import ppo
 from .sac import sac
 
 __all__ = [
-    'ddpg',
-    'ppo',
-    'sac',
+    "ddpg",
+    "ppo",
+    "sac",
 ]
diff --git a/all/presets/continuous/ddpg.py b/all/presets/continuous/ddpg.py
index 9fb2d244..62c130dd 100644
--- a/all/presets/continuous/ddpg.py
+++ b/all/presets/continuous/ddpg.py
@@ -1,16 +1,17 @@
 import copy
+
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
+
 from all.agents import DDPG, DDPGTestAgent
-from all.approximation import QContinuous, PolyakTarget
+from all.approximation import PolyakTarget, QContinuous
 from all.bodies import TimeFeature
 from all.logging import DummyLogger
-from all.policies import DeterministicPolicy
 from all.memory import ExperienceReplayBuffer
+from all.policies import DeterministicPolicy
 from all.presets.builder import PresetBuilder
+from all.presets.continuous.models import fc_deterministic_policy, fc_q
 from all.presets.preset import Preset
-from all.presets.continuous.models import fc_q, fc_deterministic_policy
-
 
 default_hyperparameters = {
     # Common settings
@@ -29,7 +30,7 @@
     "noise": 0.1,
     # Model construction
     "q_model_constructor": fc_q,
-    "policy_model_constructor": fc_deterministic_policy
+    "policy_model_constructor": fc_deterministic_policy,
 }
 
 
@@ -62,8 +63,10 @@ def __init__(self, env, name, device, **hyperparameters):
         self.policy_model = hyperparameters["policy_model_constructor"](env).to(device)
         self.action_space = env.action_space
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        n_updates = (train_steps - self.hyperparameters["replay_start_size"]) / self.hyperparameters["update_frequency"]
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        n_updates = (
+            train_steps - self.hyperparameters["replay_start_size"]
+        ) / self.hyperparameters["update_frequency"]
 
         q_optimizer = Adam(self.q_model.parameters(), lr=self.hyperparameters["lr_q"])
 
@@ -71,42 +74,39 @@ def agent(self, logger=DummyLogger(), train_steps=float('inf')):
             self.q_model,
             q_optimizer,
             target=PolyakTarget(self.hyperparameters["polyak_rate"]),
-            scheduler=CosineAnnealingLR(
-                q_optimizer,
-                n_updates
-            ),
-            logger=logger
+            scheduler=CosineAnnealingLR(q_optimizer, n_updates),
+            logger=logger,
         )
 
-        policy_optimizer = Adam(self.policy_model.parameters(), lr=self.hyperparameters["lr_pi"])
+        policy_optimizer = Adam(
+            self.policy_model.parameters(), lr=self.hyperparameters["lr_pi"]
+        )
         policy = DeterministicPolicy(
             self.policy_model,
             policy_optimizer,
             self.action_space,
             target=PolyakTarget(self.hyperparameters["polyak_rate"]),
-            scheduler=CosineAnnealingLR(
-                policy_optimizer,
-                n_updates
-            ),
-            logger=logger
+            scheduler=CosineAnnealingLR(policy_optimizer, n_updates),
+            logger=logger,
         )
 
         replay_buffer = ExperienceReplayBuffer(
-            self.hyperparameters["replay_buffer_size"],
-            device=self.device
+            self.hyperparameters["replay_buffer_size"], device=self.device
         )
 
-        return TimeFeature(DDPG(
-            q,
-            policy,
-            replay_buffer,
-            self.action_space,
-            noise=self.hyperparameters["noise"],
-            replay_start_size=self.hyperparameters["replay_start_size"],
-            discount_factor=self.hyperparameters["discount_factor"],
-            update_frequency=self.hyperparameters["update_frequency"],
-            minibatch_size=self.hyperparameters["minibatch_size"],
-        ))
+        return TimeFeature(
+            DDPG(
+                q,
+                policy,
+                replay_buffer,
+                self.action_space,
+                noise=self.hyperparameters["noise"],
+                replay_start_size=self.hyperparameters["replay_start_size"],
+                discount_factor=self.hyperparameters["discount_factor"],
+                update_frequency=self.hyperparameters["update_frequency"],
+                minibatch_size=self.hyperparameters["minibatch_size"],
+            )
+        )
 
     def test_agent(self):
         policy = DeterministicPolicy(
@@ -117,4 +117,4 @@ def test_agent(self):
         return TimeFeature(DDPGTestAgent(policy))
 
 
-ddpg = PresetBuilder('ddpg', default_hyperparameters, DDPGContinuousPreset)
+ddpg = PresetBuilder("ddpg", default_hyperparameters, DDPGContinuousPreset)
diff --git a/all/presets/continuous/models/__init__.py b/all/presets/continuous/models/__init__.py
index 012b56c9..937ca44f 100644
--- a/all/presets/continuous/models/__init__.py
+++ b/all/presets/continuous/models/__init__.py
@@ -1,11 +1,13 @@
-'''
+"""
 Pytorch models for continuous control.
 
 All models assume that a feature representing the
 current timestep is used in addition to the features
 received from the environment.
-'''
+"""
+
 import torch
+
 from all import nn
 
 
@@ -62,7 +64,7 @@ def __init__(self, env, hidden1=400, hidden2=300):
             nn.Tanh(),
             nn.Linear(hidden1, hidden2),
             nn.Tanh(),
-            nn.Linear(hidden2, env.action_space.shape[0])
+            nn.Linear(hidden2, env.action_space.shape[0]),
         )
         self.log_stds = nn.Parameter(torch.zeros(env.action_space.shape[0]))
 
diff --git a/all/presets/continuous/ppo.py b/all/presets/continuous/ppo.py
index 412c77e3..dbf73568 100644
--- a/all/presets/continuous/ppo.py
+++ b/all/presets/continuous/ppo.py
@@ -1,16 +1,17 @@
 import copy
+
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
+
 from all.agents import PPO, PPOTestAgent
-from all.approximation import VNetwork, FeatureNetwork, Identity
+from all.approximation import Identity, VNetwork
 from all.bodies import TimeFeature
 from all.logging import DummyLogger
 from all.optim import LinearScheduler
 from all.policies import GaussianPolicy
 from all.presets.builder import ParallelPresetBuilder
-from all.presets.preset import ParallelPreset
 from all.presets.continuous.models import fc_policy, fc_v
-
+from all.presets.preset import ParallelPreset
 
 default_hyperparameters = {
     # Common settings
@@ -72,66 +73,77 @@ def __init__(self, env, name, device, **hyperparameters):
         self.policy_model = hyperparameters["policy_model_constructor"](env).to(device)
         self.action_space = env.action_space
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        n_updates = train_steps * self.hyperparameters['epochs'] * self.hyperparameters['minibatches'] / (self.hyperparameters['n_steps'] * self.hyperparameters['n_envs'])
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        n_updates = (
+            train_steps
+            * self.hyperparameters["epochs"]
+            * self.hyperparameters["minibatches"]
+            / (self.hyperparameters["n_steps"] * self.hyperparameters["n_envs"])
+        )
 
-        value_optimizer = Adam(self.value_model.parameters(), lr=self.hyperparameters['lr'], eps=self.hyperparameters['eps'])
-        policy_optimizer = Adam(self.policy_model.parameters(), lr=self.hyperparameters['lr'], eps=self.hyperparameters['eps'])
+        value_optimizer = Adam(
+            self.value_model.parameters(),
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
+        )
+        policy_optimizer = Adam(
+            self.policy_model.parameters(),
+            lr=self.hyperparameters["lr"],
+            eps=self.hyperparameters["eps"],
+        )
 
         features = Identity(self.device)
 
         v = VNetwork(
             self.value_model,
             value_optimizer,
-            loss_scaling=self.hyperparameters['value_loss_scaling'],
-            clip_grad=self.hyperparameters['clip_grad'],
+            loss_scaling=self.hyperparameters["value_loss_scaling"],
+            clip_grad=self.hyperparameters["clip_grad"],
             logger=logger,
-            scheduler=CosineAnnealingLR(
-                value_optimizer,
-                n_updates
-            ),
+            scheduler=CosineAnnealingLR(value_optimizer, n_updates),
         )
 
         policy = GaussianPolicy(
             self.policy_model,
             policy_optimizer,
             self.action_space,
-            clip_grad=self.hyperparameters['clip_grad'],
+            clip_grad=self.hyperparameters["clip_grad"],
             logger=logger,
-            scheduler=CosineAnnealingLR(
-                policy_optimizer,
-                n_updates
-            ),
+            scheduler=CosineAnnealingLR(policy_optimizer, n_updates),
         )
 
-        return TimeFeature(PPO(
-            features,
-            v,
-            policy,
-            epsilon=LinearScheduler(
-                self.hyperparameters['clip_initial'],
-                self.hyperparameters['clip_final'],
-                0,
-                n_updates,
-                name='clip',
-                logger=logger
-            ),
-            epochs=self.hyperparameters['epochs'],
-            minibatches=self.hyperparameters['minibatches'],
-            n_envs=self.hyperparameters['n_envs'],
-            n_steps=self.hyperparameters['n_steps'],
-            discount_factor=self.hyperparameters['discount_factor'],
-            lam=self.hyperparameters['lam'],
-            entropy_loss_scaling=self.hyperparameters['entropy_loss_scaling'],
-            logger=logger,
-        ))
+        return TimeFeature(
+            PPO(
+                features,
+                v,
+                policy,
+                epsilon=LinearScheduler(
+                    self.hyperparameters["clip_initial"],
+                    self.hyperparameters["clip_final"],
+                    0,
+                    n_updates,
+                    name="clip",
+                    logger=logger,
+                ),
+                epochs=self.hyperparameters["epochs"],
+                minibatches=self.hyperparameters["minibatches"],
+                n_envs=self.hyperparameters["n_envs"],
+                n_steps=self.hyperparameters["n_steps"],
+                discount_factor=self.hyperparameters["discount_factor"],
+                lam=self.hyperparameters["lam"],
+                entropy_loss_scaling=self.hyperparameters["entropy_loss_scaling"],
+                logger=logger,
+            )
+        )
 
     def test_agent(self):
-        policy = GaussianPolicy(copy.deepcopy(self.policy_model), space=self.action_space)
+        policy = GaussianPolicy(
+            copy.deepcopy(self.policy_model), space=self.action_space
+        )
         return TimeFeature(PPOTestAgent(Identity(self.device), policy))
 
     def parallel_test_agent(self):
         return self.test_agent()
 
 
-ppo = ParallelPresetBuilder('ppo', default_hyperparameters, PPOContinuousPreset)
+ppo = ParallelPresetBuilder("ppo", default_hyperparameters, PPOContinuousPreset)
diff --git a/all/presets/continuous/sac.py b/all/presets/continuous/sac.py
index f4d9a76f..c33aeddf 100644
--- a/all/presets/continuous/sac.py
+++ b/all/presets/continuous/sac.py
@@ -1,16 +1,17 @@
 import copy
+
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
+
 from all.agents import SAC, SACTestAgent
-from all.approximation import QContinuous, PolyakTarget, VNetwork
+from all.approximation import PolyakTarget, QContinuous
 from all.bodies import TimeFeature
 from all.logging import DummyLogger
-from all.policies.soft_deterministic import SoftDeterministicPolicy
 from all.memory import ExperienceReplayBuffer
+from all.policies.soft_deterministic import SoftDeterministicPolicy
 from all.presets.builder import PresetBuilder
-from all.presets.preset import Preset
 from all.presets.continuous.models import fc_q, fc_soft_policy
-
+from all.presets.preset import Preset
 
 default_hyperparameters = {
     # Common settings
@@ -29,11 +30,11 @@
     "temperature_initial": 0.1,
     "lr_temperature": 1e-5,
     "entropy_backups": True,
-    "entropy_target_scaling": 1.,
+    "entropy_target_scaling": 1.0,
     # Model construction
     "q1_model_constructor": fc_q,
     "q2_model_constructor": fc_q,
-    "policy_model_constructor": fc_soft_policy
+    "policy_model_constructor": fc_soft_policy,
 }
 
 
@@ -70,71 +71,72 @@ def __init__(self, env, name, device, **hyperparameters):
         self.policy_model = hyperparameters["policy_model_constructor"](env).to(device)
         self.action_space = env.action_space
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        n_updates = (train_steps - self.hyperparameters["replay_start_size"]) / self.hyperparameters["update_frequency"]
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        n_updates = (
+            train_steps - self.hyperparameters["replay_start_size"]
+        ) / self.hyperparameters["update_frequency"]
 
         q1_optimizer = Adam(self.q1_model.parameters(), lr=self.hyperparameters["lr_q"])
         q1 = QContinuous(
             self.q1_model,
             q1_optimizer,
-            scheduler=CosineAnnealingLR(
-                q1_optimizer,
-                n_updates
-            ),
+            scheduler=CosineAnnealingLR(q1_optimizer, n_updates),
             target=PolyakTarget(self.hyperparameters["polyak_rate"]),
             logger=logger,
-            name='q1'
+            name="q1",
         )
 
         q2_optimizer = Adam(self.q2_model.parameters(), lr=self.hyperparameters["lr_q"])
         q2 = QContinuous(
             self.q2_model,
             q2_optimizer,
-            scheduler=CosineAnnealingLR(
-                q2_optimizer,
-                n_updates
-            ),
+            scheduler=CosineAnnealingLR(q2_optimizer, n_updates),
             target=PolyakTarget(self.hyperparameters["polyak_rate"]),
             logger=logger,
-            name='q2'
+            name="q2",
         )
 
-        policy_optimizer = Adam(self.policy_model.parameters(), lr=self.hyperparameters["lr_pi"])
+        policy_optimizer = Adam(
+            self.policy_model.parameters(), lr=self.hyperparameters["lr_pi"]
+        )
         policy = SoftDeterministicPolicy(
             self.policy_model,
             policy_optimizer,
             self.action_space,
-            scheduler=CosineAnnealingLR(
-                policy_optimizer,
-                n_updates
-            ),
-            logger=logger
+            scheduler=CosineAnnealingLR(policy_optimizer, n_updates),
+            logger=logger,
         )
 
         replay_buffer = ExperienceReplayBuffer(
-            self.hyperparameters["replay_buffer_size"],
-            device=self.device
+            self.hyperparameters["replay_buffer_size"], device=self.device
         )
 
-        return TimeFeature(SAC(
-            policy,
-            q1,
-            q2,
-            replay_buffer,
-            discount_factor=self.hyperparameters["discount_factor"],
-            entropy_backups=self.hyperparameters["entropy_backups"],
-            entropy_target=(-self.action_space.shape[0] * self.hyperparameters["entropy_target_scaling"]),
-            lr_temperature=self.hyperparameters["lr_temperature"],
-            minibatch_size=self.hyperparameters["minibatch_size"],
-            replay_start_size=self.hyperparameters["replay_start_size"],
-            temperature_initial=self.hyperparameters["temperature_initial"],
-            update_frequency=self.hyperparameters["update_frequency"],
-            logger=logger
-        ))
+        return TimeFeature(
+            SAC(
+                policy,
+                q1,
+                q2,
+                replay_buffer,
+                temperature_initial=self.hyperparameters["temperature_initial"],
+                entropy_backups=self.hyperparameters["entropy_backups"],
+                entropy_target=(
+                    -self.action_space.shape[0]
+                    * self.hyperparameters["entropy_target_scaling"]
+                ),
+                lr_temperature=self.hyperparameters["lr_temperature"],
+                replay_start_size=self.hyperparameters["replay_start_size"],
+                discount_factor=self.hyperparameters["discount_factor"],
+                update_frequency=self.hyperparameters["update_frequency"],
+                minibatch_size=self.hyperparameters["minibatch_size"],
+                logger=logger,
+            )
+        )
 
     def test_agent(self):
-        policy = SoftDeterministicPolicy(copy.deepcopy(self.policy_model), space=self.action_space)
+        policy = SoftDeterministicPolicy(
+            copy.deepcopy(self.policy_model), space=self.action_space
+        )
         return TimeFeature(SACTestAgent(policy))
 
 
-sac = PresetBuilder('sac', default_hyperparameters, SACContinuousPreset)
+sac = PresetBuilder("sac", default_hyperparameters, SACContinuousPreset)
diff --git a/all/presets/independent_multiagent.py b/all/presets/independent_multiagent.py
index 4eef0bd5..841045bd 100644
--- a/all/presets/independent_multiagent.py
+++ b/all/presets/independent_multiagent.py
@@ -1,21 +1,25 @@
-from .builder import PresetBuilder
-from .preset import Preset
 from all.agents import IndependentMultiagent
 from all.logging import DummyLogger
 
+from .preset import Preset
+
 
 class IndependentMultiagentPreset(Preset):
     def __init__(self, name, device, presets):
         super().__init__(name, device, presets)
 
-    def agent(self, logger=DummyLogger(), train_steps=float('inf')):
-        return IndependentMultiagent({
-            agent_id: preset.agent(logger=logger, train_steps=train_steps)
-            for agent_id, preset in self.hyperparameters.items()
-        })
+    def agent(self, logger=DummyLogger(), train_steps=float("inf")):
+        return IndependentMultiagent(
+            {
+                agent_id: preset.agent(logger=logger, train_steps=train_steps)
+                for agent_id, preset in self.hyperparameters.items()
+            }
+        )
 
     def test_agent(self):
-        return IndependentMultiagent({
-            agent_id: preset.test_agent()
-            for agent_id, preset in self.hyperparameters.items()
-        })
+        return IndependentMultiagent(
+            {
+                agent_id: preset.test_agent()
+                for agent_id, preset in self.hyperparameters.items()
+            }
+        )
diff --git a/all/presets/preset.py b/all/presets/preset.py
index 924933b5..58ef6e6f 100644
--- a/all/presets/preset.py
+++ b/all/presets/preset.py
@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+
 import torch
 
 
@@ -18,7 +19,7 @@ def __init__(self, name, device, hyperparameters):
         self.hyperparameters = hyperparameters
 
     @abstractmethod
-    def agent(self, logger=None, train_steps=float('inf')):
+    def agent(self, logger=None, train_steps=float("inf")):
         """
         Instantiate a training-mode Agent with the existing model.
 
@@ -71,7 +72,7 @@ def __init__(self, name, device, hyperparameters):
         self.hyperparameters = hyperparameters
 
     @abstractmethod
-    def agent(self, logger=None, train_steps=float('inf')):
+    def agent(self, logger=None, train_steps=float("inf")):
         """
         Instantiate a training-mode ParallelAgent with the existing model.
 
@@ -108,7 +109,7 @@ def parallel_test_agent(self):
 
     @property
     def n_envs(self):
-        return self.hyperparameters['n_envs']
+        return self.hyperparameters["n_envs"]
 
     def save(self, filename):
         """
diff --git a/all/scripts/atari.py b/all/scripts/atari.py
index da98f7ae..804f46d6 100644
--- a/all/scripts/atari.py
+++ b/all/scripts/atari.py
@@ -1,4 +1,5 @@
 import argparse
+
 from all.environments import AtariEnvironment
 from all.experiments import run_experiment
 from all.presets import atari
@@ -21,13 +22,13 @@ def main():
     parser.add_argument(
         "--render", action="store_true", default=False, help="Render the environment."
     )
+    parser.add_argument("--logdir", default="runs", help="The base logging directory.")
     parser.add_argument(
-        "--logdir", default='runs', help="The base logging directory."
-    )
-    parser.add_argument(
-        "--logger", default='tensorboard', help="The backend used for tracking experiment metrics."
+        "--logger",
+        default="tensorboard",
+        help="The backend used for tracking experiment metrics.",
     )
-    parser.add_argument('--hyperparameters', default=[], nargs='*')
+    parser.add_argument("--hyperparameters", default=[], nargs="*")
     args = parser.parse_args()
 
     env = AtariEnvironment(args.env, device=args.device)
@@ -39,7 +40,7 @@ def main():
     # parse hyperparameters
     hyperparameters = {}
     for hp in args.hyperparameters:
-        key, value = hp.split('=')
+        key, value = hp.split("=")
         hyperparameters[key] = type(agent.default_hyperparameters[key])(value)
     agent = agent.hyperparameters(**hyperparameters)
 
diff --git a/all/scripts/classic.py b/all/scripts/classic.py
index 0e80e15c..59808c84 100644
--- a/all/scripts/classic.py
+++ b/all/scripts/classic.py
@@ -1,4 +1,5 @@
 import argparse
+
 from all.environments import GymEnvironment
 from all.experiments import run_experiment
 from all.presets import classic_control
@@ -21,15 +22,17 @@ def main():
     parser.add_argument(
         "--render", action="store_true", default=False, help="Render the environment."
     )
+    parser.add_argument("--logdir", default="runs", help="The base logging directory.")
     parser.add_argument(
-        "--logdir", default='runs', help="The base logging directory."
+        "--logger",
+        default="tensorboard",
+        help="The backend used for tracking experiment metrics.",
     )
-    parser.add_argument("--logger", default='tensorboard', help="The backend used for tracking experiment metrics.")
     parser.add_argument(
-        '--hyperparameters',
+        "--hyperparameters",
         default=[],
-        nargs='*',
-        help="Custom hyperparameters, in the format hyperparameter1=value1 hyperparameter2=value2 etc."
+        nargs="*",
+        help="Custom hyperparameters, in the format hyperparameter1=value1 hyperparameter2=value2 etc.",
     )
     args = parser.parse_args()
 
@@ -42,7 +45,7 @@ def main():
     # parse hyperparameters
     hyperparameters = {}
     for hp in args.hyperparameters:
-        key, value = hp.split('=')
+        key, value = hp.split("=")
         hyperparameters[key] = type(agent.default_hyperparameters[key])(value)
     agent = agent.hyperparameters(**hyperparameters)
 
diff --git a/all/scripts/continuous.py b/all/scripts/continuous.py
index c4aa7608..a46b6be3 100644
--- a/all/scripts/continuous.py
+++ b/all/scripts/continuous.py
@@ -1,10 +1,10 @@
 # pylint: disable=unused-import
 import argparse
+
 from all.environments import GymEnvironment, PybulletEnvironment
 from all.experiments import run_experiment
 from all.presets import continuous
 
-
 # see also: PybulletEnvironment.short_names
 ENVS = {
     "mountaincar": "MountainCarContinuous-v0",
@@ -29,21 +29,23 @@ def main():
     parser.add_argument(
         "--render", action="store_true", default=False, help="Render the environment."
     )
+    parser.add_argument("--logdir", default="runs", help="The base logging directory.")
     parser.add_argument(
-        "--logdir", default='runs', help="The base logging directory."
+        "--logger",
+        default="tensorboard",
+        help="The backend used for tracking experiment metrics.",
     )
-    parser.add_argument("--logger", default='tensorboard', help="The backend used for tracking experiment metrics.")
     parser.add_argument(
-        '--hyperparameters',
+        "--hyperparameters",
         default=[],
-        nargs='*',
-        help="Custom hyperparameters, in the format hyperparameter1=value1 hyperparameter2=value2 etc."
+        nargs="*",
+        help="Custom hyperparameters, in the format hyperparameter1=value1 hyperparameter2=value2 etc.",
     )
     args = parser.parse_args()
 
     if args.env in ENVS:
         env = GymEnvironment(ENVS[args.env], device=args.device)
-    elif 'BulletEnv' in args.env or args.env in PybulletEnvironment.short_names:
+    elif "BulletEnv" in args.env or args.env in PybulletEnvironment.short_names:
         env = PybulletEnvironment(args.env, device=args.device)
     else:
         env = GymEnvironment(args.env, device=args.device)
@@ -55,7 +57,7 @@ def main():
     # parse hyperparameters
     hyperparameters = {}
     for hp in args.hyperparameters:
-        key, value = hp.split('=')
+        key, value = hp.split("=")
         hyperparameters[key] = type(agent.default_hyperparameters[key])(value)
     agent = agent.hyperparameters(**hyperparameters)
 
diff --git a/all/scripts/multiagent_atari.py b/all/scripts/multiagent_atari.py
index 6a3f1867..a76acfe7 100644
--- a/all/scripts/multiagent_atari.py
+++ b/all/scripts/multiagent_atari.py
@@ -1,11 +1,11 @@
 import argparse
+
 from all.environments import MultiagentAtariEnv
 from all.experiments.multiagent_env_experiment import MultiagentEnvExperiment
-from all.presets import atari
-from all.presets import IndependentMultiagentPreset
+from all.presets import IndependentMultiagentPreset, atari
 
 
-class DummyEnv():
+class DummyEnv:
     def __init__(self, state_space, action_space):
         self.state_space = state_space
         self.action_space = action_space
@@ -14,9 +14,7 @@ def __init__(self, state_space, action_space):
 def main():
     parser = argparse.ArgumentParser(description="Run an multiagent Atari benchmark.")
     parser.add_argument("env", help="Name of the Atari game (e.g. pong_v2).")
-    parser.add_argument(
-        "agents", nargs='*', help="List of agents."
-    )
+    parser.add_argument("agents", nargs="*", help="List of agents.")
     parser.add_argument(
         "--device",
         default="cuda",
@@ -34,13 +32,17 @@ def main():
         "--render", action="store_true", default=False, help="Render the environment."
     )
     parser.add_argument(
-        "--logger", default='tensorboard', help="The backend used for tracking experiment metrics."
+        "--logger",
+        default="tensorboard",
+        help="The backend used for tracking experiment metrics.",
     )
     args = parser.parse_args()
 
     env = MultiagentAtariEnv(args.env, device=args.device)
 
-    assert len(env.agents) == len(args.agents), f"Must specify {len(env.agents)} agents for this environment."
+    assert len(env.agents) == len(
+        args.agents
+    ), f"Must specify {len(env.agents)} agents for this environment."
 
     presets = {
         agent_id: getattr(atari, agent_type)
@@ -52,7 +54,7 @@ def main():
     }
 
     experiment = MultiagentEnvExperiment(
-        IndependentMultiagentPreset('Independent', args.device, presets),
+        IndependentMultiagentPreset("Independent", args.device, presets),
         env,
         verbose=False,
         render=args.render,
diff --git a/all/scripts/plot.py b/all/scripts/plot.py
index b041657a..0ba4472a 100644
--- a/all/scripts/plot.py
+++ b/all/scripts/plot.py
@@ -1,11 +1,17 @@
 import argparse
+
 from all.experiments import plot_returns_100
 
 
 def main():
     parser = argparse.ArgumentParser(description="Plots the results of experiments.")
-    parser.add_argument("--logdir", help="Output directory", default='runs')
-    parser.add_argument("--timesteps", type=int, default=-1, help="The final point will be fixed to this x-value")
+    parser.add_argument("--logdir", help="Output directory", default="runs")
+    parser.add_argument(
+        "--timesteps",
+        type=int,
+        default=-1,
+        help="The final point will be fixed to this x-value",
+    )
     args = parser.parse_args()
     plot_returns_100(args.logdir, timesteps=args.timesteps)
 
diff --git a/all/scripts/release.py b/all/scripts/release.py
index 475e5b80..01a342f7 100644
--- a/all/scripts/release.py
+++ b/all/scripts/release.py
@@ -1,4 +1,5 @@
-'''Create slurm tasks to run release test suite'''
+"""Create slurm tasks to run release test suite"""
+
 from all.environments import AtariEnvironment, GymEnvironment
 from all.experiments import SlurmExperiment
 from all.presets import atari, classic_control, continuous
@@ -6,7 +7,7 @@
 
 def main():
     # run on gpu
-    device = 'cuda'
+    device = "cuda"
 
     def get_agents(preset):
         agents = [getattr(preset, agent_name) for agent_name in preset.__all__]
@@ -14,29 +15,23 @@ def get_agents(preset):
 
     SlurmExperiment(
         get_agents(atari),
-        AtariEnvironment('Breakout', device=device),
+        AtariEnvironment("Breakout", device=device),
         10e7,
-        sbatch_args={
-            'partition': '1080ti-long'
-        }
+        sbatch_args={"partition": "1080ti-long"},
     )
 
     SlurmExperiment(
         get_agents(classic_control),
-        GymEnvironment('CartPole-v0', device=device),
+        GymEnvironment("CartPole-v0", device=device),
         100000,
-        sbatch_args={
-            'partition': '1080ti-short'
-        }
+        sbatch_args={"partition": "1080ti-short"},
     )
 
     SlurmExperiment(
         get_agents(continuous),
-        GymEnvironment('LunarLanderContinuous-v2', device=device),
+        GymEnvironment("LunarLanderContinuous-v2", device=device),
         500000,
-        sbatch_args={
-            'partition': '1080ti-short'
-        }
+        sbatch_args={"partition": "1080ti-short"},
     )
 
 
diff --git a/all/scripts/watch_atari.py b/all/scripts/watch_atari.py
index b5cd6a28..cd278afc 100644
--- a/all/scripts/watch_atari.py
+++ b/all/scripts/watch_atari.py
@@ -1,5 +1,5 @@
 import argparse
-from all.bodies import DeepmindAtariBody
+
 from all.environments import AtariEnvironment
 from all.experiments import load_and_watch
 
diff --git a/all/scripts/watch_classic.py b/all/scripts/watch_classic.py
index 0b2e7416..239f0964 100644
--- a/all/scripts/watch_classic.py
+++ b/all/scripts/watch_classic.py
@@ -1,11 +1,14 @@
 import argparse
+
 from all.environments import GymEnvironment
 from all.experiments import load_and_watch
 
 
 def main():
     parser = argparse.ArgumentParser(description="Run an Atari benchmark.")
-    parser.add_argument("env", help="Name of the environment (e.g. RoboschoolHalfCheetah-v1")
+    parser.add_argument(
+        "env", help="Name of the environment (e.g. RoboschoolHalfCheetah-v1"
+    )
     parser.add_argument("filename", help="File where the model was saved.")
     parser.add_argument(
         "--device",
diff --git a/all/scripts/watch_continuous.py b/all/scripts/watch_continuous.py
index 2cd17250..903f83d7 100644
--- a/all/scripts/watch_continuous.py
+++ b/all/scripts/watch_continuous.py
@@ -1,8 +1,9 @@
 # pylint: disable=unused-import
 import argparse
-from all.bodies import TimeFeature
+
 from all.environments import GymEnvironment, PybulletEnvironment
 from all.experiments import load_and_watch
+
 from .continuous import ENVS
 
 
@@ -24,7 +25,7 @@ def main():
 
     if args.env in ENVS:
         env = GymEnvironment(args.env, device=args.device, render_mode="human")
-    elif 'BulletEnv' in args.env or args.env in PybulletEnvironment.short_names:
+    elif "BulletEnv" in args.env or args.env in PybulletEnvironment.short_names:
         env = PybulletEnvironment(args.env, device=args.device, render_mode="human")
     else:
         env = GymEnvironment(args.env, device=args.device, render_mode="human")
diff --git a/all/scripts/watch_multiagent_atari.py b/all/scripts/watch_multiagent_atari.py
index fe6095eb..9d16f1c9 100644
--- a/all/scripts/watch_multiagent_atari.py
+++ b/all/scripts/watch_multiagent_atari.py
@@ -1,9 +1,9 @@
 import argparse
 import time
+
 import torch
-from all.bodies import DeepmindAtariBody
+
 from all.environments import MultiagentAtariEnv
-from all.experiments import load_and_watch
 
 
 def watch(env, filename, fps, reload):
@@ -15,7 +15,7 @@ def watch(env, filename, fps, reload):
             try:
                 agent = torch.load(filename).test_agent()
             except Exception as e:
-                print('Warning: error reloading model: {}'.format(filename))
+                print("Warning: error reloading model: {}".format(filename))
                 print(e)
 
 
@@ -49,7 +49,9 @@ def main():
     )
     parser.add_argument(
         "--reload",
-        action="store_true", default=False, help="Reload the model from disk after every episode"
+        action="store_true",
+        default=False,
+        help="Reload the model from disk after every episode",
     )
     args = parser.parse_args()
     env = MultiagentAtariEnv(args.env, device=args.device, render_mode="human")
diff --git a/benchmarks/atari40.py b/benchmarks/atari40.py
index 4a1ffdae..3dc88c7c 100644
--- a/benchmarks/atari40.py
+++ b/benchmarks/atari40.py
@@ -1,6 +1,6 @@
+from all.environments import AtariEnvironment
 from all.experiments import SlurmExperiment
 from all.presets import atari
-from all.environments import AtariEnvironment
 
 
 def main():
@@ -12,10 +12,17 @@ def main():
         atari.ppo,
         atari.rainbow,
     ]
-    envs = [AtariEnvironment(env, device='cuda') for env in ['BeamRider', 'Breakout', 'Pong', 'Qbert', 'SpaceInvaders']]
-    SlurmExperiment(agents, envs, 10e6, logdir='benchmarks/atari40', sbatch_args={
-        'partition': 'gpu-long'
-    })
+    envs = [
+        AtariEnvironment(env, device="cuda")
+        for env in ["BeamRider", "Breakout", "Pong", "Qbert", "SpaceInvaders"]
+    ]
+    SlurmExperiment(
+        agents,
+        envs,
+        10e6,
+        logdir="benchmarks/atari40",
+        sbatch_args={"partition": "gpu-long"},
+    )
 
 
 if __name__ == "__main__":
diff --git a/benchmarks/pybullet.py b/benchmarks/pybullet.py
index 0762add8..dc045fd6 100644
--- a/benchmarks/pybullet.py
+++ b/benchmarks/pybullet.py
@@ -1,22 +1,25 @@
+from all.environments import PybulletEnvironment
 from all.experiments import SlurmExperiment
 from all.presets.continuous import ddpg, ppo, sac
-from all.environments import PybulletEnvironment
 
 
 def main():
     frames = int(1e7)
 
-    agents = [
-        ddpg,
-        ppo,
-        sac
-    ]
+    agents = [ddpg, ppo, sac]
 
-    envs = [PybulletEnvironment(env, device='cuda') for env in PybulletEnvironment.short_names]
+    envs = [
+        PybulletEnvironment(env, device="cuda")
+        for env in PybulletEnvironment.short_names
+    ]
 
-    SlurmExperiment(agents, envs, frames, logdir='benchmarks/pybullet', sbatch_args={
-        'partition': 'gpu-long'
-    })
+    SlurmExperiment(
+        agents,
+        envs,
+        frames,
+        logdir="benchmarks/pybullet",
+        sbatch_args={"partition": "gpu-long"},
+    )
 
 
 if __name__ == "__main__":
diff --git a/examples/experiment.py b/examples/experiment.py
index 9492aea4..e03d1398 100644
--- a/examples/experiment.py
+++ b/examples/experiment.py
@@ -1,13 +1,14 @@
-'''
+"""
 Quick example of usage of the run_experiment API.
-'''
-from all.experiments import run_experiment, plot_returns_100
-from all.presets.classic_control import dqn, a2c
+"""
+
 from all.environments import GymEnvironment
+from all.experiments import plot_returns_100, run_experiment
+from all.presets.classic_control import a2c, dqn
 
 
 def main():
-    DEVICE = 'cpu'
+    DEVICE = "cpu"
     # DEVICE = 'cuda' # uncomment for gpu support
     timesteps = 40000
     run_experiment(
@@ -15,14 +16,16 @@ def main():
             # DQN with default hyperparameters
             dqn.device(DEVICE),
             # DQN with a custom hyperparameters and a custom name.
-            dqn.device(DEVICE).hyperparameters(replay_buffer_size=100).name('dqn-small-buffer'),
+            dqn.device(DEVICE)
+            .hyperparameters(replay_buffer_size=100)
+            .name("dqn-small-buffer"),
             # A2C with a custom name
-            a2c.device(DEVICE).name('not-dqn')
+            a2c.device(DEVICE).name("not-dqn"),
         ],
-        [GymEnvironment('CartPole-v0', DEVICE), GymEnvironment('Acrobot-v1', DEVICE)],
+        [GymEnvironment("CartPole-v0", DEVICE), GymEnvironment("Acrobot-v1", DEVICE)],
         timesteps,
     )
-    plot_returns_100('runs', timesteps=timesteps)
+    plot_returns_100("runs", timesteps=timesteps)
 
 
 if __name__ == "__main__":
diff --git a/examples/slurm_experiment.py b/examples/slurm_experiment.py
index 155fdb5b..f938cc0b 100644
--- a/examples/slurm_experiment.py
+++ b/examples/slurm_experiment.py
@@ -1,19 +1,25 @@
-'''
+"""
 Quick example of a2c running on slurm, a distributed cluster.
 Note that it only runs for 1 million frames.
 For real experiments, you will surely need a modified version of this script.
-'''
+"""
+
+from all.environments import AtariEnvironment
 from all.experiments import SlurmExperiment
 from all.presets.atari import a2c, dqn
-from all.environments import AtariEnvironment
 
 
 def main():
-    device = 'cuda'
-    envs = [AtariEnvironment(env, device) for env in ['Pong', 'Breakout', 'SpaceInvaders']]
-    SlurmExperiment([a2c.device(device), dqn.device(device)], envs, 1e6, sbatch_args={
-        'partition': '1080ti-short'
-    })
+    device = "cuda"
+    envs = [
+        AtariEnvironment(env, device) for env in ["Pong", "Breakout", "SpaceInvaders"]
+    ]
+    SlurmExperiment(
+        [a2c.device(device), dqn.device(device)],
+        envs,
+        1e6,
+        sbatch_args={"partition": "1080ti-short"},
+    )
 
 
 if __name__ == "__main__":
diff --git a/integration/validate_agent.py b/integration/validate_agent.py
index d0b753bd..1b786efe 100644
--- a/integration/validate_agent.py
+++ b/integration/validate_agent.py
@@ -1,7 +1,12 @@
 import os
+
+from all.experiments import (
+    MultiagentEnvExperiment,
+    ParallelEnvExperiment,
+    SingleEnvExperiment,
+)
 from all.logging import DummyLogger
-from all.experiments import SingleEnvExperiment, ParallelEnvExperiment, MultiagentEnvExperiment
-from all.presets import ParallelPreset, Preset
+from all.presets import ParallelPreset
 
 
 class TestSingleEnvExperiment(SingleEnvExperiment):
diff --git a/setup.py b/setup.py
index 42f5b3a7..3c136805 100644
--- a/setup.py
+++ b/setup.py
@@ -1,5 +1,4 @@
-from setuptools import setup, find_packages
-
+from setuptools import find_packages, setup
 
 GYM_VERSION = "0.29.1"
 PETTINGZOO_VERSION = "1.24.2"
@@ -24,8 +23,9 @@
         "supersuit~=3.9.1",
     ],
     "test": [
-        "flake8>=3.8",  # linter for pep8 compliance
-        "autopep8>=1.5",  # automatically fixes some pep8 errors
+        "black>=24.1.1",  # linting/formatting
+        "isort>=5.13.2",  # sort imports
+        "flake8>=7.0.0",  # more linting
         "torch-testing>=0.0.2",  # pytorch assertion library
     ],
     "docs": [
@@ -71,13 +71,13 @@
         ],
     },
     install_requires=[
-        "gymnasium~={}".format(GYM_VERSION),    # common environment interface
-        "numpy>=1.22.3",                        # math library
-        "matplotlib>=3.5.1",                    # plotting library
-        "opencv-python-headless>=4.0.0",        # used by atari wrappers
-        "torch>=1.11.0",                        # core deep learning library
-        "tensorboard>=2.8.0",                   # logging and visualization
-        "cloudpickle>=2.0.0",                   # used to copy environments
+        "gymnasium~={}".format(GYM_VERSION),  # common environment interface
+        "numpy>=1.22.3",  # math library
+        "matplotlib>=3.5.1",  # plotting library
+        "opencv-python-headless>=4.0.0",  # used by atari wrappers
+        "torch>=1.11.0",  # core deep learning library
+        "tensorboard>=2.8.0",  # logging and visualization
+        "cloudpickle>=2.0.0",  # used to copy environments
     ],
     extras_require=extras,
 )


Test patch:
diff --git a/all/approximation/feature_network_test.py b/all/approximation/feature_network_test.py
index e8190e99..8a76b42d 100644
--- a/all/approximation/feature_network_test.py
+++ b/all/approximation/feature_network_test.py
@@ -1,10 +1,11 @@
 import unittest
+
 import torch
-from torch import nn
 import torch_testing as tt
-from all.core import State
-from all.approximation.feature_network import FeatureNetwork
+from torch import nn
 
+from all.approximation.feature_network import FeatureNetwork
+from all.core import State
 
 STATE_DIM = 2
 
@@ -16,20 +17,21 @@ def setUp(self):
 
         optimizer = torch.optim.SGD(self.model.parameters(), lr=0.1)
         self.features = FeatureNetwork(self.model, optimizer)
-        self.states = State({
-            'observation': torch.randn(3, STATE_DIM),
-            'mask': torch.tensor([1, 0, 1])
-        })
-        self.expected_features = State({
-            'observation': torch.tensor(
-                [
-                    [-0.2385, -0.7263, -0.0340],
-                    [-0.3569, -0.6612, 0.3485],
-                    [-0.0296, -0.7566, -0.4624],
-                ]
-            ),
-            'mask': torch.tensor([1, 0, 1])
-        })
+        self.states = State(
+            {"observation": torch.randn(3, STATE_DIM), "mask": torch.tensor([1, 0, 1])}
+        )
+        self.expected_features = State(
+            {
+                "observation": torch.tensor(
+                    [
+                        [-0.2385, -0.7263, -0.0340],
+                        [-0.3569, -0.6612, 0.3485],
+                        [-0.0296, -0.7566, -0.4624],
+                    ]
+                ),
+                "mask": torch.tensor([1, 0, 1]),
+            }
+        )
 
     def test_forward(self):
         features = self.features(self.states)
@@ -41,14 +43,14 @@ def test_backward(self):
         loss = torch.sum(states.observation)
         self.features.reinforce(loss)
         features = self.features(self.states)
-        expected = State({
-            'observation': torch.tensor([
-                [-0.71, -1.2, -0.5],
-                [-0.72, -1.03, -0.02],
-                [-0.57, -1.3, -1.01]
-            ]),
-            'mask': torch.tensor([1, 0, 1]),
-        })
+        expected = State(
+            {
+                "observation": torch.tensor(
+                    [[-0.71, -1.2, -0.5], [-0.72, -1.03, -0.02], [-0.57, -1.3, -1.01]]
+                ),
+                "mask": torch.tensor([1, 0, 1]),
+            }
+        )
         self.assert_state_equal(features, expected)
 
     def test_eval(self):
diff --git a/all/approximation/identity_test.py b/all/approximation/identity_test.py
index 7af58afd..a5ef493b 100644
--- a/all/approximation/identity_test.py
+++ b/all/approximation/identity_test.py
@@ -1,13 +1,15 @@
 import unittest
+
 import torch
 import torch_testing as tt
+
+from all.approximation import FixedTarget, Identity
 from all.core import State
-from all.approximation import Identity, FixedTarget
 
 
 class TestIdentityNetwork(unittest.TestCase):
     def setUp(self):
-        self.model = Identity('cpu', target=FixedTarget(10))
+        self.model = Identity("cpu", target=FixedTarget(10))
 
     def test_forward_tensor(self):
         inputs = torch.tensor([1, 2, 3])
@@ -15,9 +17,7 @@ def test_forward_tensor(self):
         tt.assert_equal(inputs, outputs)
 
     def test_forward_state(self):
-        inputs = State({
-            'observation': torch.tensor([1, 2, 3])
-        })
+        inputs = State({"observation": torch.tensor([1, 2, 3])})
         outputs = self.model(inputs)
         self.assertEqual(inputs, outputs)
 
diff --git a/all/approximation/q_dist_test.py b/all/approximation/q_dist_test.py
index 778db4ee..06df95a4 100644
--- a/all/approximation/q_dist_test.py
+++ b/all/approximation/q_dist_test.py
@@ -1,9 +1,11 @@
 import unittest
+
 import torch
-from torch import nn
 import torch_testing as tt
-from all.core import StateArray
+from torch import nn
+
 from all.approximation import QDist
+from all.core import StateArray
 
 STATE_DIM = 1
 ACTIONS = 2
@@ -73,7 +75,9 @@ def test_single_q_values(self):
         )
 
     def test_done(self):
-        states = StateArray(torch.randn((3, STATE_DIM)), (3,), mask=torch.tensor([1, 0, 1]))
+        states = StateArray(
+            torch.randn((3, STATE_DIM)), (3,), mask=torch.tensor([1, 0, 1])
+        )
         probs = self.q(states)
         self.assertEqual(probs.shape, (3, ACTIONS, ATOMS))
         tt.assert_almost_equal(
@@ -134,141 +138,997 @@ def _loss(dist, target_dist):
     def test_project_dist(self):
         # This gave problems in the past between different cuda version,
         # so a test was added.
-        q = QDist(self.model, self.optimizer, ACTIONS, 51, -10., 10.)
-        dist = torch.tensor([
-            [0.0190, 0.0197, 0.0200, 0.0190, 0.0195, 0.0198, 0.0194, 0.0192, 0.0201,
-            0.0203, 0.0189, 0.0190, 0.0199, 0.0193, 0.0192, 0.0199, 0.0198, 0.0197,
-            0.0193, 0.0198, 0.0192, 0.0191, 0.0200, 0.0202, 0.0191, 0.0202, 0.0198,
-            0.0200, 0.0198, 0.0193, 0.0192, 0.0202, 0.0192, 0.0194, 0.0199, 0.0197,
-            0.0197, 0.0201, 0.0199, 0.0190, 0.0192, 0.0195, 0.0202, 0.0194, 0.0203,
-            0.0201, 0.0190, 0.0192, 0.0201, 0.0201, 0.0192],
-            [0.0191, 0.0197, 0.0200, 0.0190, 0.0195, 0.0198, 0.0194, 0.0192, 0.0201,
-            0.0203, 0.0190, 0.0190, 0.0199, 0.0193, 0.0192, 0.0199, 0.0198, 0.0197,
-            0.0193, 0.0198, 0.0192, 0.0191, 0.0200, 0.0202, 0.0191, 0.0202, 0.0198,
-            0.0200, 0.0198, 0.0193, 0.0192, 0.0202, 0.0192, 0.0194, 0.0199, 0.0197,
-            0.0197, 0.0200, 0.0199, 0.0190, 0.0192, 0.0195, 0.0202, 0.0194, 0.0203,
-            0.0201, 0.0190, 0.0192, 0.0201, 0.0200, 0.0192],
-            [0.0191, 0.0197, 0.0200, 0.0190, 0.0195, 0.0198, 0.0194, 0.0192, 0.0200,
-            0.0203, 0.0190, 0.0191, 0.0199, 0.0193, 0.0192, 0.0199, 0.0198, 0.0197,
-            0.0193, 0.0198, 0.0192, 0.0191, 0.0199, 0.0202, 0.0192, 0.0202, 0.0198,
-            0.0200, 0.0198, 0.0193, 0.0192, 0.0202, 0.0192, 0.0194, 0.0199, 0.0197,
-            0.0197, 0.0200, 0.0199, 0.0190, 0.0192, 0.0195, 0.0202, 0.0194, 0.0203,
-            0.0201, 0.0190, 0.0192, 0.0201, 0.0200, 0.0192]
-        ])
-        support = torch.tensor([
-            [-9.7030, -9.3149, -8.9268, -8.5386, -8.1505, -7.7624, -7.3743, -6.9862,
-            -6.5980, -6.2099, -5.8218, -5.4337, -5.0456, -4.6574, -4.2693, -3.8812,
-            -3.4931, -3.1050, -2.7168, -2.3287, -1.9406, -1.5525, -1.1644, -0.7762,
-            -0.3881, 0.0000, 0.3881, 0.7762, 1.1644, 1.5525, 1.9406, 2.3287,
-            2.7168, 3.1050, 3.4931, 3.8812, 4.2693, 4.6574, 5.0456, 5.4337,
-            5.8218, 6.2099, 6.5980, 6.9862, 7.3743, 7.7624, 8.1505, 8.5386,
-            8.9268, 9.3149, 9.7030],
-            [-9.7030, -9.3149, -8.9268, -8.5386, -8.1505, -7.7624, -7.3743, -6.9862,
-            -6.5980, -6.2099, -5.8218, -5.4337, -5.0456, -4.6574, -4.2693, -3.8812,
-            -3.4931, -3.1050, -2.7168, -2.3287, -1.9406, -1.5525, -1.1644, -0.7762,
-            -0.3881, 0.0000, 0.3881, 0.7762, 1.1644, 1.5525, 1.9406, 2.3287,
-            2.7168, 3.1050, 3.4931, 3.8812, 4.2693, 4.6574, 5.0456, 5.4337,
-            5.8218, 6.2099, 6.5980, 6.9862, 7.3743, 7.7624, 8.1505, 8.5386,
-            8.9268, 9.3149, 9.7030],
-            [-9.7030, -9.3149, -8.9268, -8.5386, -8.1505, -7.7624, -7.3743, -6.9862,
-            -6.5980, -6.2099, -5.8218, -5.4337, -5.0456, -4.6574, -4.2693, -3.8812,
-            -3.4931, -3.1050, -2.7168, -2.3287, -1.9406, -1.5525, -1.1644, -0.7762,
-            -0.3881, 0.0000, 0.3881, 0.7762, 1.1644, 1.5525, 1.9406, 2.3287,
-            2.7168, 3.1050, 3.4931, 3.8812, 4.2693, 4.6574, 5.0456, 5.4337,
-            5.8218, 6.2099, 6.5980, 6.9862, 7.3743, 7.7624, 8.1505, 8.5386,
-            8.9268, 9.3149, 9.7030]
-        ])
-        expected = torch.tensor([
-            [0.0049, 0.0198, 0.0204, 0.0202, 0.0198, 0.0202, 0.0202, 0.0199, 0.0202,
-            0.0208, 0.0201, 0.0195, 0.0201, 0.0201, 0.0198, 0.0203, 0.0204, 0.0203,
-            0.0200, 0.0203, 0.0199, 0.0197, 0.0205, 0.0208, 0.0197, 0.0214, 0.0204,
-            0.0206, 0.0203, 0.0199, 0.0199, 0.0206, 0.0198, 0.0201, 0.0204, 0.0203,
-            0.0204, 0.0206, 0.0201, 0.0197, 0.0199, 0.0204, 0.0204, 0.0205, 0.0208,
-            0.0200, 0.0197, 0.0204, 0.0207, 0.0200, 0.0049],
-            [0.0049, 0.0198, 0.0204, 0.0202, 0.0198, 0.0202, 0.0202, 0.0199, 0.0202,
-            0.0208, 0.0202, 0.0196, 0.0201, 0.0201, 0.0198, 0.0203, 0.0204, 0.0203,
-            0.0200, 0.0203, 0.0199, 0.0197, 0.0205, 0.0208, 0.0197, 0.0214, 0.0204,
-            0.0206, 0.0203, 0.0199, 0.0199, 0.0206, 0.0198, 0.0201, 0.0204, 0.0203,
-            0.0204, 0.0206, 0.0201, 0.0197, 0.0199, 0.0204, 0.0204, 0.0205, 0.0208,
-            0.0200, 0.0197, 0.0204, 0.0206, 0.0200, 0.0049],
-            [0.0049, 0.0198, 0.0204, 0.0202, 0.0198, 0.0202, 0.0202, 0.0199, 0.0202,
-            0.0208, 0.0202, 0.0196, 0.0202, 0.0201, 0.0198, 0.0203, 0.0204, 0.0203,
-            0.0200, 0.0203, 0.0199, 0.0197, 0.0204, 0.0208, 0.0198, 0.0214, 0.0204,
-            0.0206, 0.0203, 0.0199, 0.0199, 0.0206, 0.0198, 0.0201, 0.0204, 0.0203,
-            0.0204, 0.0206, 0.0201, 0.0197, 0.0199, 0.0204, 0.0204, 0.0205, 0.0208,
-            0.0200, 0.0197, 0.0204, 0.0206, 0.0200, 0.0049]
-        ])
-        tt.assert_almost_equal(q.project(dist, support).cpu(), expected.cpu(), decimal=3)
+        q = QDist(self.model, self.optimizer, ACTIONS, 51, -10.0, 10.0)
+        dist = torch.tensor(
+            [
+                [
+                    0.0190,
+                    0.0197,
+                    0.0200,
+                    0.0190,
+                    0.0195,
+                    0.0198,
+                    0.0194,
+                    0.0192,
+                    0.0201,
+                    0.0203,
+                    0.0189,
+                    0.0190,
+                    0.0199,
+                    0.0193,
+                    0.0192,
+                    0.0199,
+                    0.0198,
+                    0.0197,
+                    0.0193,
+                    0.0198,
+                    0.0192,
+                    0.0191,
+                    0.0200,
+                    0.0202,
+                    0.0191,
+                    0.0202,
+                    0.0198,
+                    0.0200,
+                    0.0198,
+                    0.0193,
+                    0.0192,
+                    0.0202,
+                    0.0192,
+                    0.0194,
+                    0.0199,
+                    0.0197,
+                    0.0197,
+                    0.0201,
+                    0.0199,
+                    0.0190,
+                    0.0192,
+                    0.0195,
+                    0.0202,
+                    0.0194,
+                    0.0203,
+                    0.0201,
+                    0.0190,
+                    0.0192,
+                    0.0201,
+                    0.0201,
+                    0.0192,
+                ],
+                [
+                    0.0191,
+                    0.0197,
+                    0.0200,
+                    0.0190,
+                    0.0195,
+                    0.0198,
+                    0.0194,
+                    0.0192,
+                    0.0201,
+                    0.0203,
+                    0.0190,
+                    0.0190,
+                    0.0199,
+                    0.0193,
+                    0.0192,
+                    0.0199,
+                    0.0198,
+                    0.0197,
+                    0.0193,
+                    0.0198,
+                    0.0192,
+                    0.0191,
+                    0.0200,
+                    0.0202,
+                    0.0191,
+                    0.0202,
+                    0.0198,
+                    0.0200,
+                    0.0198,
+                    0.0193,
+                    0.0192,
+                    0.0202,
+                    0.0192,
+                    0.0194,
+                    0.0199,
+                    0.0197,
+                    0.0197,
+                    0.0200,
+                    0.0199,
+                    0.0190,
+                    0.0192,
+                    0.0195,
+                    0.0202,
+                    0.0194,
+                    0.0203,
+                    0.0201,
+                    0.0190,
+                    0.0192,
+                    0.0201,
+                    0.0200,
+                    0.0192,
+                ],
+                [
+                    0.0191,
+                    0.0197,
+                    0.0200,
+                    0.0190,
+                    0.0195,
+                    0.0198,
+                    0.0194,
+                    0.0192,
+                    0.0200,
+                    0.0203,
+                    0.0190,
+                    0.0191,
+                    0.0199,
+                    0.0193,
+                    0.0192,
+                    0.0199,
+                    0.0198,
+                    0.0197,
+                    0.0193,
+                    0.0198,
+                    0.0192,
+                    0.0191,
+                    0.0199,
+                    0.0202,
+                    0.0192,
+                    0.0202,
+                    0.0198,
+                    0.0200,
+                    0.0198,
+                    0.0193,
+                    0.0192,
+                    0.0202,
+                    0.0192,
+                    0.0194,
+                    0.0199,
+                    0.0197,
+                    0.0197,
+                    0.0200,
+                    0.0199,
+                    0.0190,
+                    0.0192,
+                    0.0195,
+                    0.0202,
+                    0.0194,
+                    0.0203,
+                    0.0201,
+                    0.0190,
+                    0.0192,
+                    0.0201,
+                    0.0200,
+                    0.0192,
+                ],
+            ]
+        )
+        support = torch.tensor(
+            [
+                [
+                    -9.7030,
+                    -9.3149,
+                    -8.9268,
+                    -8.5386,
+                    -8.1505,
+                    -7.7624,
+                    -7.3743,
+                    -6.9862,
+                    -6.5980,
+                    -6.2099,
+                    -5.8218,
+                    -5.4337,
+                    -5.0456,
+                    -4.6574,
+                    -4.2693,
+                    -3.8812,
+                    -3.4931,
+                    -3.1050,
+                    -2.7168,
+                    -2.3287,
+                    -1.9406,
+                    -1.5525,
+                    -1.1644,
+                    -0.7762,
+                    -0.3881,
+                    0.0000,
+                    0.3881,
+                    0.7762,
+                    1.1644,
+                    1.5525,
+                    1.9406,
+                    2.3287,
+                    2.7168,
+                    3.1050,
+                    3.4931,
+                    3.8812,
+                    4.2693,
+                    4.6574,
+                    5.0456,
+                    5.4337,
+                    5.8218,
+                    6.2099,
+                    6.5980,
+                    6.9862,
+                    7.3743,
+                    7.7624,
+                    8.1505,
+                    8.5386,
+                    8.9268,
+                    9.3149,
+                    9.7030,
+                ],
+                [
+                    -9.7030,
+                    -9.3149,
+                    -8.9268,
+                    -8.5386,
+                    -8.1505,
+                    -7.7624,
+                    -7.3743,
+                    -6.9862,
+                    -6.5980,
+                    -6.2099,
+                    -5.8218,
+                    -5.4337,
+                    -5.0456,
+                    -4.6574,
+                    -4.2693,
+                    -3.8812,
+                    -3.4931,
+                    -3.1050,
+                    -2.7168,
+                    -2.3287,
+                    -1.9406,
+                    -1.5525,
+                    -1.1644,
+                    -0.7762,
+                    -0.3881,
+                    0.0000,
+                    0.3881,
+                    0.7762,
+                    1.1644,
+                    1.5525,
+                    1.9406,
+                    2.3287,
+                    2.7168,
+                    3.1050,
+                    3.4931,
+                    3.8812,
+                    4.2693,
+                    4.6574,
+                    5.0456,
+                    5.4337,
+                    5.8218,
+                    6.2099,
+                    6.5980,
+                    6.9862,
+                    7.3743,
+                    7.7624,
+                    8.1505,
+                    8.5386,
+                    8.9268,
+                    9.3149,
+                    9.7030,
+                ],
+                [
+                    -9.7030,
+                    -9.3149,
+                    -8.9268,
+                    -8.5386,
+                    -8.1505,
+                    -7.7624,
+                    -7.3743,
+                    -6.9862,
+                    -6.5980,
+                    -6.2099,
+                    -5.8218,
+                    -5.4337,
+                    -5.0456,
+                    -4.6574,
+                    -4.2693,
+                    -3.8812,
+                    -3.4931,
+                    -3.1050,
+                    -2.7168,
+                    -2.3287,
+                    -1.9406,
+                    -1.5525,
+                    -1.1644,
+                    -0.7762,
+                    -0.3881,
+                    0.0000,
+                    0.3881,
+                    0.7762,
+                    1.1644,
+                    1.5525,
+                    1.9406,
+                    2.3287,
+                    2.7168,
+                    3.1050,
+                    3.4931,
+                    3.8812,
+                    4.2693,
+                    4.6574,
+                    5.0456,
+                    5.4337,
+                    5.8218,
+                    6.2099,
+                    6.5980,
+                    6.9862,
+                    7.3743,
+                    7.7624,
+                    8.1505,
+                    8.5386,
+                    8.9268,
+                    9.3149,
+                    9.7030,
+                ],
+            ]
+        )
+        expected = torch.tensor(
+            [
+                [
+                    0.0049,
+                    0.0198,
+                    0.0204,
+                    0.0202,
+                    0.0198,
+                    0.0202,
+                    0.0202,
+                    0.0199,
+                    0.0202,
+                    0.0208,
+                    0.0201,
+                    0.0195,
+                    0.0201,
+                    0.0201,
+                    0.0198,
+                    0.0203,
+                    0.0204,
+                    0.0203,
+                    0.0200,
+                    0.0203,
+                    0.0199,
+                    0.0197,
+                    0.0205,
+                    0.0208,
+                    0.0197,
+                    0.0214,
+                    0.0204,
+                    0.0206,
+                    0.0203,
+                    0.0199,
+                    0.0199,
+                    0.0206,
+                    0.0198,
+                    0.0201,
+                    0.0204,
+                    0.0203,
+                    0.0204,
+                    0.0206,
+                    0.0201,
+                    0.0197,
+                    0.0199,
+                    0.0204,
+                    0.0204,
+                    0.0205,
+                    0.0208,
+                    0.0200,
+                    0.0197,
+                    0.0204,
+                    0.0207,
+                    0.0200,
+                    0.0049,
+                ],
+                [
+                    0.0049,
+                    0.0198,
+                    0.0204,
+                    0.0202,
+                    0.0198,
+                    0.0202,
+                    0.0202,
+                    0.0199,
+                    0.0202,
+                    0.0208,
+                    0.0202,
+                    0.0196,
+                    0.0201,
+                    0.0201,
+                    0.0198,
+                    0.0203,
+                    0.0204,
+                    0.0203,
+                    0.0200,
+                    0.0203,
+                    0.0199,
+                    0.0197,
+                    0.0205,
+                    0.0208,
+                    0.0197,
+                    0.0214,
+                    0.0204,
+                    0.0206,
+                    0.0203,
+                    0.0199,
+                    0.0199,
+                    0.0206,
+                    0.0198,
+                    0.0201,
+                    0.0204,
+                    0.0203,
+                    0.0204,
+                    0.0206,
+                    0.0201,
+                    0.0197,
+                    0.0199,
+                    0.0204,
+                    0.0204,
+                    0.0205,
+                    0.0208,
+                    0.0200,
+                    0.0197,
+                    0.0204,
+                    0.0206,
+                    0.0200,
+                    0.0049,
+                ],
+                [
+                    0.0049,
+                    0.0198,
+                    0.0204,
+                    0.0202,
+                    0.0198,
+                    0.0202,
+                    0.0202,
+                    0.0199,
+                    0.0202,
+                    0.0208,
+                    0.0202,
+                    0.0196,
+                    0.0202,
+                    0.0201,
+                    0.0198,
+                    0.0203,
+                    0.0204,
+                    0.0203,
+                    0.0200,
+                    0.0203,
+                    0.0199,
+                    0.0197,
+                    0.0204,
+                    0.0208,
+                    0.0198,
+                    0.0214,
+                    0.0204,
+                    0.0206,
+                    0.0203,
+                    0.0199,
+                    0.0199,
+                    0.0206,
+                    0.0198,
+                    0.0201,
+                    0.0204,
+                    0.0203,
+                    0.0204,
+                    0.0206,
+                    0.0201,
+                    0.0197,
+                    0.0199,
+                    0.0204,
+                    0.0204,
+                    0.0205,
+                    0.0208,
+                    0.0200,
+                    0.0197,
+                    0.0204,
+                    0.0206,
+                    0.0200,
+                    0.0049,
+                ],
+            ]
+        )
+        tt.assert_almost_equal(
+            q.project(dist, support).cpu(), expected.cpu(), decimal=3
+        )
 
     def test_project_dist_cuda(self):
         if torch.cuda.is_available():
             # This gave problems in the past between different cuda version,
             # so a test was added.
-            q = QDist(self.model.cuda(), self.optimizer, ACTIONS, 51, -10., 10.)
-            dist = torch.tensor([
-                [0.0190, 0.0197, 0.0200, 0.0190, 0.0195, 0.0198, 0.0194, 0.0192, 0.0201,
-                0.0203, 0.0189, 0.0190, 0.0199, 0.0193, 0.0192, 0.0199, 0.0198, 0.0197,
-                0.0193, 0.0198, 0.0192, 0.0191, 0.0200, 0.0202, 0.0191, 0.0202, 0.0198,
-                0.0200, 0.0198, 0.0193, 0.0192, 0.0202, 0.0192, 0.0194, 0.0199, 0.0197,
-                0.0197, 0.0201, 0.0199, 0.0190, 0.0192, 0.0195, 0.0202, 0.0194, 0.0203,
-                0.0201, 0.0190, 0.0192, 0.0201, 0.0201, 0.0192],
-                [0.0191, 0.0197, 0.0200, 0.0190, 0.0195, 0.0198, 0.0194, 0.0192, 0.0201,
-                0.0203, 0.0190, 0.0190, 0.0199, 0.0193, 0.0192, 0.0199, 0.0198, 0.0197,
-                0.0193, 0.0198, 0.0192, 0.0191, 0.0200, 0.0202, 0.0191, 0.0202, 0.0198,
-                0.0200, 0.0198, 0.0193, 0.0192, 0.0202, 0.0192, 0.0194, 0.0199, 0.0197,
-                0.0197, 0.0200, 0.0199, 0.0190, 0.0192, 0.0195, 0.0202, 0.0194, 0.0203,
-                0.0201, 0.0190, 0.0192, 0.0201, 0.0200, 0.0192],
-                [0.0191, 0.0197, 0.0200, 0.0190, 0.0195, 0.0198, 0.0194, 0.0192, 0.0200,
-                0.0203, 0.0190, 0.0191, 0.0199, 0.0193, 0.0192, 0.0199, 0.0198, 0.0197,
-                0.0193, 0.0198, 0.0192, 0.0191, 0.0199, 0.0202, 0.0192, 0.0202, 0.0198,
-                0.0200, 0.0198, 0.0193, 0.0192, 0.0202, 0.0192, 0.0194, 0.0199, 0.0197,
-                0.0197, 0.0200, 0.0199, 0.0190, 0.0192, 0.0195, 0.0202, 0.0194, 0.0203,
-                0.0201, 0.0190, 0.0192, 0.0201, 0.0200, 0.0192]
-            ]).cuda()
-            support = torch.tensor([
-                [-9.7030, -9.3149, -8.9268, -8.5386, -8.1505, -7.7624, -7.3743, -6.9862,
-                -6.5980, -6.2099, -5.8218, -5.4337, -5.0456, -4.6574, -4.2693, -3.8812,
-                -3.4931, -3.1050, -2.7168, -2.3287, -1.9406, -1.5525, -1.1644, -0.7762,
-                -0.3881, 0.0000, 0.3881, 0.7762, 1.1644, 1.5525, 1.9406, 2.3287,
-                2.7168, 3.1050, 3.4931, 3.8812, 4.2693, 4.6574, 5.0456, 5.4337,
-                5.8218, 6.2099, 6.5980, 6.9862, 7.3743, 7.7624, 8.1505, 8.5386,
-                8.9268, 9.3149, 9.7030],
-                [-9.7030, -9.3149, -8.9268, -8.5386, -8.1505, -7.7624, -7.3743, -6.9862,
-                -6.5980, -6.2099, -5.8218, -5.4337, -5.0456, -4.6574, -4.2693, -3.8812,
-                -3.4931, -3.1050, -2.7168, -2.3287, -1.9406, -1.5525, -1.1644, -0.7762,
-                -0.3881, 0.0000, 0.3881, 0.7762, 1.1644, 1.5525, 1.9406, 2.3287,
-                2.7168, 3.1050, 3.4931, 3.8812, 4.2693, 4.6574, 5.0456, 5.4337,
-                5.8218, 6.2099, 6.5980, 6.9862, 7.3743, 7.7624, 8.1505, 8.5386,
-                8.9268, 9.3149, 9.7030],
-                [-9.7030, -9.3149, -8.9268, -8.5386, -8.1505, -7.7624, -7.3743, -6.9862,
-                -6.5980, -6.2099, -5.8218, -5.4337, -5.0456, -4.6574, -4.2693, -3.8812,
-                -3.4931, -3.1050, -2.7168, -2.3287, -1.9406, -1.5525, -1.1644, -0.7762,
-                -0.3881, 0.0000, 0.3881, 0.7762, 1.1644, 1.5525, 1.9406, 2.3287,
-                2.7168, 3.1050, 3.4931, 3.8812, 4.2693, 4.6574, 5.0456, 5.4337,
-                5.8218, 6.2099, 6.5980, 6.9862, 7.3743, 7.7624, 8.1505, 8.5386,
-                8.9268, 9.3149, 9.7030]
-            ]).cuda()
-            expected = torch.tensor([
-                [0.0049, 0.0198, 0.0204, 0.0202, 0.0198, 0.0202, 0.0202, 0.0199, 0.0202,
-                0.0208, 0.0201, 0.0195, 0.0201, 0.0201, 0.0198, 0.0203, 0.0204, 0.0203,
-                0.0200, 0.0203, 0.0199, 0.0197, 0.0205, 0.0208, 0.0197, 0.0214, 0.0204,
-                0.0206, 0.0203, 0.0199, 0.0199, 0.0206, 0.0198, 0.0201, 0.0204, 0.0203,
-                0.0204, 0.0206, 0.0201, 0.0197, 0.0199, 0.0204, 0.0204, 0.0205, 0.0208,
-                0.0200, 0.0197, 0.0204, 0.0207, 0.0200, 0.0049],
-                [0.0049, 0.0198, 0.0204, 0.0202, 0.0198, 0.0202, 0.0202, 0.0199, 0.0202,
-                0.0208, 0.0202, 0.0196, 0.0201, 0.0201, 0.0198, 0.0203, 0.0204, 0.0203,
-                0.0200, 0.0203, 0.0199, 0.0197, 0.0205, 0.0208, 0.0197, 0.0214, 0.0204,
-                0.0206, 0.0203, 0.0199, 0.0199, 0.0206, 0.0198, 0.0201, 0.0204, 0.0203,
-                0.0204, 0.0206, 0.0201, 0.0197, 0.0199, 0.0204, 0.0204, 0.0205, 0.0208,
-                0.0200, 0.0197, 0.0204, 0.0206, 0.0200, 0.0049],
-                [0.0049, 0.0198, 0.0204, 0.0202, 0.0198, 0.0202, 0.0202, 0.0199, 0.0202,
-                0.0208, 0.0202, 0.0196, 0.0202, 0.0201, 0.0198, 0.0203, 0.0204, 0.0203,
-                0.0200, 0.0203, 0.0199, 0.0197, 0.0204, 0.0208, 0.0198, 0.0214, 0.0204,
-                0.0206, 0.0203, 0.0199, 0.0199, 0.0206, 0.0198, 0.0201, 0.0204, 0.0203,
-                0.0204, 0.0206, 0.0201, 0.0197, 0.0199, 0.0204, 0.0204, 0.0205, 0.0208,
-                0.0200, 0.0197, 0.0204, 0.0206, 0.0200, 0.0049]
-            ])
-            tt.assert_almost_equal(q.project(dist, support).cpu(), expected.cpu(), decimal=3)
+            q = QDist(self.model.cuda(), self.optimizer, ACTIONS, 51, -10.0, 10.0)
+            dist = torch.tensor(
+                [
+                    [
+                        0.0190,
+                        0.0197,
+                        0.0200,
+                        0.0190,
+                        0.0195,
+                        0.0198,
+                        0.0194,
+                        0.0192,
+                        0.0201,
+                        0.0203,
+                        0.0189,
+                        0.0190,
+                        0.0199,
+                        0.0193,
+                        0.0192,
+                        0.0199,
+                        0.0198,
+                        0.0197,
+                        0.0193,
+                        0.0198,
+                        0.0192,
+                        0.0191,
+                        0.0200,
+                        0.0202,
+                        0.0191,
+                        0.0202,
+                        0.0198,
+                        0.0200,
+                        0.0198,
+                        0.0193,
+                        0.0192,
+                        0.0202,
+                        0.0192,
+                        0.0194,
+                        0.0199,
+                        0.0197,
+                        0.0197,
+                        0.0201,
+                        0.0199,
+                        0.0190,
+                        0.0192,
+                        0.0195,
+                        0.0202,
+                        0.0194,
+                        0.0203,
+                        0.0201,
+                        0.0190,
+                        0.0192,
+                        0.0201,
+                        0.0201,
+                        0.0192,
+                    ],
+                    [
+                        0.0191,
+                        0.0197,
+                        0.0200,
+                        0.0190,
+                        0.0195,
+                        0.0198,
+                        0.0194,
+                        0.0192,
+                        0.0201,
+                        0.0203,
+                        0.0190,
+                        0.0190,
+                        0.0199,
+                        0.0193,
+                        0.0192,
+                        0.0199,
+                        0.0198,
+                        0.0197,
+                        0.0193,
+                        0.0198,
+                        0.0192,
+                        0.0191,
+                        0.0200,
+                        0.0202,
+                        0.0191,
+                        0.0202,
+                        0.0198,
+                        0.0200,
+                        0.0198,
+                        0.0193,
+                        0.0192,
+                        0.0202,
+                        0.0192,
+                        0.0194,
+                        0.0199,
+                        0.0197,
+                        0.0197,
+                        0.0200,
+                        0.0199,
+                        0.0190,
+                        0.0192,
+                        0.0195,
+                        0.0202,
+                        0.0194,
+                        0.0203,
+                        0.0201,
+                        0.0190,
+                        0.0192,
+                        0.0201,
+                        0.0200,
+                        0.0192,
+                    ],
+                    [
+                        0.0191,
+                        0.0197,
+                        0.0200,
+                        0.0190,
+                        0.0195,
+                        0.0198,
+                        0.0194,
+                        0.0192,
+                        0.0200,
+                        0.0203,
+                        0.0190,
+                        0.0191,
+                        0.0199,
+                        0.0193,
+                        0.0192,
+                        0.0199,
+                        0.0198,
+                        0.0197,
+                        0.0193,
+                        0.0198,
+                        0.0192,
+                        0.0191,
+                        0.0199,
+                        0.0202,
+                        0.0192,
+                        0.0202,
+                        0.0198,
+                        0.0200,
+                        0.0198,
+                        0.0193,
+                        0.0192,
+                        0.0202,
+                        0.0192,
+                        0.0194,
+                        0.0199,
+                        0.0197,
+                        0.0197,
+                        0.0200,
+                        0.0199,
+                        0.0190,
+                        0.0192,
+                        0.0195,
+                        0.0202,
+                        0.0194,
+                        0.0203,
+                        0.0201,
+                        0.0190,
+                        0.0192,
+                        0.0201,
+                        0.0200,
+                        0.0192,
+                    ],
+                ]
+            ).cuda()
+            support = torch.tensor(
+                [
+                    [
+                        -9.7030,
+                        -9.3149,
+                        -8.9268,
+                        -8.5386,
+                        -8.1505,
+                        -7.7624,
+                        -7.3743,
+                        -6.9862,
+                        -6.5980,
+                        -6.2099,
+                        -5.8218,
+                        -5.4337,
+                        -5.0456,
+                        -4.6574,
+                        -4.2693,
+                        -3.8812,
+                        -3.4931,
+                        -3.1050,
+                        -2.7168,
+                        -2.3287,
+                        -1.9406,
+                        -1.5525,
+                        -1.1644,
+                        -0.7762,
+                        -0.3881,
+                        0.0000,
+                        0.3881,
+                        0.7762,
+                        1.1644,
+                        1.5525,
+                        1.9406,
+                        2.3287,
+                        2.7168,
+                        3.1050,
+                        3.4931,
+                        3.8812,
+                        4.2693,
+                        4.6574,
+                        5.0456,
+                        5.4337,
+                        5.8218,
+                        6.2099,
+                        6.5980,
+                        6.9862,
+                        7.3743,
+                        7.7624,
+                        8.1505,
+                        8.5386,
+                        8.9268,
+                        9.3149,
+                        9.7030,
+                    ],
+                    [
+                        -9.7030,
+                        -9.3149,
+                        -8.9268,
+                        -8.5386,
+                        -8.1505,
+                        -7.7624,
+                        -7.3743,
+                        -6.9862,
+                        -6.5980,
+                        -6.2099,
+                        -5.8218,
+                        -5.4337,
+                        -5.0456,
+                        -4.6574,
+                        -4.2693,
+                        -3.8812,
+                        -3.4931,
+                        -3.1050,
+                        -2.7168,
+                        -2.3287,
+                        -1.9406,
+                        -1.5525,
+                        -1.1644,
+                        -0.7762,
+                        -0.3881,
+                        0.0000,
+                        0.3881,
+                        0.7762,
+                        1.1644,
+                        1.5525,
+                        1.9406,
+                        2.3287,
+                        2.7168,
+                        3.1050,
+                        3.4931,
+                        3.8812,
+                        4.2693,
+                        4.6574,
+                        5.0456,
+                        5.4337,
+                        5.8218,
+                        6.2099,
+                        6.5980,
+                        6.9862,
+                        7.3743,
+                        7.7624,
+                        8.1505,
+                        8.5386,
+                        8.9268,
+                        9.3149,
+                        9.7030,
+                    ],
+                    [
+                        -9.7030,
+                        -9.3149,
+                        -8.9268,
+                        -8.5386,
+                        -8.1505,
+                        -7.7624,
+                        -7.3743,
+                        -6.9862,
+                        -6.5980,
+                        -6.2099,
+                        -5.8218,
+                        -5.4337,
+                        -5.0456,
+                        -4.6574,
+                        -4.2693,
+                        -3.8812,
+                        -3.4931,
+                        -3.1050,
+                        -2.7168,
+                        -2.3287,
+                        -1.9406,
+                        -1.5525,
+                        -1.1644,
+                        -0.7762,
+                        -0.3881,
+                        0.0000,
+                        0.3881,
+                        0.7762,
+                        1.1644,
+                        1.5525,
+                        1.9406,
+                        2.3287,
+                        2.7168,
+                        3.1050,
+                        3.4931,
+                        3.8812,
+                        4.2693,
+                        4.6574,
+                        5.0456,
+                        5.4337,
+                        5.8218,
+                        6.2099,
+                        6.5980,
+                        6.9862,
+                        7.3743,
+                        7.7624,
+                        8.1505,
+                        8.5386,
+                        8.9268,
+                        9.3149,
+                        9.7030,
+                    ],
+                ]
+            ).cuda()
+            expected = torch.tensor(
+                [
+                    [
+                        0.0049,
+                        0.0198,
+                        0.0204,
+                        0.0202,
+                        0.0198,
+                        0.0202,
+                        0.0202,
+                        0.0199,
+                        0.0202,
+                        0.0208,
+                        0.0201,
+                        0.0195,
+                        0.0201,
+                        0.0201,
+                        0.0198,
+                        0.0203,
+                        0.0204,
+                        0.0203,
+                        0.0200,
+                        0.0203,
+                        0.0199,
+                        0.0197,
+                        0.0205,
+                        0.0208,
+                        0.0197,
+                        0.0214,
+                        0.0204,
+                        0.0206,
+                        0.0203,
+                        0.0199,
+                        0.0199,
+                        0.0206,
+                        0.0198,
+                        0.0201,
+                        0.0204,
+                        0.0203,
+                        0.0204,
+                        0.0206,
+                        0.0201,
+                        0.0197,
+                        0.0199,
+                        0.0204,
+                        0.0204,
+                        0.0205,
+                        0.0208,
+                        0.0200,
+                        0.0197,
+                        0.0204,
+                        0.0207,
+                        0.0200,
+                        0.0049,
+                    ],
+                    [
+                        0.0049,
+                        0.0198,
+                        0.0204,
+                        0.0202,
+                        0.0198,
+                        0.0202,
+                        0.0202,
+                        0.0199,
+                        0.0202,
+                        0.0208,
+                        0.0202,
+                        0.0196,
+                        0.0201,
+                        0.0201,
+                        0.0198,
+                        0.0203,
+                        0.0204,
+                        0.0203,
+                        0.0200,
+                        0.0203,
+                        0.0199,
+                        0.0197,
+                        0.0205,
+                        0.0208,
+                        0.0197,
+                        0.0214,
+                        0.0204,
+                        0.0206,
+                        0.0203,
+                        0.0199,
+                        0.0199,
+                        0.0206,
+                        0.0198,
+                        0.0201,
+                        0.0204,
+                        0.0203,
+                        0.0204,
+                        0.0206,
+                        0.0201,
+                        0.0197,
+                        0.0199,
+                        0.0204,
+                        0.0204,
+                        0.0205,
+                        0.0208,
+                        0.0200,
+                        0.0197,
+                        0.0204,
+                        0.0206,
+                        0.0200,
+                        0.0049,
+                    ],
+                    [
+                        0.0049,
+                        0.0198,
+                        0.0204,
+                        0.0202,
+                        0.0198,
+                        0.0202,
+                        0.0202,
+                        0.0199,
+                        0.0202,
+                        0.0208,
+                        0.0202,
+                        0.0196,
+                        0.0202,
+                        0.0201,
+                        0.0198,
+                        0.0203,
+                        0.0204,
+                        0.0203,
+                        0.0200,
+                        0.0203,
+                        0.0199,
+                        0.0197,
+                        0.0204,
+                        0.0208,
+                        0.0198,
+                        0.0214,
+                        0.0204,
+                        0.0206,
+                        0.0203,
+                        0.0199,
+                        0.0199,
+                        0.0206,
+                        0.0198,
+                        0.0201,
+                        0.0204,
+                        0.0203,
+                        0.0204,
+                        0.0206,
+                        0.0201,
+                        0.0197,
+                        0.0199,
+                        0.0204,
+                        0.0204,
+                        0.0205,
+                        0.0208,
+                        0.0200,
+                        0.0197,
+                        0.0204,
+                        0.0206,
+                        0.0200,
+                        0.0049,
+                    ],
+                ]
+            )
+            tt.assert_almost_equal(
+                q.project(dist, support).cpu(), expected.cpu(), decimal=3
+            )
 
 
 if __name__ == "__main__":
diff --git a/all/approximation/q_network_test.py b/all/approximation/q_network_test.py
index dee5a061..911715ea 100644
--- a/all/approximation/q_network_test.py
+++ b/all/approximation/q_network_test.py
@@ -1,11 +1,13 @@
 import unittest
+
+import numpy as np
 import torch
+import torch_testing as tt
 from torch import nn
 from torch.nn.functional import smooth_l1_loss
-import torch_testing as tt
-import numpy as np
+
+from all.approximation import FixedTarget, QNetwork
 from all.core import State, StateArray
-from all.approximation import QNetwork, FixedTarget
 
 STATE_DIM = 2
 ACTIONS = 3
@@ -14,31 +16,30 @@
 class TestQNetwork(unittest.TestCase):
     def setUp(self):
         torch.manual_seed(2)
-        self.model = nn.Sequential(
-            nn.Linear(STATE_DIM, ACTIONS)
-        )
+        self.model = nn.Sequential(nn.Linear(STATE_DIM, ACTIONS))
 
         def optimizer(params):
             return torch.optim.SGD(params, lr=0.1)
+
         self.q = QNetwork(self.model, optimizer)
 
     def test_eval_list(self):
         states = StateArray(
-            torch.randn(5, STATE_DIM),
-            (5,),
-            mask=torch.tensor([1, 1, 0, 1, 0])
+            torch.randn(5, STATE_DIM), (5,), mask=torch.tensor([1, 1, 0, 1, 0])
         )
         result = self.q.eval(states)
         tt.assert_almost_equal(
             result,
-            torch.tensor([
-                [-0.238509, -0.726287, -0.034026],
-                [-0.35688755, -0.6612102, 0.34849477],
-                [0., 0., 0.],
-                [0.1944, -0.5536, -0.2345],
-                [0., 0., 0.]
-            ]),
-            decimal=2
+            torch.tensor(
+                [
+                    [-0.238509, -0.726287, -0.034026],
+                    [-0.35688755, -0.6612102, 0.34849477],
+                    [0.0, 0.0, 0.0],
+                    [0.1944, -0.5536, -0.2345],
+                    [0.0, 0.0, 0.0],
+                ]
+            ),
+            decimal=2,
         )
 
     def test_eval_actions(self):
@@ -46,20 +47,16 @@ def test_eval_actions(self):
         actions = [1, 2, 0]
         result = self.q.eval(states, actions)
         self.assertEqual(result.shape, torch.Size([3]))
-        tt.assert_almost_equal(result, torch.tensor([-0.7262873, 0.3484948, -0.0296164]))
+        tt.assert_almost_equal(
+            result, torch.tensor([-0.7262873, 0.3484948, -0.0296164])
+        )
 
     def test_target_net(self):
         torch.manual_seed(2)
-        model = nn.Sequential(
-            nn.Linear(1, 1)
-        )
+        model = nn.Sequential(nn.Linear(1, 1))
         optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
-        q = QNetwork(
-            model,
-            optimizer,
-            target=FixedTarget(3)
-        )
-        inputs = State(torch.tensor([1.]))
+        q = QNetwork(model, optimizer, target=FixedTarget(3))
+        inputs = State(torch.tensor([1.0]))
 
         def loss(policy_value):
             target = policy_value - 1
@@ -95,5 +92,5 @@ def loss(policy_value):
         np.testing.assert_equal(target_value, -0.6085841655731201)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/all/approximation/v_network_test.py b/all/approximation/v_network_test.py
index cc281ed7..3c05320c 100644
--- a/all/approximation/v_network_test.py
+++ b/all/approximation/v_network_test.py
@@ -1,7 +1,9 @@
 import unittest
+
 import torch
-from torch import nn
 import torch_testing as tt
+from torch import nn
+
 from all.approximation.v_network import VNetwork
 from all.core import StateArray
 
@@ -16,31 +18,29 @@ def loss(value, error):
 class TestVNetwork(unittest.TestCase):
     def setUp(self):
         torch.manual_seed(2)
-        self.model = nn.Sequential(
-            nn.Linear(STATE_DIM, 1)
-        )
+        self.model = nn.Sequential(nn.Linear(STATE_DIM, 1))
 
         optimizer = torch.optim.SGD(self.model.parameters(), lr=0.1)
         self.v = VNetwork(self.model, optimizer)
 
     def test_reinforce_list(self):
         states = StateArray(
-            torch.randn(5, STATE_DIM),
-            (5,),
-            mask=torch.tensor([1, 1, 0, 1, 0])
+            torch.randn(5, STATE_DIM), (5,), mask=torch.tensor([1, 1, 0, 1, 0])
         )
         result = self.v(states)
-        tt.assert_almost_equal(result, torch.tensor([0.7053187, 0.3975691, 0., 0.2701665, 0.]))
+        tt.assert_almost_equal(
+            result, torch.tensor([0.7053187, 0.3975691, 0.0, 0.2701665, 0.0])
+        )
 
         self.v.reinforce(loss(result, torch.tensor([1, -1, 1, 1, 1])).float())
         result = self.v(states)
-        tt.assert_almost_equal(result, torch.tensor([0.9732854, 0.5453826, 0., 0.4344811, 0.]))
+        tt.assert_almost_equal(
+            result, torch.tensor([0.9732854, 0.5453826, 0.0, 0.4344811, 0.0])
+        )
 
     def test_multi_reinforce(self):
         states = StateArray(
-            torch.randn(6, STATE_DIM),
-            (6,),
-            mask=torch.tensor([1, 1, 0, 1, 0, 0, 0])
+            torch.randn(6, STATE_DIM), (6,), mask=torch.tensor([1, 1, 0, 1, 0, 0, 0])
         )
         result1 = self.v(states[0:2])
         self.v.reinforce(loss(result1, torch.tensor([1, 2])).float())
@@ -52,5 +52,5 @@ def test_multi_reinforce(self):
             self.v.reinforce(loss(result3, torch.tensor([1, 2])).float())
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/all/bodies/time_test.py b/all/bodies/time_test.py
index 31dc3ad8..d0971278 100644
--- a/all/bodies/time_test.py
+++ b/all/bodies/time_test.py
@@ -1,11 +1,13 @@
 import unittest
+
 import torch
 import torch_testing as tt
-from all.core import State, StateArray
+
 from all.bodies import TimeFeature
+from all.core import State, StateArray
 
 
-class TestAgent():
+class TestAgent:
     def __init__(self):
         self.last_state = None
 
@@ -23,57 +25,101 @@ def setUp(self):
     def test_init(self):
         state = State(torch.randn(4))
         self.agent.act(state)
-        tt.assert_allclose(self.test_agent.last_state.observation, torch.tensor(
-            [0.3923, -0.2236, -0.3195, -1.2050, 0.0000]), atol=1e-04)
+        tt.assert_allclose(
+            self.test_agent.last_state.observation,
+            torch.tensor([0.3923, -0.2236, -0.3195, -1.2050, 0.0000]),
+            atol=1e-04,
+        )
 
     def test_single_env(self):
         state = State(torch.randn(4))
         self.agent.act(state)
-        tt.assert_allclose(self.test_agent.last_state.observation, torch.tensor(
-            [0.3923, -0.2236, -0.3195, -1.2050, 0.]), atol=1e-04)
+        tt.assert_allclose(
+            self.test_agent.last_state.observation,
+            torch.tensor([0.3923, -0.2236, -0.3195, -1.2050, 0.0]),
+            atol=1e-04,
+        )
         self.agent.act(state)
-        tt.assert_allclose(self.test_agent.last_state.observation, torch.tensor(
-            [0.3923, -0.2236, -0.3195, -1.2050, 1e-3]), atol=1e-04)
+        tt.assert_allclose(
+            self.test_agent.last_state.observation,
+            torch.tensor([0.3923, -0.2236, -0.3195, -1.2050, 1e-3]),
+            atol=1e-04,
+        )
         self.agent.act(state)
-        tt.assert_allclose(self.test_agent.last_state.observation, torch.tensor(
-            [0.3923, -0.2236, -0.3195, -1.2050, 2e-3]), atol=1e-04)
+        tt.assert_allclose(
+            self.test_agent.last_state.observation,
+            torch.tensor([0.3923, -0.2236, -0.3195, -1.2050, 2e-3]),
+            atol=1e-04,
+        )
 
     def test_reset(self):
         state = State(torch.randn(4))
         self.agent.act(state)
-        tt.assert_allclose(self.test_agent.last_state.observation, torch.tensor(
-            [0.3923, -0.2236, -0.3195, -1.2050, 0.0000]), atol=1e-04)
+        tt.assert_allclose(
+            self.test_agent.last_state.observation,
+            torch.tensor([0.3923, -0.2236, -0.3195, -1.2050, 0.0000]),
+            atol=1e-04,
+        )
         self.agent.act(state)
-        tt.assert_allclose(self.test_agent.last_state.observation, torch.tensor(
-            [0.3923, -0.2236, -0.3195, -1.2050, 1e-3]), atol=1e-04)
+        tt.assert_allclose(
+            self.test_agent.last_state.observation,
+            torch.tensor([0.3923, -0.2236, -0.3195, -1.2050, 1e-3]),
+            atol=1e-04,
+        )
         self.agent.act(State(state.observation, done=True))
-        tt.assert_allclose(self.test_agent.last_state.observation, torch.tensor(
-            [0.3923, -0.2236, -0.3195, -1.2050, 2e-3]), atol=1e-04)
+        tt.assert_allclose(
+            self.test_agent.last_state.observation,
+            torch.tensor([0.3923, -0.2236, -0.3195, -1.2050, 2e-3]),
+            atol=1e-04,
+        )
         self.agent.act(State(state.observation))
-        tt.assert_allclose(self.test_agent.last_state.observation, torch.tensor(
-            [0.3923, -0.2236, -0.3195, -1.2050, 0.0000]), atol=1e-04)
+        tt.assert_allclose(
+            self.test_agent.last_state.observation,
+            torch.tensor([0.3923, -0.2236, -0.3195, -1.2050, 0.0000]),
+            atol=1e-04,
+        )
         self.agent.act(state)
-        tt.assert_allclose(self.test_agent.last_state.observation, torch.tensor(
-            [0.3923, -0.2236, -0.3195, -1.2050, 1e-3]), atol=1e-04)
+        tt.assert_allclose(
+            self.test_agent.last_state.observation,
+            torch.tensor([0.3923, -0.2236, -0.3195, -1.2050, 1e-3]),
+            atol=1e-04,
+        )
 
     def test_multi_env(self):
         state = StateArray(torch.randn(2, 2), (2,))
         self.agent.act(state)
-        tt.assert_allclose(self.test_agent.last_state.observation, torch.tensor(
-            [[0.3923, -0.2236, 0.], [-0.3195, -1.2050, 0.]]), atol=1e-04)
+        tt.assert_allclose(
+            self.test_agent.last_state.observation,
+            torch.tensor([[0.3923, -0.2236, 0.0], [-0.3195, -1.2050, 0.0]]),
+            atol=1e-04,
+        )
         self.agent.act(state)
-        tt.assert_allclose(self.test_agent.last_state.observation, torch.tensor(
-            [[0.3923, -0.2236, 1e-3], [-0.3195, -1.2050, 1e-3]]), atol=1e-04)
-        self.agent.act(StateArray(state.observation, (2,), done=torch.tensor([False, True])))
-        tt.assert_allclose(self.test_agent.last_state.observation, torch.tensor(
-            [[0.3923, -0.2236, 2e-3], [-0.3195, -1.2050, 2e-3]]), atol=1e-04)
+        tt.assert_allclose(
+            self.test_agent.last_state.observation,
+            torch.tensor([[0.3923, -0.2236, 1e-3], [-0.3195, -1.2050, 1e-3]]),
+            atol=1e-04,
+        )
+        self.agent.act(
+            StateArray(state.observation, (2,), done=torch.tensor([False, True]))
+        )
+        tt.assert_allclose(
+            self.test_agent.last_state.observation,
+            torch.tensor([[0.3923, -0.2236, 2e-3], [-0.3195, -1.2050, 2e-3]]),
+            atol=1e-04,
+        )
         self.agent.act(state)
-        tt.assert_allclose(self.test_agent.last_state.observation, torch.tensor(
-            [[0.3923, -0.2236, 3e-3], [-0.3195, -1.2050, 0.]]), atol=1e-04)
+        tt.assert_allclose(
+            self.test_agent.last_state.observation,
+            torch.tensor([[0.3923, -0.2236, 3e-3], [-0.3195, -1.2050, 0.0]]),
+            atol=1e-04,
+        )
         self.agent.act(state)
-        tt.assert_allclose(self.test_agent.last_state.observation, torch.tensor(
-            [[0.3923, -0.2236, 4e-3], [-0.3195, -1.2050, 1e-3]]), atol=1e-04)
+        tt.assert_allclose(
+            self.test_agent.last_state.observation,
+            torch.tensor([[0.3923, -0.2236, 4e-3], [-0.3195, -1.2050, 1e-3]]),
+            atol=1e-04,
+        )
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/all/core/state_test.py b/all/core/state_test.py
index 6670b553..00c2c4ee 100644
--- a/all/core/state_test.py
+++ b/all/core/state_test.py
@@ -1,8 +1,10 @@
 import unittest
 import warnings
+
 import numpy as np
 import torch
 import torch_testing as tt
+
 from all.core import State, StateArray
 
 
@@ -11,86 +13,75 @@ def test_constructor_defaults(self):
         observation = torch.randn(3, 4)
         state = State(observation)
         tt.assert_equal(state.observation, observation)
-        self.assertEqual(state.mask, 1.)
+        self.assertEqual(state.mask, 1.0)
         self.assertEqual(state.done, False)
-        self.assertEqual(state.reward, 0.)
+        self.assertEqual(state.reward, 0.0)
         self.assertEqual(state.shape, ())
 
     def test_from_dict(self):
         observation = torch.randn(3, 4)
-        state = State({
-            'observation': observation,
-            'done': True,
-            'mask': 1,
-            'reward': 5.
-        })
+        state = State(
+            {"observation": observation, "done": True, "mask": 1, "reward": 5.0}
+        )
         tt.assert_equal(state.observation, observation)
         self.assertEqual(state.done, True)
-        self.assertEqual(state.mask, 1.)
-        self.assertEqual(state.reward, 5.)
+        self.assertEqual(state.mask, 1.0)
+        self.assertEqual(state.reward, 5.0)
 
     def test_auto_mask_true(self):
         observation = torch.randn(3, 4)
-        state = State({
-            'observation': observation,
-            'done': True,
-            'reward': 5.
-        })
-        self.assertEqual(state.mask, 0.)
+        state = State({"observation": observation, "done": True, "reward": 5.0})
+        self.assertEqual(state.mask, 0.0)
 
     def test_auto_mask_false(self):
         observation = torch.randn(3, 4)
-        state = State({
-            'observation': observation,
-            'done': False,
-            'reward': 5.
-        })
-        self.assertEqual(state.mask, 1.)
+        state = State({"observation": observation, "done": False, "reward": 5.0})
+        self.assertEqual(state.mask, 1.0)
 
     def test_from_gym_reset(self):
         observation = np.array([1, 2, 3])
-        state = State.from_gym((observation, {'coolInfo': 3}))
+        state = State.from_gym((observation, {"coolInfo": 3}))
         tt.assert_equal(state.observation, torch.from_numpy(observation))
-        self.assertEqual(state.mask, 1.)
+        self.assertEqual(state.mask, 1.0)
         self.assertEqual(state.done, False)
-        self.assertEqual(state.reward, 0.)
+        self.assertEqual(state.reward, 0.0)
         self.assertEqual(state.shape, ())
-        self.assertEqual(state['coolInfo'], 3.)
+        self.assertEqual(state["coolInfo"], 3.0)
 
     def test_from_gym_step(self):
         observation = np.array([1, 2, 3])
-        state = State.from_gym((observation, 2., True, False, {'coolInfo': 3.}))
+        state = State.from_gym((observation, 2.0, True, False, {"coolInfo": 3.0}))
         tt.assert_equal(state.observation, torch.from_numpy(observation))
-        self.assertEqual(state.mask, 0.)
+        self.assertEqual(state.mask, 0.0)
         self.assertEqual(state.done, True)
-        self.assertEqual(state.reward, 2.)
-        self.assertEqual(state['coolInfo'], 3.)
+        self.assertEqual(state.reward, 2.0)
+        self.assertEqual(state["coolInfo"], 3.0)
         self.assertEqual(state.shape, ())
 
     def test_from_truncated_gym_step(self):
         observation = np.array([1, 2, 3])
-        state = State.from_gym((observation, 2., False, True, {'coolInfo': 3.}))
+        state = State.from_gym((observation, 2.0, False, True, {"coolInfo": 3.0}))
         tt.assert_equal(state.observation, torch.from_numpy(observation))
-        self.assertEqual(state.mask, 1.)
+        self.assertEqual(state.mask, 1.0)
         self.assertEqual(state.done, True)
-        self.assertEqual(state.reward, 2.)
-        self.assertEqual(state['coolInfo'], 3.)
+        self.assertEqual(state.reward, 2.0)
+        self.assertEqual(state["coolInfo"], 3.0)
         self.assertEqual(state.shape, ())
 
     def test_legacy_gym_step(self):
         observation = np.array([1, 2, 3])
-        state = State.from_gym((observation, 2., True, {'coolInfo': 3.}))
+        state = State.from_gym((observation, 2.0, True, {"coolInfo": 3.0}))
         tt.assert_equal(state.observation, torch.from_numpy(observation))
-        self.assertEqual(state.mask, 0.)
+        self.assertEqual(state.mask, 0.0)
         self.assertEqual(state.done, True)
-        self.assertEqual(state.reward, 2.)
-        self.assertEqual(state['coolInfo'], 3.)
+        self.assertEqual(state.reward, 2.0)
+        self.assertEqual(state["coolInfo"], 3.0)
         self.assertEqual(state.shape, ())
 
     def test_as_input(self):
         observation = torch.randn(3, 4)
         state = State(observation)
-        self.assertEqual(state.as_input('observation').shape, (1, 3, 4))
+        self.assertEqual(state.as_input("observation").shape, (1, 3, 4))
 
     def test_as_output(self):
         observation = torch.randn(3, 4)
@@ -100,30 +91,30 @@ def test_as_output(self):
 
     def test_apply_mask(self):
         observation = torch.randn(3, 4)
-        state = State.from_gym((observation, 0., True, False, {}))
+        state = State.from_gym((observation, 0.0, True, False, {}))
         tt.assert_equal(state.apply_mask(observation), torch.zeros(3, 4))
 
     def test_apply(self):
         observation = torch.randn(3, 4)
         state = State(observation)
         model = torch.nn.Conv1d(3, 5, 2)
-        output = state.apply(model, 'observation')
+        output = state.apply(model, "observation")
         self.assertEqual(output.shape, (5, 3))
         self.assertNotEqual(output.sum().item(), 0)
 
     def test_apply_done(self):
         observation = torch.randn(3, 4)
-        state = State.from_gym((observation, 0., True, False, {}))
+        state = State.from_gym((observation, 0.0, True, False, {}))
         model = torch.nn.Conv1d(3, 5, 2)
-        output = state.apply(model, 'observation')
+        output = state.apply(model, "observation")
         self.assertEqual(output.shape, (5, 3))
         self.assertEqual(output.sum().item(), 0)
 
     def test_to_device(self):
         observation = torch.randn(3, 4)
-        state = State(observation, device=torch.device('cpu'))
+        state = State(observation, device=torch.device("cpu"))
         state_cpu = state.to("cpu")
-        self.assertTrue(torch.equal(state['observation'], state_cpu['observation']))
+        self.assertTrue(torch.equal(state["observation"], state_cpu["observation"]))
         self.assertFalse(state is state_cpu)
 
 
@@ -132,23 +123,23 @@ def test_constructor_defaults(self):
         raw = torch.randn(3, 4)
         state = State(raw, (3,))
         tt.assert_equal(state.observation, raw)
-        self.assertEqual(state.mask, 1.)
+        self.assertEqual(state.mask, 1.0)
         self.assertEqual(state.done, False)
-        self.assertEqual(state.reward, 0.)
+        self.assertEqual(state.reward, 0.0)
 
     def test_apply(self):
         observation = torch.randn(3, 4)
         state = StateArray(observation, (3,))
         model = torch.nn.Linear(4, 2)
-        output = state.apply(model, 'observation')
+        output = state.apply(model, "observation")
         self.assertEqual(output.shape, (3, 2))
         self.assertNotEqual(output.sum().item(), 0)
 
     def test_apply_done(self):
         observation = torch.randn(3, 4)
-        state = StateArray(observation, (3,), mask=torch.tensor([0., 0., 0.]))
+        state = StateArray(observation, (3,), mask=torch.tensor([0.0, 0.0, 0.0]))
         model = torch.nn.Linear(4, 2)
-        output = state.apply(model, 'observation')
+        output = state.apply(model, "observation")
         self.assertEqual(output.shape, (3, 2))
         self.assertEqual(output.sum().item(), 0)
 
@@ -160,17 +151,19 @@ def test_as_output(self):
 
     def test_auto_mask(self):
         observation = torch.randn(3, 4)
-        state = StateArray({
-            'observation': observation,
-            'done': torch.tensor([True, False, True]),
-        }, (3,))
-        tt.assert_equal(state.mask, torch.tensor([0., 1., 0.]))
+        state = StateArray(
+            {
+                "observation": observation,
+                "done": torch.tensor([True, False, True]),
+            },
+            (3,),
+        )
+        tt.assert_equal(state.mask, torch.tensor([0.0, 1.0, 0.0]))
 
     def test_multi_dim(self):
-        state = StateArray.array([
-            State(torch.randn((3, 4))),
-            State(torch.randn((3, 4)))
-        ])
+        state = StateArray.array(
+            [State(torch.randn((3, 4))), State(torch.randn((3, 4)))]
+        )
         self.assertEqual(state.shape, (2,))
         state = StateArray.array([state] * 3)
         self.assertEqual(state.shape, (3, 2))
@@ -181,10 +174,9 @@ def test_multi_dim(self):
         tt.assert_equal(state.reward, torch.zeros((5, 3, 2)))
 
     def test_view(self):
-        state = StateArray.array([
-            State(torch.randn((3, 4))),
-            State(torch.randn((3, 4)))
-        ])
+        state = StateArray.array(
+            [State(torch.randn((3, 4))), State(torch.randn((3, 4)))]
+        )
         self.assertEqual(state.shape, (2,))
         state = StateArray.array([state] * 3)
         self.assertEqual(state.shape, (3, 2))
@@ -193,28 +185,45 @@ def test_view(self):
         self.assertEqual(state.observation.shape, (2, 3, 3, 4))
 
     def test_batch_exec(self):
-        zeros = StateArray.array([
-            State(torch.zeros((3, 4))),
-            State(torch.zeros((3, 4))),
-            State(torch.zeros((3, 4)))
-        ])
-        ones_state = zeros.batch_execute(2, lambda x: StateArray({'observation': x.observation + 1}, x.shape, x.device))
+        zeros = StateArray.array(
+            [
+                State(torch.zeros((3, 4))),
+                State(torch.zeros((3, 4))),
+                State(torch.zeros((3, 4))),
+            ]
+        )
+        ones_state = zeros.batch_execute(
+            2,
+            lambda x: StateArray({"observation": x.observation + 1}, x.shape, x.device),
+        )
         ones_tensor = zeros.batch_execute(2, lambda x: x.observation + 1)
         self.assertEqual(ones_state.shape, (3,))
         self.assertTrue(torch.equal(ones_state.observation, torch.ones((3, 3, 4))))
         self.assertTrue(torch.equal(ones_tensor, torch.ones((3, 3, 4))))
 
     def test_cat(self):
-        i1 = StateArray({'observation': torch.zeros((2, 3, 4)), 'reward': torch.ones((2,))}, shape=(2,))
-        i2 = StateArray({'observation': torch.zeros((1, 3, 4)), 'reward': torch.ones((1,))}, shape=(1,))
+        i1 = StateArray(
+            {"observation": torch.zeros((2, 3, 4)), "reward": torch.ones((2,))},
+            shape=(2,),
+        )
+        i2 = StateArray(
+            {"observation": torch.zeros((1, 3, 4)), "reward": torch.ones((1,))},
+            shape=(1,),
+        )
         cat = StateArray.cat([i1, i2])
         self.assertEqual(cat.shape, (3,))
         self.assertEqual(cat.observation.shape, (3, 3, 4))
         self.assertEqual(cat.reward.shape, (3,))
 
     def test_cat_axis1(self):
-        i1 = StateArray({'observation': torch.zeros((2, 3, 4)), 'reward': torch.ones((2, 3))}, shape=(2, 3))
-        i2 = StateArray({'observation': torch.zeros((2, 2, 4)), 'reward': torch.ones((2, 2))}, shape=(2, 2))
+        i1 = StateArray(
+            {"observation": torch.zeros((2, 3, 4)), "reward": torch.ones((2, 3))},
+            shape=(2, 3),
+        )
+        i2 = StateArray(
+            {"observation": torch.zeros((2, 2, 4)), "reward": torch.ones((2, 2))},
+            shape=(2, 2),
+        )
         cat = StateArray.cat([i1, i2], axis=1)
         self.assertEqual(cat.shape, (2, 5))
         self.assertEqual(cat.observation.shape, (2, 5, 4))
@@ -223,33 +232,41 @@ def test_cat_axis1(self):
     def test_key_error(self):
         with warnings.catch_warnings(record=True) as w:
             warnings.simplefilter("always")
-            StateArray.array([
-                State({
-                    'observation': torch.tensor([1, 2]),
-                    'other_key': True
-                }),
-                State({
-                    'observation': torch.tensor([1, 2]),
-                }),
-            ])
+            StateArray.array(
+                [
+                    State({"observation": torch.tensor([1, 2]), "other_key": True}),
+                    State(
+                        {
+                            "observation": torch.tensor([1, 2]),
+                        }
+                    ),
+                ]
+            )
             self.assertEqual(len(w), 1)
-            self.assertEqual(w[0].message.args[0], 'KeyError while creating StateArray for key "other_key", omitting.')
+            self.assertEqual(
+                w[0].message.args[0],
+                'KeyError while creating StateArray for key "other_key", omitting.',
+            )
 
     def test_type_error(self):
         with warnings.catch_warnings(record=True) as w:
             warnings.simplefilter("always")
-            StateArray.array([
-                State({
-                    'observation': torch.tensor([1, 2]),
-                    'other_key': torch.tensor([1])
-                }),
-                State({
-                    'observation': torch.tensor([1, 2]),
-                    'other_key': 5.
-                }),
-            ])
+            StateArray.array(
+                [
+                    State(
+                        {
+                            "observation": torch.tensor([1, 2]),
+                            "other_key": torch.tensor([1]),
+                        }
+                    ),
+                    State({"observation": torch.tensor([1, 2]), "other_key": 5.0}),
+                ]
+            )
             self.assertEqual(len(w), 1)
-            self.assertEqual(w[0].message.args[0], 'TypeError while creating StateArray for key "other_key", omitting.')
+            self.assertEqual(
+                w[0].message.args[0],
+                'TypeError while creating StateArray for key "other_key", omitting.',
+            )
 
 
 if __name__ == "__main__":
diff --git a/all/environments/atari_test.py b/all/environments/atari_test.py
index 4b6bf4aa..fef2dfd6 100644
--- a/all/environments/atari_test.py
+++ b/all/environments/atari_test.py
@@ -1,10 +1,11 @@
 import unittest
+
 from all.environments import AtariEnvironment
 
 
 class AtariEnvironmentTest(unittest.TestCase):
     def test_reset(self):
-        env = AtariEnvironment('Breakout')
+        env = AtariEnvironment("Breakout")
         state = env.reset()
         self.assertEqual(state.observation.shape, (1, 84, 84))
         self.assertEqual(state.reward, 0)
@@ -12,30 +13,30 @@ def test_reset(self):
         self.assertEqual(state.mask, 1)
 
     def test_step(self):
-        env = AtariEnvironment('Breakout')
+        env = AtariEnvironment("Breakout")
         env.reset()
         state = env.step(1)
         self.assertEqual(state.observation.shape, (1, 84, 84))
         self.assertEqual(state.reward, 0)
         self.assertFalse(state.done)
         self.assertEqual(state.mask, 1)
-        self.assertEqual(state['life_lost'], False)
+        self.assertEqual(state["life_lost"], False)
 
     def test_step_until_life_lost(self):
-        env = AtariEnvironment('Breakout')
+        env = AtariEnvironment("Breakout")
         env.reset()
         for _ in range(100):
             state = env.step(1)
-            if state['life_lost']:
+            if state["life_lost"]:
                 break
         self.assertEqual(state.observation.shape, (1, 84, 84))
         self.assertEqual(state.reward, 0)
         self.assertFalse(state.done)
         self.assertEqual(state.mask, 1)
-        self.assertEqual(state['life_lost'], True)
+        self.assertEqual(state["life_lost"], True)
 
     def test_step_until_done(self):
-        env = AtariEnvironment('Breakout')
+        env = AtariEnvironment("Breakout")
         env.reset()
         for _ in range(1000):
             state = env.step(1)
@@ -45,4 +46,4 @@ def test_step_until_done(self):
         self.assertEqual(state.reward, 0)
         self.assertTrue(state.done)
         self.assertEqual(state.mask, 0)
-        self.assertEqual(state['life_lost'], False)
+        self.assertEqual(state["life_lost"], False)
diff --git a/all/environments/duplicate_env_test.py b/all/environments/duplicate_env_test.py
index 1c0c750e..e506db17 100644
--- a/all/environments/duplicate_env_test.py
+++ b/all/environments/duplicate_env_test.py
@@ -1,18 +1,19 @@
 import unittest
-import gymnasium
+
 import torch
+
 from all.environments import DuplicateEnvironment, GymEnvironment
 
 
 def make_vec_env(num_envs=3):
-    env = [GymEnvironment('CartPole-v0') for i in range(num_envs)]
+    env = [GymEnvironment("CartPole-v0") for i in range(num_envs)]
     return env
 
 
 class DuplicateEnvironmentTest(unittest.TestCase):
     def test_env_name(self):
         env = DuplicateEnvironment(make_vec_env())
-        self.assertEqual(env.name, 'CartPole-v0')
+        self.assertEqual(env.name, "CartPole-v0")
 
     def test_num_envs(self):
         num_envs = 5
@@ -25,9 +26,30 @@ def test_reset(self):
         env = DuplicateEnvironment(make_vec_env(num_envs))
         state = env.reset()
         self.assertEqual(state.observation.shape, (num_envs, 4))
-        self.assertTrue((state.reward == torch.zeros(num_envs, )).all())
-        self.assertTrue((state.done == torch.zeros(num_envs, )).all())
-        self.assertTrue((state.mask == torch.ones(num_envs, )).all())
+        self.assertTrue(
+            (
+                state.reward
+                == torch.zeros(
+                    num_envs,
+                )
+            ).all()
+        )
+        self.assertTrue(
+            (
+                state.done
+                == torch.zeros(
+                    num_envs,
+                )
+            ).all()
+        )
+        self.assertTrue(
+            (
+                state.mask
+                == torch.ones(
+                    num_envs,
+                )
+            ).all()
+        )
 
     def test_step(self):
         num_envs = 5
@@ -35,9 +57,30 @@ def test_step(self):
         env.reset()
         state = env.step(torch.ones(num_envs, dtype=torch.int32))
         self.assertEqual(state.observation.shape, (num_envs, 4))
-        self.assertTrue((state.reward == torch.ones(num_envs, )).all())
-        self.assertTrue((state.done == torch.zeros(num_envs, )).all())
-        self.assertTrue((state.mask == torch.ones(num_envs, )).all())
+        self.assertTrue(
+            (
+                state.reward
+                == torch.ones(
+                    num_envs,
+                )
+            ).all()
+        )
+        self.assertTrue(
+            (
+                state.done
+                == torch.zeros(
+                    num_envs,
+                )
+            ).all()
+        )
+        self.assertTrue(
+            (
+                state.mask
+                == torch.ones(
+                    num_envs,
+                )
+            ).all()
+        )
 
     def test_step_until_done(self):
         num_envs = 3
@@ -48,6 +91,6 @@ def test_step_until_done(self):
             if state.done[0]:
                 break
         self.assertEqual(state[0].observation.shape, (4,))
-        self.assertEqual(state[0].reward, 1.)
+        self.assertEqual(state[0].reward, 1.0)
         self.assertTrue(state[0].done)
         self.assertEqual(state[0].mask, 0)
diff --git a/all/environments/gym_test.py b/all/environments/gym_test.py
index 31dea2e2..29abcd34 100644
--- a/all/environments/gym_test.py
+++ b/all/environments/gym_test.py
@@ -1,15 +1,15 @@
 import unittest
-import gymnasium
+
 from all.environments import GymEnvironment
 
 
 class GymEnvironmentTest(unittest.TestCase):
     def test_env_name(self):
-        env = GymEnvironment('CartPole-v0')
-        self.assertEqual(env.name, 'CartPole-v0')
+        env = GymEnvironment("CartPole-v0")
+        self.assertEqual(env.name, "CartPole-v0")
 
     def test_reset(self):
-        env = GymEnvironment('CartPole-v0')
+        env = GymEnvironment("CartPole-v0")
         state = env.reset()
         self.assertEqual(state.observation.shape, (4,))
         self.assertEqual(state.reward, 0)
@@ -17,22 +17,22 @@ def test_reset(self):
         self.assertEqual(state.mask, 1)
 
     def test_step(self):
-        env = GymEnvironment('CartPole-v0')
+        env = GymEnvironment("CartPole-v0")
         env.reset()
         state = env.step(1)
         self.assertEqual(state.observation.shape, (4,))
-        self.assertEqual(state.reward, 1.)
+        self.assertEqual(state.reward, 1.0)
         self.assertFalse(state.done)
         self.assertEqual(state.mask, 1)
 
     def test_step_until_done(self):
-        env = GymEnvironment('CartPole-v0')
+        env = GymEnvironment("CartPole-v0")
         env.reset()
         for _ in range(100):
             state = env.step(1)
             if state.done:
                 break
         self.assertEqual(state.observation.shape, (4,))
-        self.assertEqual(state.reward, 1.)
+        self.assertEqual(state.reward, 1.0)
         self.assertTrue(state.done)
         self.assertEqual(state.mask, 0)
diff --git a/all/environments/mujoco_test.py b/all/environments/mujoco_test.py
index 7a207e8c..7fb240b7 100644
--- a/all/environments/mujoco_test.py
+++ b/all/environments/mujoco_test.py
@@ -1,11 +1,12 @@
 import unittest
-from all.environments import MujocoEnvironment, GymEnvironment
+
+from all.environments import MujocoEnvironment
 
 
 class MujocoEnvironmentTest(unittest.TestCase):
     def test_load_env(self):
         env = MujocoEnvironment("Ant-v4")
-        self.assertEqual(env.name, 'Ant-v4')
+        self.assertEqual(env.name, "Ant-v4")
 
     def test_observation_space(self):
         env = MujocoEnvironment("Ant-v4")
@@ -19,7 +20,7 @@ def test_reset(self):
         env = MujocoEnvironment("Ant-v4")
         state = env.reset(seed=0)
         self.assertEqual(state.observation.shape, (27,))
-        self.assertEqual(state.reward, 0.)
+        self.assertEqual(state.reward, 0.0)
         self.assertFalse(state.done)
         self.assertEqual(state.mask, 1)
 
@@ -28,8 +29,8 @@ def test_step(self):
         state = env.reset(seed=0)
         state = env.step(env.action_space.sample())
         self.assertEqual(state.observation.shape, (27,))
-        self.assertGreater(state.reward, -2.)
+        self.assertGreater(state.reward, -2.0)
         self.assertLess(state.reward, 2)
-        self.assertNotEqual(state.reward, 0.)
+        self.assertNotEqual(state.reward, 0.0)
         self.assertFalse(state.done)
         self.assertEqual(state.mask, 1)
diff --git a/all/environments/multiagent_atari_test.py b/all/environments/multiagent_atari_test.py
index c5c5e7a8..78c779c9 100644
--- a/all/environments/multiagent_atari_test.py
+++ b/all/environments/multiagent_atari_test.py
@@ -1,82 +1,84 @@
 import unittest
+
 import torch
+
 from all.environments import MultiagentAtariEnv
 
 
 class MultiagentAtariEnvTest(unittest.TestCase):
     def test_init(self):
-        MultiagentAtariEnv('pong_v3', device='cpu')
-        MultiagentAtariEnv('mario_bros_v3', device='cpu')
-        MultiagentAtariEnv('entombed_cooperative_v3', device='cpu')
+        MultiagentAtariEnv("pong_v3", device="cpu")
+        MultiagentAtariEnv("mario_bros_v3", device="cpu")
+        MultiagentAtariEnv("entombed_cooperative_v3", device="cpu")
 
     def test_reset(self):
-        env = MultiagentAtariEnv('pong_v3', device='cpu')
+        env = MultiagentAtariEnv("pong_v3", device="cpu")
         state = env.reset()
         self.assertEqual(state.observation.shape, (1, 84, 84))
         self.assertEqual(state.reward, 0)
         self.assertEqual(state.done, False)
-        self.assertEqual(state.mask, 1.)
-        self.assertEqual(state['agent'], 'first_0')
+        self.assertEqual(state.mask, 1.0)
+        self.assertEqual(state["agent"], "first_0")
 
     def test_step(self):
-        env = MultiagentAtariEnv('pong_v3', device='cpu')
+        env = MultiagentAtariEnv("pong_v3", device="cpu")
         env.reset()
         state = env.step(0)
         self.assertEqual(state.observation.shape, (1, 84, 84))
         self.assertEqual(state.reward, 0)
         self.assertEqual(state.done, False)
-        self.assertEqual(state.mask, 1.)
-        self.assertEqual(state['agent'], 'second_0')
+        self.assertEqual(state.mask, 1.0)
+        self.assertEqual(state["agent"], "second_0")
 
     def test_step_tensor(self):
-        env = MultiagentAtariEnv('pong_v3', device='cpu')
+        env = MultiagentAtariEnv("pong_v3", device="cpu")
         env.reset()
         state = env.step(torch.tensor([0]))
         self.assertEqual(state.observation.shape, (1, 84, 84))
         self.assertEqual(state.reward, 0)
         self.assertEqual(state.done, False)
-        self.assertEqual(state.mask, 1.)
-        self.assertEqual(state['agent'], 'second_0')
+        self.assertEqual(state.mask, 1.0)
+        self.assertEqual(state["agent"], "second_0")
 
     def test_name(self):
-        env = MultiagentAtariEnv('pong_v3', device='cpu')
-        self.assertEqual(env.name, 'pong_v3')
+        env = MultiagentAtariEnv("pong_v3", device="cpu")
+        self.assertEqual(env.name, "pong_v3")
 
     def test_agent_iter(self):
-        env = MultiagentAtariEnv('pong_v3', device='cpu')
+        env = MultiagentAtariEnv("pong_v3", device="cpu")
         env.reset()
         it = iter(env.agent_iter())
-        self.assertEqual(next(it), 'first_0')
+        self.assertEqual(next(it), "first_0")
 
     def test_state_spaces(self):
-        env = MultiagentAtariEnv('pong_v3', device='cpu')
-        self.assertEqual(env.state_space('first_0').shape, (1, 84, 84))
-        self.assertEqual(env.state_space('second_0').shape, (1, 84, 84))
+        env = MultiagentAtariEnv("pong_v3", device="cpu")
+        self.assertEqual(env.state_space("first_0").shape, (1, 84, 84))
+        self.assertEqual(env.state_space("second_0").shape, (1, 84, 84))
 
     def test_action_spaces(self):
-        env = MultiagentAtariEnv('pong_v3', device='cpu')
-        self.assertEqual(env.action_space('first_0').n, 6)
-        self.assertEqual(env.action_space('second_0').n, 6)
+        env = MultiagentAtariEnv("pong_v3", device="cpu")
+        self.assertEqual(env.action_space("first_0").n, 6)
+        self.assertEqual(env.action_space("second_0").n, 6)
 
     def test_list_agents(self):
-        env = MultiagentAtariEnv('pong_v3', device='cpu')
-        self.assertEqual(env.agents, ['first_0', 'second_0'])
+        env = MultiagentAtariEnv("pong_v3", device="cpu")
+        self.assertEqual(env.agents, ["first_0", "second_0"])
 
     def test_is_done(self):
-        env = MultiagentAtariEnv('pong_v3', device='cpu')
+        env = MultiagentAtariEnv("pong_v3", device="cpu")
         env.reset()
-        self.assertFalse(env.is_done('first_0'))
-        self.assertFalse(env.is_done('second_0'))
+        self.assertFalse(env.is_done("first_0"))
+        self.assertFalse(env.is_done("second_0"))
 
     def test_last(self):
-        env = MultiagentAtariEnv('pong_v3', device='cpu')
+        env = MultiagentAtariEnv("pong_v3", device="cpu")
         env.reset()
         state = env.last()
         self.assertEqual(state.observation.shape, (1, 84, 84))
         self.assertEqual(state.reward, 0)
         self.assertEqual(state.done, False)
-        self.assertEqual(state.mask, 1.)
-        self.assertEqual(state['agent'], 'first_0')
+        self.assertEqual(state.mask, 1.0)
+        self.assertEqual(state["agent"], "first_0")
 
 
 if __name__ == "__main__":
diff --git a/all/environments/multiagent_pettingzoo_test.py b/all/environments/multiagent_pettingzoo_test.py
index 482e50e5..184728af 100644
--- a/all/environments/multiagent_pettingzoo_test.py
+++ b/all/environments/multiagent_pettingzoo_test.py
@@ -1,8 +1,9 @@
 import unittest
-import torch
-from all.environments import MultiagentPettingZooEnv
+
 from pettingzoo.mpe import simple_world_comm_v3
 
+from all.environments import MultiagentPettingZooEnv
+
 
 class MultiagentPettingZooEnvTest(unittest.TestCase):
     def test_init(self):
@@ -14,8 +15,8 @@ def test_reset(self):
         self.assertEqual(state.observation.shape, (34,))
         self.assertEqual(state.reward, 0)
         self.assertEqual(state.done, False)
-        self.assertEqual(state.mask, 1.)
-        self.assertEqual(state['agent'], 'leadadversary_0')
+        self.assertEqual(state.mask, 1.0)
+        self.assertEqual(state["agent"], "leadadversary_0")
 
     def test_step(self):
         env = self._make_env()
@@ -24,8 +25,8 @@ def test_step(self):
         self.assertEqual(state.observation.shape, (34,))
         self.assertEqual(state.reward, 0)
         self.assertEqual(state.done, False)
-        self.assertEqual(state.mask, 1.)
-        self.assertEqual(state['agent'], 'adversary_0')
+        self.assertEqual(state.mask, 1.0)
+        self.assertEqual(state["agent"], "adversary_0")
 
     def test_step_tensor(self):
         env = self._make_env()
@@ -34,38 +35,48 @@ def test_step_tensor(self):
         self.assertEqual(state.observation.shape, (34,))
         self.assertEqual(state.reward, 0)
         self.assertEqual(state.done, False)
-        self.assertEqual(state.mask, 1.)
-        self.assertEqual(state['agent'], 'adversary_0')
+        self.assertEqual(state.mask, 1.0)
+        self.assertEqual(state["agent"], "adversary_0")
 
     def test_name(self):
         env = self._make_env()
-        self.assertEqual(env.name, 'simple_world_comm_v3')
+        self.assertEqual(env.name, "simple_world_comm_v3")
 
     def test_agent_iter(self):
         env = self._make_env()
         env.reset()
         it = iter(env.agent_iter())
-        self.assertEqual(next(it), 'leadadversary_0')
+        self.assertEqual(next(it), "leadadversary_0")
 
     def test_state_spaces(self):
         env = self._make_env()
-        self.assertEqual(env.state_space('leadadversary_0').shape, (34,))
-        self.assertEqual(env.state_space('adversary_0').shape, (34,))
+        self.assertEqual(env.state_space("leadadversary_0").shape, (34,))
+        self.assertEqual(env.state_space("adversary_0").shape, (34,))
 
     def test_action_spaces(self):
         env = self._make_env()
-        self.assertEqual(env.action_space('leadadversary_0').n, 20)
-        self.assertEqual(env.action_space('adversary_0').n, 5)
+        self.assertEqual(env.action_space("leadadversary_0").n, 20)
+        self.assertEqual(env.action_space("adversary_0").n, 5)
 
     def test_list_agents(self):
         env = self._make_env()
-        self.assertEqual(env.agents, ['leadadversary_0', 'adversary_0', 'adversary_1', 'adversary_2', 'agent_0', 'agent_1'])
+        self.assertEqual(
+            env.agents,
+            [
+                "leadadversary_0",
+                "adversary_0",
+                "adversary_1",
+                "adversary_2",
+                "agent_0",
+                "agent_1",
+            ],
+        )
 
     def test_terminated(self):
         env = self._make_env()
         env.reset()
-        self.assertFalse(env.is_done('leadadversary_0'))
-        self.assertFalse(env.is_done('adversary_0'))
+        self.assertFalse(env.is_done("leadadversary_0"))
+        self.assertFalse(env.is_done("adversary_0"))
 
     def test_last(self):
         env = self._make_env()
@@ -74,20 +85,28 @@ def test_last(self):
         self.assertEqual(state.observation.shape, (34,))
         self.assertEqual(state.reward, 0)
         self.assertEqual(state.done, False)
-        self.assertEqual(state.mask, 1.)
-        self.assertEqual(state['agent'], 'leadadversary_0')
+        self.assertEqual(state.mask, 1.0)
+        self.assertEqual(state["agent"], "leadadversary_0")
 
     def test_variable_spaces(self):
-        env = MultiagentPettingZooEnv(simple_world_comm_v3.env(), name="simple_world_comm_v2", device='cpu')
+        env = MultiagentPettingZooEnv(
+            simple_world_comm_v3.env(), name="simple_world_comm_v2", device="cpu"
+        )
         state = env.reset()
         # tests that action spaces work
         for agent in env.agents:
             state = env.last()
-            self.assertTrue(env.observation_space(agent).contains(state['observation'].cpu().detach().numpy()))
+            self.assertTrue(
+                env.observation_space(agent).contains(
+                    state["observation"].cpu().detach().numpy()
+                )
+            )
             env.step(env.action_space(env.agent_selection).sample())
 
     def _make_env(self):
-        return MultiagentPettingZooEnv(simple_world_comm_v3.env(), name="simple_world_comm_v3", device='cpu')
+        return MultiagentPettingZooEnv(
+            simple_world_comm_v3.env(), name="simple_world_comm_v3", device="cpu"
+        )
 
 
 if __name__ == "__main__":
diff --git a/all/environments/pybullet_test.py b/all/environments/pybullet_test.py
index d08e5bfb..7dcda1dd 100644
--- a/all/environments/pybullet_test.py
+++ b/all/environments/pybullet_test.py
@@ -1,5 +1,6 @@
 import unittest
-from all.environments import PybulletEnvironment, GymEnvironment
+
+from all.environments import PybulletEnvironment
 
 
 class PybulletEnvironmentTest(unittest.TestCase):
@@ -9,25 +10,25 @@ def test_env_short_name(self):
             self.assertEqual(env.name, long_name)
 
     def test_env_full_name(self):
-        env = PybulletEnvironment('HalfCheetahBulletEnv-v0')
-        self.assertEqual(env.name, 'HalfCheetahBulletEnv-v0')
+        env = PybulletEnvironment("HalfCheetahBulletEnv-v0")
+        self.assertEqual(env.name, "HalfCheetahBulletEnv-v0")
 
     def test_reset(self):
-        env = PybulletEnvironment('cheetah')
+        env = PybulletEnvironment("cheetah")
         state = env.reset()
         self.assertEqual(state.observation.shape, (26,))
-        self.assertEqual(state.reward, 0.)
+        self.assertEqual(state.reward, 0.0)
         self.assertFalse(state.done)
         self.assertEqual(state.mask, 1)
 
     def test_step(self):
-        env = PybulletEnvironment('cheetah')
+        env = PybulletEnvironment("cheetah")
         env.seed(0)
         state = env.reset()
         state = env.step(env.action_space.sample())
         self.assertEqual(state.observation.shape, (26,))
-        self.assertGreater(state.reward, -1.)
+        self.assertGreater(state.reward, -1.0)
         self.assertLess(state.reward, 1)
-        self.assertNotEqual(state.reward, 0.)
+        self.assertNotEqual(state.reward, 0.0)
         self.assertFalse(state.done)
         self.assertEqual(state.mask, 1)
diff --git a/all/environments/vector_env_test.py b/all/environments/vector_env_test.py
index a4cfba77..117c8a6f 100644
--- a/all/environments/vector_env_test.py
+++ b/all/environments/vector_env_test.py
@@ -1,18 +1,22 @@
 import unittest
+
 import gymnasium
 import torch
-from all.environments import GymVectorEnvironment, GymEnvironment, DuplicateEnvironment
+
+from all.environments import DuplicateEnvironment, GymEnvironment, GymVectorEnvironment
 
 
 def make_vec_env(num_envs=3):
-    env = gymnasium.vector.SyncVectorEnv([lambda: gymnasium.make('CartPole-v0')] * num_envs)
+    env = gymnasium.vector.SyncVectorEnv(
+        [lambda: gymnasium.make("CartPole-v0")] * num_envs
+    )
     return env
 
 
 class GymVectorEnvironmentTest(unittest.TestCase):
     def test_env_name(self):
         env = GymVectorEnvironment(make_vec_env(), "CartPole")
-        self.assertEqual(env.name, 'CartPole')
+        self.assertEqual(env.name, "CartPole")
 
     def test_num_envs(self):
         num_envs = 5
@@ -25,9 +29,30 @@ def test_reset(self):
         env = GymVectorEnvironment(make_vec_env(num_envs), "CartPole")
         state = env.reset()
         self.assertEqual(state.observation.shape, (num_envs, 4))
-        self.assertTrue((state.reward == torch.zeros(num_envs, )).all())
-        self.assertTrue((state.done == torch.zeros(num_envs, )).all())
-        self.assertTrue((state.mask == torch.ones(num_envs, )).all())
+        self.assertTrue(
+            (
+                state.reward
+                == torch.zeros(
+                    num_envs,
+                )
+            ).all()
+        )
+        self.assertTrue(
+            (
+                state.done
+                == torch.zeros(
+                    num_envs,
+                )
+            ).all()
+        )
+        self.assertTrue(
+            (
+                state.mask
+                == torch.ones(
+                    num_envs,
+                )
+            ).all()
+        )
 
     def test_step(self):
         num_envs = 5
@@ -35,9 +60,30 @@ def test_step(self):
         env.reset()
         state = env.step(torch.ones(num_envs, dtype=torch.int32))
         self.assertEqual(state.observation.shape, (num_envs, 4))
-        self.assertTrue((state.reward == torch.ones(num_envs, )).all())
-        self.assertTrue((state.done == torch.zeros(num_envs, )).all())
-        self.assertTrue((state.mask == torch.ones(num_envs, )).all())
+        self.assertTrue(
+            (
+                state.reward
+                == torch.ones(
+                    num_envs,
+                )
+            ).all()
+        )
+        self.assertTrue(
+            (
+                state.done
+                == torch.zeros(
+                    num_envs,
+                )
+            ).all()
+        )
+        self.assertTrue(
+            (
+                state.mask
+                == torch.ones(
+                    num_envs,
+                )
+            ).all()
+        )
 
     def test_step_until_done(self):
         num_envs = 3
@@ -50,14 +96,16 @@ def test_step_until_done(self):
         else:
             self.assertTrue(False)
         self.assertEqual(state[0].observation.shape, (4,))
-        self.assertEqual(state[0].reward, 1.)
+        self.assertEqual(state[0].reward, 1.0)
         self.assertTrue(state[0].done)
         self.assertEqual(state[0].mask, 0)
 
     def test_same_as_duplicate(self):
         n_envs = 3
         torch.manual_seed(42)
-        env1 = DuplicateEnvironment([GymEnvironment('CartPole-v0') for i in range(n_envs)])
+        env1 = DuplicateEnvironment(
+            [GymEnvironment("CartPole-v0") for i in range(n_envs)]
+        )
         env2 = GymVectorEnvironment(make_vec_env(n_envs), "CartPole-v0")
         state1 = env1.reset(seed=42)
         state2 = env2.reset(seed=42)
diff --git a/all/experiments/multiagent_env_experiment_test.py b/all/experiments/multiagent_env_experiment_test.py
index fbac6bbd..ad782c7b 100644
--- a/all/experiments/multiagent_env_experiment_test.py
+++ b/all/experiments/multiagent_env_experiment_test.py
@@ -1,17 +1,19 @@
 import random
 import unittest
+
 import numpy as np
 import torch
-from all.presets.atari import dqn
-from all.presets import IndependentMultiagentPreset
+
 from all.environments import MultiagentAtariEnv
 from all.experiments import MultiagentEnvExperiment
 from all.experiments.single_env_experiment_test import MockLogger
+from all.presets import IndependentMultiagentPreset
+from all.presets.atari import dqn
 
 
 class MockExperiment(MultiagentEnvExperiment):
     def _make_logger(self, logdir, agent_name, env_name, verbose, logger):
-        self._logger = MockLogger(self, agent_name + '_' + env_name, verbose)
+        self._logger = MockLogger(self, agent_name + "_" + env_name, verbose)
         return self._logger
 
 
@@ -20,57 +22,96 @@ def setUp(self):
         random.seed(0)
         np.random.seed(0)
         torch.manual_seed(0)
-        self.env = MultiagentAtariEnv('space_invaders_v2', device='cpu', seed=0)
+        self.env = MultiagentAtariEnv("space_invaders_v2", device="cpu", seed=0)
         self.env.reset(seed=0)
         self.experiment = None
 
     def test_adds_default_name(self):
-        experiment = MockExperiment(self.make_preset(), self.env, quiet=True, save_freq=float('inf'))
+        experiment = MockExperiment(
+            self.make_preset(), self.env, quiet=True, save_freq=float("inf")
+        )
         self.assertEqual(experiment._logger.label, "independent_space_invaders_v2")
 
     def test_adds_custom_name(self):
-        experiment = MockExperiment(self.make_preset(), self.env, name='custom', quiet=True, save_freq=float('inf'))
+        experiment = MockExperiment(
+            self.make_preset(),
+            self.env,
+            name="custom",
+            quiet=True,
+            save_freq=float("inf"),
+        )
         self.assertEqual(experiment._logger.label, "custom_space_invaders_v2")
 
     def test_writes_training_returns(self):
-        experiment = MockExperiment(self.make_preset(), self.env, quiet=True, save_freq=float('inf'))
+        experiment = MockExperiment(
+            self.make_preset(), self.env, quiet=True, save_freq=float("inf")
+        )
         experiment.train(episodes=3)
         self.maxDiff = None
         # could not get the exact numbers to be reproducible across enviornments :(
-        self.assertEqual(len(experiment._logger.data['eval/first_0/returns/frame']['values']), 3)
-        self.assertEqual(len(experiment._logger.data['eval/first_0/returns/frame']['steps']), 3)
-        self.assertEqual(len(experiment._logger.data['eval/second_0/returns/frame']['values']), 3)
-        self.assertEqual(len(experiment._logger.data['eval/second_0/returns/frame']['steps']), 3)
+        self.assertEqual(
+            len(experiment._logger.data["eval/first_0/returns/frame"]["values"]), 3
+        )
+        self.assertEqual(
+            len(experiment._logger.data["eval/first_0/returns/frame"]["steps"]), 3
+        )
+        self.assertEqual(
+            len(experiment._logger.data["eval/second_0/returns/frame"]["values"]), 3
+        )
+        self.assertEqual(
+            len(experiment._logger.data["eval/second_0/returns/frame"]["steps"]), 3
+        )
 
     def test_writes_test_returns(self):
-        experiment = MockExperiment(self.make_preset(), self.env, quiet=True, save_freq=float('inf'))
+        experiment = MockExperiment(
+            self.make_preset(), self.env, quiet=True, save_freq=float("inf")
+        )
         experiment.train(episodes=3)
         experiment._logger.data = {}
         experiment.test(episodes=3)
-        self.assertEqual(list(experiment._logger.data.keys()), [
-            'summary/first_0/returns-test/mean',
-            'summary/first_0/returns-test/std',
-            'summary/second_0/returns-test/mean',
-            'summary/second_0/returns-test/std'
-        ])
-        steps = experiment._logger.data['summary/first_0/returns-test/mean']['steps'][0]
+        self.assertEqual(
+            list(experiment._logger.data.keys()),
+            [
+                "summary/first_0/returns-test/mean",
+                "summary/first_0/returns-test/std",
+                "summary/second_0/returns-test/mean",
+                "summary/second_0/returns-test/std",
+            ],
+        )
+        steps = experiment._logger.data["summary/first_0/returns-test/mean"]["steps"][0]
         for datum in experiment._logger.data.values():
-            self.assertEqual(len(datum['values']), 1)
-            self.assertGreaterEqual(datum['values'][0], 0.0)
-            self.assertEqual(len(datum['steps']), 1)
-            self.assertEqual(datum['steps'][0], steps)
+            self.assertEqual(len(datum["values"]), 1)
+            self.assertGreaterEqual(datum["values"][0], 0.0)
+            self.assertEqual(len(datum["steps"]), 1)
+            self.assertEqual(datum["steps"][0], steps)
 
     def test_writes_loss(self):
-        experiment = MockExperiment(self.make_preset(), self.env, quiet=True, verbose=True, save_freq=float('inf'))
+        experiment = MockExperiment(
+            self.make_preset(),
+            self.env,
+            quiet=True,
+            verbose=True,
+            save_freq=float("inf"),
+        )
         self.assertTrue(experiment._logger.verbose)
-        experiment = MockExperiment(self.make_preset(), self.env, quiet=True, verbose=False, save_freq=float('inf'))
+        experiment = MockExperiment(
+            self.make_preset(),
+            self.env,
+            quiet=True,
+            verbose=False,
+            save_freq=float("inf"),
+        )
         self.assertFalse(experiment._logger.verbose)
 
     def make_preset(self):
-        return IndependentMultiagentPreset('independent', 'cpu', {
-            agent: dqn.device('cpu').env(env).build()
-            for agent, env in self.env.subenvs.items()
-        })
+        return IndependentMultiagentPreset(
+            "independent",
+            "cpu",
+            {
+                agent: dqn.device("cpu").env(env).build()
+                for agent, env in self.env.subenvs.items()
+            },
+        )
 
 
 if __name__ == "__main__":
diff --git a/all/experiments/parallel_env_experiment_test.py b/all/experiments/parallel_env_experiment_test.py
index 6fd88843..f65d9e49 100644
--- a/all/experiments/parallel_env_experiment_test.py
+++ b/all/experiments/parallel_env_experiment_test.py
@@ -1,10 +1,12 @@
 import unittest
+
 import numpy as np
 import torch
-from all.presets.classic_control import a2c
+
 from all.environments import GymEnvironment
 from all.experiments import ParallelEnvExperiment
 from all.experiments.single_env_experiment_test import MockLogger
+from all.presets.classic_control import a2c
 
 
 class MockExperiment(ParallelEnvExperiment):
diff --git a/all/experiments/single_env_experiment_test.py b/all/experiments/single_env_experiment_test.py
index e2e6adc6..1a7e50f5 100644
--- a/all/experiments/single_env_experiment_test.py
+++ b/all/experiments/single_env_experiment_test.py
@@ -1,10 +1,12 @@
 import unittest
+
 import numpy as np
 import torch
-from all.presets.classic_control import dqn
+
 from all.environments import GymEnvironment
 from all.experiments import SingleEnvExperiment
 from all.logging import Logger
+from all.presets.classic_control import dqn
 
 
 class MockLogger(Logger):
@@ -76,7 +78,7 @@ def test_writes_training_returns_frame(self):
         experiment.train(episodes=3)
         np.testing.assert_equal(
             experiment._logger.data["eval/returns/frame"]["values"],
-            np.array([22., 17., 28.]),
+            np.array([22.0, 17.0, 28.0]),
         )
         np.testing.assert_equal(
             experiment._logger.data["eval/returns/frame"]["steps"],
diff --git a/all/experiments/watch_test.py b/all/experiments/watch_test.py
index 6327cd69..2aaab832 100644
--- a/all/experiments/watch_test.py
+++ b/all/experiments/watch_test.py
@@ -1,17 +1,19 @@
 import unittest
 from unittest import mock
+
 import torch
+
 from all.environments import GymEnvironment
 from all.experiments.watch import load_and_watch
 
 
-class MockAgent():
+class MockAgent:
     def act(self):
         # sample from cartpole action space
         return torch.randint(0, 2, [])
 
 
-class MockPreset():
+class MockPreset:
     def __init__(self, filename):
         self.filename = filename
 
@@ -20,9 +22,9 @@ def test_agent(self):
 
 
 class WatchTest(unittest.TestCase):
-    @mock.patch('torch.load', lambda filename: MockPreset(filename))
-    @mock.patch('time.sleep', mock.MagicMock())
-    @mock.patch('sys.stdout', mock.MagicMock())
+    @mock.patch("torch.load", lambda filename: MockPreset(filename))
+    @mock.patch("time.sleep", mock.MagicMock())
+    @mock.patch("sys.stdout", mock.MagicMock())
     def test_load_and_watch(self):
         env = mock.MagicMock(GymEnvironment("CartPole-v0", render_mode="rgb_array"))
         load_and_watch("file.name", env, n_episodes=3)
diff --git a/all/memory/advantage_test.py b/all/memory/advantage_test.py
index 2223e4d4..368eecd4 100644
--- a/all/memory/advantage_test.py
+++ b/all/memory/advantage_test.py
@@ -1,9 +1,11 @@
 import unittest
+
 import torch
 import torch_testing as tt
+
 from all import nn
+from all.approximation import FeatureNetwork, VNetwork
 from all.core import StateArray
-from all.approximation import VNetwork, FeatureNetwork
 from all.memory import NStepAdvantageBuffer
 
 
@@ -30,21 +32,22 @@ def test_rollout(self):
 
         expected_states = StateArray(torch.arange(0, 6).unsqueeze(1).float(), (6,))
         expected_next_states = StateArray(
-            torch.cat((torch.arange(6, 9), torch.arange(6, 9))).unsqueeze(1).float(), (6,)
+            torch.cat((torch.arange(6, 9), torch.arange(6, 9))).unsqueeze(1).float(),
+            (6,),
         )
-        expected_returns = torch.tensor([
-            0.5, 0.5, 0.5,
-            1, 1, 1
-        ]).float()
-        expected_lengths = torch.tensor([
-            2., 2, 2,
-            1, 1, 1
-        ])
+        expected_returns = torch.tensor([0.5, 0.5, 0.5, 1, 1, 1]).float()
+        expected_lengths = torch.tensor([2.0, 2, 2, 1, 1, 1])
 
         self.assert_states_equal(states, expected_states)
-        tt.assert_allclose(advantages, self._compute_expected_advantages(
-            expected_states, expected_returns, expected_next_states, expected_lengths
-        ))
+        tt.assert_allclose(
+            advantages,
+            self._compute_expected_advantages(
+                expected_states,
+                expected_returns,
+                expected_next_states,
+                expected_lengths,
+            ),
+        )
 
     def test_rollout_with_dones(self):
         buffer = NStepAdvantageBuffer(self.v, self.features, 3, 3, discount_factor=0.5)
@@ -59,31 +62,31 @@ def test_rollout_with_dones(self):
         buffer.store(states[6:9], actions, 2 * torch.ones(3))
         states, actions, advantages = buffer.advantages(states[9:12])
 
-        expected_states = StateArray(torch.arange(0, 9).unsqueeze(1).float(), (9,), done=done[0:9])
+        expected_states = StateArray(
+            torch.arange(0, 9).unsqueeze(1).float(), (9,), done=done[0:9]
+        )
         expected_next_done = torch.tensor([True] * 9)
         expected_next_done[5] = False
         expected_next_done[7] = False
         expected_next_done[8] = False
-        expected_next_states = StateArray(torch.tensor([
-            9, 7, 5,
-            9, 7, 11,
-            9, 10, 11
-        ]).unsqueeze(1).float(), (9,), done=expected_next_done)
-        expected_returns = torch.tensor([
-            1, 0.5, 0,
-            2, 1, 2,
-            2, 2, 2
-        ]).float()
-        expected_lengths = torch.tensor([
-            3, 2, 1,
-            2, 1, 2,
-            1, 1, 1
-        ]).float()
+        expected_next_states = StateArray(
+            torch.tensor([9, 7, 5, 9, 7, 11, 9, 10, 11]).unsqueeze(1).float(),
+            (9,),
+            done=expected_next_done,
+        )
+        expected_returns = torch.tensor([1, 0.5, 0, 2, 1, 2, 2, 2, 2]).float()
+        expected_lengths = torch.tensor([3, 2, 1, 2, 1, 2, 1, 1, 1]).float()
 
         self.assert_states_equal(states, expected_states)
-        tt.assert_allclose(advantages, self._compute_expected_advantages(
-            expected_states, expected_returns, expected_next_states, expected_lengths
-        ))
+        tt.assert_allclose(
+            advantages,
+            self._compute_expected_advantages(
+                expected_states,
+                expected_returns,
+                expected_next_states,
+                expected_lengths,
+            ),
+        )
 
     def test_multi_rollout(self):
         buffer = NStepAdvantageBuffer(self.v, self.features, 2, 2, discount_factor=0.5)
@@ -95,15 +98,20 @@ def test_multi_rollout(self):
         states, actions, advantages = buffer.advantages(raw_states[4:6])
         expected_states = StateArray(torch.arange(0, 4).unsqueeze(1).float(), (4,))
         expected_returns = torch.tensor([1.5, 1.5, 1, 1])
-        expected_next_states = StateArray(torch.tensor([4., 5, 4, 5]).unsqueeze(1), (4,))
-        expected_lengths = torch.tensor([2., 2, 1, 1])
+        expected_next_states = StateArray(
+            torch.tensor([4.0, 5, 4, 5]).unsqueeze(1), (4,)
+        )
+        expected_lengths = torch.tensor([2.0, 2, 1, 1])
         self.assert_states_equal(states, expected_states)
-        tt.assert_allclose(advantages, self._compute_expected_advantages(
-            expected_states,
-            expected_returns,
-            expected_next_states,
-            expected_lengths
-        ))
+        tt.assert_allclose(
+            advantages,
+            self._compute_expected_advantages(
+                expected_states,
+                expected_returns,
+                expected_next_states,
+                expected_lengths,
+            ),
+        )
 
         buffer.store(raw_states[4:6], actions, torch.ones(2))
         buffer.store(raw_states[6:8], actions, torch.ones(2))
@@ -111,22 +119,28 @@ def test_multi_rollout(self):
         states, actions, advantages = buffer.advantages(raw_states[8:10])
         expected_states = StateArray(torch.arange(4, 8).unsqueeze(1).float(), (4,))
         self.assert_states_equal(states, expected_states)
-        tt.assert_allclose(advantages, self._compute_expected_advantages(
-            expected_states,
-            torch.tensor([1.5, 1.5, 1, 1]),
-            StateArray(torch.tensor([8, 9, 8, 9]).unsqueeze(1).float(), (4,)),
-            torch.tensor([2., 2, 1, 1])
-        ))
+        tt.assert_allclose(
+            advantages,
+            self._compute_expected_advantages(
+                expected_states,
+                torch.tensor([1.5, 1.5, 1, 1]),
+                StateArray(torch.tensor([8, 9, 8, 9]).unsqueeze(1).float(), (4,)),
+                torch.tensor([2.0, 2, 1, 1]),
+            ),
+        )
 
     def assert_array_equal(self, actual, expected):
         for i, exp in enumerate(expected):
-            self.assertEqual(actual[i], exp, msg=(
-                ("\nactual: %s\nexpected: %s") % (actual, expected)))
+            self.assertEqual(
+                actual[i],
+                exp,
+                msg=(("\nactual: %s\nexpected: %s") % (actual, expected)),
+            )
 
     def assert_states_equal(self, actual, expected):
         tt.assert_almost_equal(actual.observation, expected.observation)
         tt.assert_equal(actual.mask, expected.mask)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/all/memory/generalized_advantage_test.py b/all/memory/generalized_advantage_test.py
index 593938a4..37cd0b05 100644
--- a/all/memory/generalized_advantage_test.py
+++ b/all/memory/generalized_advantage_test.py
@@ -1,9 +1,11 @@
 import unittest
+
 import torch
 import torch_testing as tt
+
 from all import nn
+from all.approximation import FeatureNetwork, VNetwork
 from all.core import State
-from all.approximation import VNetwork, FeatureNetwork
 from all.memory import GeneralizedAdvantageBuffer
 
 
@@ -16,27 +18,26 @@ def setUp(self):
     def _compute_expected_advantages(self, states, returns, next_states, lengths):
         return (
             returns
-            + (0.5 ** lengths) * self.v.eval(self.features.eval(next_states))
+            + (0.5**lengths) * self.v.eval(self.features.eval(next_states))
             - self.v.eval(self.features.eval(states))
         )
 
     def test_simple(self):
         buffer = GeneralizedAdvantageBuffer(
-            self.v,
-            self.features,
-            2,
-            1,
-            discount_factor=0.5,
-            lam=0.5
+            self.v, self.features, 2, 1, discount_factor=0.5, lam=0.5
         )
         actions = torch.ones((1))
-        states = State.array([State({'observation': torch.tensor([float(x)])}) for x in range(3)])
-        rewards = torch.tensor([1., 2, 4])
+        states = State.array(
+            [State({"observation": torch.tensor([float(x)])}) for x in range(3)]
+        )
+        rewards = torch.tensor([1.0, 2, 4])
         buffer.store(states[0], actions, rewards[0])
         buffer.store(states[1], actions, rewards[1])
 
         values = self.v.eval(self.features.eval(states))
-        tt.assert_almost_equal(values, torch.tensor([0.1826, -0.3476, -0.8777]), decimal=3)
+        tt.assert_almost_equal(
+            values, torch.tensor([0.1826, -0.3476, -0.8777]), decimal=3
+        )
 
         td_errors = torch.zeros(2)
         td_errors[0] = rewards[0] + 0.5 * values[1] - values[0]
@@ -54,66 +55,66 @@ def test_simple(self):
 
     def test_parallel(self):
         buffer = GeneralizedAdvantageBuffer(
-            self.v,
-            self.features,
-            2,
-            2,
-            discount_factor=0.5,
-            lam=0.5
+            self.v, self.features, 2, 2, discount_factor=0.5, lam=0.5
         )
         actions = torch.ones((2))
 
         def make_states(x, y):
-            return State.array([
-                State({'observation': torch.tensor([float(x)])}),
-                State({'observation': torch.tensor([float(y)])})
-            ])
-
-        states = State.array([
-            make_states(0, 3),
-            make_states(1, 4),
-            make_states(2, 5),
-        ])
+            return State.array(
+                [
+                    State({"observation": torch.tensor([float(x)])}),
+                    State({"observation": torch.tensor([float(y)])}),
+                ]
+            )
+
+        states = State.array(
+            [
+                make_states(0, 3),
+                make_states(1, 4),
+                make_states(2, 5),
+            ]
+        )
         self.assertEqual(states.shape, (3, 2))
-        rewards = torch.tensor([[1., 1], [2, 1], [4, 1]])
+        rewards = torch.tensor([[1.0, 1], [2, 1], [4, 1]])
         buffer.store(states[0], actions, rewards[0])
         buffer.store(states[1], actions, rewards[1])
 
         values = self.v.eval(self.features.eval(states)).view(3, -1)
-        tt.assert_almost_equal(values, torch.tensor([
-            [0.183, -1.408],
-            [-0.348, -1.938],
-            [-0.878, -2.468]
-        ]), decimal=3)
+        tt.assert_almost_equal(
+            values,
+            torch.tensor([[0.183, -1.408], [-0.348, -1.938], [-0.878, -2.468]]),
+            decimal=3,
+        )
 
         td_errors = torch.zeros(2, 2)
         td_errors[0] = rewards[0] + 0.5 * values[1] - values[0]
         td_errors[1] = rewards[1] + 0.5 * values[2] - values[1]
-        tt.assert_almost_equal(td_errors, torch.tensor([
-            [0.6436, 1.439],
-            [1.909, 1.704]
-        ]), decimal=3)
+        tt.assert_almost_equal(
+            td_errors, torch.tensor([[0.6436, 1.439], [1.909, 1.704]]), decimal=3
+        )
 
         advantages = torch.zeros(2, 2)
         advantages[0] = td_errors[0] + 0.25 * td_errors[1]
         advantages[1] = td_errors[1]
-        tt.assert_almost_equal(advantages, torch.tensor([
-            [1.121, 1.865],
-            [1.909, 1.704]
-        ]), decimal=3)
+        tt.assert_almost_equal(
+            advantages, torch.tensor([[1.121, 1.865], [1.909, 1.704]]), decimal=3
+        )
 
         _states, _actions, _advantages = buffer.advantages(states[2])
         tt.assert_almost_equal(_advantages, advantages.view(-1))
 
     def assert_array_equal(self, actual, expected):
         for i, exp in enumerate(expected):
-            self.assertEqual(actual[i], exp, msg=(
-                ("\nactual: %s\nexpected: %s") % (actual, expected)))
+            self.assertEqual(
+                actual[i],
+                exp,
+                msg=(("\nactual: %s\nexpected: %s") % (actual, expected)),
+            )
 
     def assert_states_equal(self, actual, expected):
         tt.assert_almost_equal(actual.raw, expected.raw)
         tt.assert_equal(actual.mask, expected.mask)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/all/memory/replay_buffer_test.py b/all/memory/replay_buffer_test.py
index ccaa6202..52a4f40c 100644
--- a/all/memory/replay_buffer_test.py
+++ b/all/memory/replay_buffer_test.py
@@ -1,13 +1,15 @@
-import unittest
 import random
-import torch
+import unittest
+
 import numpy as np
+import torch
 import torch_testing as tt
+
 from all.core import State, StateArray
 from all.memory import (
     ExperienceReplayBuffer,
-    PrioritizedReplayBuffer,
     NStepReplayBuffer,
+    PrioritizedReplayBuffer,
 )
 
 
@@ -52,17 +54,21 @@ def test_run(self):
 
     def test_store_device(self):
         if torch.cuda.is_available():
-            self.replay_buffer = ExperienceReplayBuffer(5, device='cuda', store_device='cpu')
+            self.replay_buffer = ExperienceReplayBuffer(
+                5, device="cuda", store_device="cpu"
+            )
 
-            states = torch.arange(0, 20).to('cuda')
-            actions = torch.arange(0, 20).view((-1, 1)).to('cuda')
-            rewards = torch.arange(0, 20).to('cuda')
+            states = torch.arange(0, 20).to("cuda")
+            actions = torch.arange(0, 20).view((-1, 1)).to("cuda")
+            rewards = torch.arange(0, 20).to("cuda")
             state = State(states[0])
             next_state = State(states[1], reward=rewards[1])
             self.replay_buffer.store(state, actions[0], next_state)
             sample = self.replay_buffer.sample(3)
-            self.assertEqual(sample[0].device, torch.device('cuda'))
-            self.assertEqual(self.replay_buffer.buffer[0][0].device, torch.device('cpu'))
+            self.assertEqual(sample[0].device, torch.device("cuda"))
+            self.assertEqual(
+                self.replay_buffer.buffer[0][0].device, torch.device("cpu")
+            )
 
 
 class TestPrioritizedReplayBuffer(unittest.TestCase):
@@ -73,7 +79,9 @@ def setUp(self):
         self.replay_buffer = PrioritizedReplayBuffer(5, 0.6)
 
     def test_run(self):
-        states = StateArray(torch.arange(0, 20), (20,), reward=torch.arange(-1, 19).float())
+        states = StateArray(
+            torch.arange(0, 20), (20,), reward=torch.arange(-1, 19).float()
+        )
         actions = torch.arange(0, 20).view((-1, 1))
         expected_samples = State(
             torch.tensor(
@@ -127,7 +135,9 @@ def setUp(self):
         self.replay_buffer = NStepReplayBuffer(4, 0.5, ExperienceReplayBuffer(100))
 
     def test_run(self):
-        states = StateArray(torch.arange(0, 20), (20,), reward=torch.arange(-1, 19).float())
+        states = StateArray(
+            torch.arange(0, 20), (20,), reward=torch.arange(-1, 19).float()
+        )
         actions = torch.arange(0, 20)
 
         for i in range(3):
@@ -141,22 +151,24 @@ def test_run(self):
         sample = self.replay_buffer.buffer.buffer[0]
         self.assert_states_equal(sample[0], states[0])
         tt.assert_equal(sample[1], actions[0])
-        tt.assert_equal(sample[2].reward, torch.tensor(0 + 1 * 0.5 + 2 * 0.25 + 3 * 0.125))
+        tt.assert_equal(
+            sample[2].reward, torch.tensor(0 + 1 * 0.5 + 2 * 0.25 + 3 * 0.125)
+        )
         tt.assert_equal(
             self.replay_buffer.buffer.buffer[1][2].reward,
             torch.tensor(1 + 2 * 0.5 + 3 * 0.25 + 4 * 0.125),
         )
 
     def test_done(self):
-        state = State(torch.tensor(1), reward=1.)
+        state = State(torch.tensor(1), reward=1.0)
         action = torch.tensor(0)
-        done_state = State(torch.tensor(1), reward=1., done=True)
+        done_state = State(torch.tensor(1), reward=1.0, done=True)
 
         self.replay_buffer.store(state, action, done_state)
         self.assertEqual(len(self.replay_buffer), 1)
         sample = self.replay_buffer.buffer.buffer[0]
         self.assert_states_equal(state, sample[0])
-        self.assertEqual(sample[2].reward, 1.)
+        self.assertEqual(sample[2].reward, 1.0)
 
         self.replay_buffer.store(state, action, state)
         self.replay_buffer.store(state, action, state)
diff --git a/all/nn/nn_test.py b/all/nn/nn_test.py
index 5e465532..4380e14f 100644
--- a/all/nn/nn_test.py
+++ b/all/nn/nn_test.py
@@ -1,8 +1,10 @@
 import unittest
+
+import gymnasium
 import numpy as np
 import torch
 import torch_testing as tt
-import gymnasium
+
 from all import nn
 from all.core import StateArray
 
diff --git a/all/optim/scheduler_test.py b/all/optim/scheduler_test.py
index 3a66d237..e96a2152 100644
--- a/all/optim/scheduler_test.py
+++ b/all/optim/scheduler_test.py
@@ -1,6 +1,8 @@
 import unittest
+
 import numpy as np
-from all.optim import Schedulable, LinearScheduler
+
+from all.optim import LinearScheduler, Schedulable
 
 
 class Obj(Schedulable):
@@ -17,5 +19,5 @@ def test_linear_scheduler(self):
         np.testing.assert_allclose(actual, expected)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/all/policies/deterministic_test.py b/all/policies/deterministic_test.py
index 683800a6..9b42077e 100644
--- a/all/policies/deterministic_test.py
+++ b/all/policies/deterministic_test.py
@@ -1,10 +1,12 @@
 import unittest
+
+import numpy as np
 import torch
 import torch_testing as tt
-import numpy as np
 from gymnasium.spaces import Box
+
 from all import nn
-from all.approximation import FixedTarget, DummyCheckpointer
+from all.approximation import DummyCheckpointer, FixedTarget
 from all.core import State
 from all.policies import DeterministicPolicy
 
@@ -15,16 +17,11 @@
 class TestDeterministic(unittest.TestCase):
     def setUp(self):
         torch.manual_seed(2)
-        self.model = nn.Sequential(
-            nn.Linear0(STATE_DIM, ACTION_DIM)
-        )
+        self.model = nn.Sequential(nn.Linear0(STATE_DIM, ACTION_DIM))
         self.optimizer = torch.optim.RMSprop(self.model.parameters(), lr=0.01)
         self.space = Box(np.array([-1, -1, -1]), np.array([1, 1, 1]))
         self.policy = DeterministicPolicy(
-            self.model,
-            self.optimizer,
-            self.space,
-            checkpointer=DummyCheckpointer()
+            self.model, self.optimizer, self.space, checkpointer=DummyCheckpointer()
         )
 
     def test_output_shape(self):
@@ -54,10 +51,7 @@ def test_converge(self):
 
     def test_target(self):
         self.policy = DeterministicPolicy(
-            self.model,
-            self.optimizer,
-            self.space,
-            target=FixedTarget(3)
+            self.model, self.optimizer, self.space, target=FixedTarget(3)
         )
         state = State(torch.ones(1, STATE_DIM))
 
@@ -81,5 +75,5 @@ def test_target(self):
         )
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/all/policies/gaussian_test.py b/all/policies/gaussian_test.py
index 3bcaf83f..e83be068 100644
--- a/all/policies/gaussian_test.py
+++ b/all/policies/gaussian_test.py
@@ -1,9 +1,11 @@
 import unittest
+
 import numpy as np
 import torch
-from torch import nn
 import torch_testing as tt
 from gymnasium.spaces import Box
+from torch import nn
+
 from all.approximation import DummyCheckpointer
 from all.core import State
 from all.policies import GaussianPolicy
@@ -16,11 +18,11 @@ class TestGaussian(unittest.TestCase):
     def setUp(self):
         torch.manual_seed(2)
         self.space = Box(np.array([-1, -1, -1]), np.array([1, 1, 1]))
-        self.model = nn.Sequential(
-            nn.Linear(STATE_DIM, ACTION_DIM * 2)
-        )
+        self.model = nn.Sequential(nn.Linear(STATE_DIM, ACTION_DIM * 2))
         optimizer = torch.optim.RMSprop(self.model.parameters(), lr=0.01)
-        self.policy = GaussianPolicy(self.model, optimizer, self.space, checkpointer=DummyCheckpointer())
+        self.policy = GaussianPolicy(
+            self.model, optimizer, self.space, checkpointer=DummyCheckpointer()
+        )
 
     def test_output_shape(self):
         state = State(torch.randn(1, STATE_DIM))
@@ -45,7 +47,7 @@ def test_reinforce_one(self):
 
     def test_converge(self):
         state = State(torch.randn(1, STATE_DIM))
-        target = torch.tensor([1., 2., -1.])
+        target = torch.tensor([1.0, 2.0, -1.0])
 
         for _ in range(0, 1000):
             dist = self.policy(state)
@@ -60,11 +62,13 @@ def test_converge(self):
     def test_eval(self):
         state = State(torch.randn(1, STATE_DIM))
         dist = self.policy.no_grad(state)
-        tt.assert_almost_equal(dist.mean, torch.tensor([[-0.237, 0.497, -0.058]]), decimal=3)
+        tt.assert_almost_equal(
+            dist.mean, torch.tensor([[-0.237, 0.497, -0.058]]), decimal=3
+        )
         tt.assert_almost_equal(dist.entropy(), torch.tensor([4.254]), decimal=3)
         best = self.policy.eval(state).sample()
         tt.assert_almost_equal(best, torch.tensor([[-0.888, -0.887, 0.404]]), decimal=3)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/all/policies/soft_deterministic_test.py b/all/policies/soft_deterministic_test.py
index f8c4a62a..5479a35e 100644
--- a/all/policies/soft_deterministic_test.py
+++ b/all/policies/soft_deterministic_test.py
@@ -1,8 +1,10 @@
 import unittest
-import torch
+
 import numpy as np
+import torch
 import torch_testing as tt
 from gymnasium.spaces import Box
+
 from all import nn
 from all.approximation import DummyCheckpointer
 from all.core import State
@@ -15,16 +17,11 @@
 class TestSoftDeterministic(unittest.TestCase):
     def setUp(self):
         torch.manual_seed(2)
-        self.model = nn.Sequential(
-            nn.Linear0(STATE_DIM, ACTION_DIM * 2)
-        )
+        self.model = nn.Sequential(nn.Linear0(STATE_DIM, ACTION_DIM * 2))
         self.optimizer = torch.optim.RMSprop(self.model.parameters(), lr=0.01)
         self.space = Box(np.array([-1, -1, -1]), np.array([1, 1, 1]))
         self.policy = SoftDeterministicPolicy(
-            self.model,
-            self.optimizer,
-            self.space,
-            checkpointer=DummyCheckpointer()
+            self.model, self.optimizer, self.space, checkpointer=DummyCheckpointer()
         )
 
     def test_output_shape(self):
@@ -61,7 +58,7 @@ def test_scaling(self):
         policy1 = SoftDeterministicPolicy(
             self.model,
             self.optimizer,
-            Box(np.array([-1., -1., -1.]), np.array([1., 1., 1.]))
+            Box(np.array([-1.0, -1.0, -1.0]), np.array([1.0, 1.0, 1.0])),
         )
         action1, log_prob1 = policy1(state)
 
@@ -71,7 +68,7 @@ def test_scaling(self):
         policy2 = SoftDeterministicPolicy(
             self.model,
             self.optimizer,
-            Box(np.array([-2., -1., -1.]), np.array([2., 1., 1.]))
+            Box(np.array([-2.0, -1.0, -1.0]), np.array([2.0, 1.0, 1.0])),
         )
         action2, log_prob2 = policy2(state)
 
@@ -80,5 +77,5 @@ def test_scaling(self):
         tt.assert_allclose(log_prob1 - np.log(2), log_prob2)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/all/policies/softmax_test.py b/all/policies/softmax_test.py
index 6584b44c..a543ac15 100644
--- a/all/policies/softmax_test.py
+++ b/all/policies/softmax_test.py
@@ -1,7 +1,9 @@
 import unittest
+
 import torch
-from torch import nn
 import torch_testing as tt
+from torch import nn
+
 from all.core import State
 from all.policies import SoftmaxPolicy
 
@@ -12,9 +14,7 @@
 class TestSoftmax(unittest.TestCase):
     def setUp(self):
         torch.manual_seed(2)
-        self.model = nn.Sequential(
-            nn.Linear(STATE_DIM, ACTIONS)
-        )
+        self.model = nn.Sequential(nn.Linear(STATE_DIM, ACTIONS))
         optimizer = torch.optim.SGD(self.model.parameters(), lr=0.1)
         self.policy = SoftmaxPolicy(self.model, optimizer)
 
@@ -63,25 +63,33 @@ def loss(log_probs):
 
         # notice the values increase with each successive reinforce
         log_probs = self.policy(states).log_prob(actions)
-        tt.assert_almost_equal(log_probs, torch.tensor([-0.84, -0.62, -0.757]), decimal=3)
+        tt.assert_almost_equal(
+            log_probs, torch.tensor([-0.84, -0.62, -0.757]), decimal=3
+        )
         self.policy.reinforce(loss(log_probs))
         log_probs = self.policy(states).log_prob(actions)
-        tt.assert_almost_equal(log_probs, torch.tensor([-0.811, -0.561, -0.701]), decimal=3)
+        tt.assert_almost_equal(
+            log_probs, torch.tensor([-0.811, -0.561, -0.701]), decimal=3
+        )
         self.policy.reinforce(loss(log_probs))
         log_probs = self.policy(states).log_prob(actions)
-        tt.assert_almost_equal(log_probs, torch.tensor([-0.785, -0.51, -0.651]), decimal=3)
+        tt.assert_almost_equal(
+            log_probs, torch.tensor([-0.785, -0.51, -0.651]), decimal=3
+        )
 
     def test_eval(self):
         states = State(torch.randn(3, STATE_DIM), torch.tensor([1, 1, 1]))
         dist = self.policy.no_grad(states)
-        tt.assert_almost_equal(dist.probs, torch.tensor([
-            [0.352, 0.216, 0.432],
-            [0.266, 0.196, 0.538],
-            [0.469, 0.227, 0.304]
-        ]), decimal=3)
+        tt.assert_almost_equal(
+            dist.probs,
+            torch.tensor(
+                [[0.352, 0.216, 0.432], [0.266, 0.196, 0.538], [0.469, 0.227, 0.304]]
+            ),
+            decimal=3,
+        )
         best = self.policy.eval(states).sample()
         tt.assert_equal(best, torch.tensor([2, 2, 0]))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/all/presets/atari_test.py b/all/presets/atari_test.py
index 5840d94a..65df9987 100644
--- a/all/presets/atari_test.py
+++ b/all/presets/atari_test.py
@@ -1,33 +1,26 @@
 import os
 import unittest
+
 import torch
+
 from all.environments import AtariEnvironment, DuplicateEnvironment
 from all.logging import DummyLogger
-from all.presets import Preset, ParallelPreset
-from all.presets.atari import (
-    a2c,
-    c51,
-    ddqn,
-    dqn,
-    ppo,
-    rainbow,
-    vac,
-    vpg,
-    vsarsa,
-    vqn
-)
+from all.presets import ParallelPreset
+from all.presets.atari import a2c, c51, ddqn, dqn, ppo, rainbow, vac, vpg, vqn, vsarsa
 
 
 class TestAtariPresets(unittest.TestCase):
     def setUp(self):
-        self.env = AtariEnvironment('Breakout')
+        self.env = AtariEnvironment("Breakout")
         self.env.reset()
-        self.parallel_env = DuplicateEnvironment([AtariEnvironment('Breakout'), AtariEnvironment('Breakout')])
+        self.parallel_env = DuplicateEnvironment(
+            [AtariEnvironment("Breakout"), AtariEnvironment("Breakout")]
+        )
         self.parallel_env.reset()
 
     def tearDown(self):
-        if os.path.exists('test_preset.pt'):
-            os.remove('test_preset.pt')
+        if os.path.exists("test_preset.pt"):
+            os.remove("test_preset.pt")
 
     def test_a2c(self):
         self.validate_preset(a2c)
@@ -60,7 +53,7 @@ def test_vqn(self):
         self.validate_preset(vqn)
 
     def validate_preset(self, builder):
-        preset = builder.device('cpu').env(self.env).build()
+        preset = builder.device("cpu").env(self.env).build()
         if isinstance(preset, ParallelPreset):
             return self.validate_parallel_preset(preset)
         return self.validate_standard_preset(preset)
@@ -73,8 +66,8 @@ def validate_standard_preset(self, preset):
         test_agent = preset.test_agent()
         test_agent.act(self.env.state)
         # test save/load
-        preset.save('test_preset.pt')
-        preset = torch.load('test_preset.pt')
+        preset.save("test_preset.pt")
+        preset = torch.load("test_preset.pt")
         test_agent = preset.test_agent()
         test_agent.act(self.env.state)
 
@@ -89,8 +82,8 @@ def validate_parallel_preset(self, preset):
         parallel_test_agent = preset.test_agent()
         parallel_test_agent.act(self.parallel_env.state_array)
         # test save/load
-        preset.save('test_preset.pt')
-        preset = torch.load('test_preset.pt')
+        preset.save("test_preset.pt")
+        preset = torch.load("test_preset.pt")
         test_agent = preset.test_agent()
         test_agent.act(self.env.state)
 
diff --git a/all/presets/builder_test.py b/all/presets/builder_test.py
index 7250436f..ba14f092 100644
--- a/all/presets/builder_test.py
+++ b/all/presets/builder_test.py
@@ -1,23 +1,24 @@
 import unittest
 from unittest.mock import Mock
+
 from all.presets import PresetBuilder
 
 
 class TestPresetBuilder(unittest.TestCase):
     def setUp(self):
         self.name = "my_preset"
-        self.default_hyperparameters = {
-            "lr": 1e-4,
-            "gamma": 0.99
-        }
+        self.default_hyperparameters = {"lr": 1e-4, "gamma": 0.99}
 
-        class MockPreset():
+        class MockPreset:
             def __init__(self, env, name, device, **hyperparameters):
                 self.env = env
                 self.name = name
                 self.device = device
                 self.hyperparameters = hyperparameters
-        self.builder = PresetBuilder(self.name, self.default_hyperparameters, MockPreset)
+
+        self.builder = PresetBuilder(
+            self.name, self.default_hyperparameters, MockPreset
+        )
 
     def test_default_name(self):
         agent = self.builder.env(Mock).build()
@@ -33,7 +34,9 @@ def test_default_hyperparameters(self):
 
     def test_override_hyperparameters(self):
         agent = self.builder.hyperparameters(lr=0.01).env(Mock).build()
-        self.assertEqual(agent.hyperparameters, {**self.default_hyperparameters, "lr": 0.01})
+        self.assertEqual(
+            agent.hyperparameters, {**self.default_hyperparameters, "lr": 0.01}
+        )
 
     def test_bad_hyperparameters(self):
         with self.assertRaises(KeyError):
@@ -41,26 +44,34 @@ def test_bad_hyperparameters(self):
 
     def test_default_device(self):
         agent = self.builder.env(Mock).build()
-        self.assertEqual(agent.device, 'cuda')
+        self.assertEqual(agent.device, "cuda")
 
     def test_override_device(self):
-        agent = self.builder.device('cpu').env(Mock).build()
-        self.assertEqual(agent.device, 'cpu')
+        agent = self.builder.device("cpu").env(Mock).build()
+        self.assertEqual(agent.device, "cpu")
 
     def test_no_side_effects(self):
-        self.builder.device('cpu').hyperparameters(lr=0.01).device('cpu').env(Mock).build()
+        self.builder.device("cpu").hyperparameters(lr=0.01).device("cpu").env(
+            Mock
+        ).build()
         my_env = Mock
         agent = self.builder.env(Mock).build()
         self.assertEqual(agent.name, self.name)
         self.assertEqual(agent.hyperparameters, self.default_hyperparameters)
-        self.assertEqual(agent.device, 'cuda')
+        self.assertEqual(agent.device, "cuda")
         self.assertEqual(agent.env, my_env)
 
     def test_call_api(self):
-        agent = self.builder(device='cpu', hyperparameters={"lr": 0.01}, name="cool_name").env(Mock).build()
+        agent = (
+            self.builder(device="cpu", hyperparameters={"lr": 0.01}, name="cool_name")
+            .env(Mock)
+            .build()
+        )
         self.assertEqual(agent.name, "cool_name")
-        self.assertEqual(agent.hyperparameters, {**self.default_hyperparameters, "lr": 0.01})
-        self.assertEqual(agent.device, 'cpu')
+        self.assertEqual(
+            agent.hyperparameters, {**self.default_hyperparameters, "lr": 0.01}
+        )
+        self.assertEqual(agent.device, "cpu")
 
 
 if __name__ == "__main__":
diff --git a/all/presets/classic_control_test.py b/all/presets/classic_control_test.py
index 39b93faa..2f06eff5 100644
--- a/all/presets/classic_control_test.py
+++ b/all/presets/classic_control_test.py
@@ -1,9 +1,11 @@
 import os
 import unittest
+
 import torch
-from all.environments import GymEnvironment, DuplicateEnvironment
+
+from all.environments import DuplicateEnvironment, GymEnvironment
 from all.logging import DummyLogger
-from all.presets import Preset, ParallelPreset
+from all.presets import ParallelPreset
 from all.presets.classic_control import (
     a2c,
     c51,
@@ -20,14 +22,16 @@
 
 class TestClassicControlPresets(unittest.TestCase):
     def setUp(self):
-        self.env = GymEnvironment('CartPole-v0')
+        self.env = GymEnvironment("CartPole-v0")
         self.env.reset()
-        self.parallel_env = DuplicateEnvironment([GymEnvironment('CartPole-v0'), GymEnvironment('CartPole-v0')])
+        self.parallel_env = DuplicateEnvironment(
+            [GymEnvironment("CartPole-v0"), GymEnvironment("CartPole-v0")]
+        )
         self.parallel_env.reset()
 
     def tearDown(self):
-        if os.path.exists('test_preset.pt'):
-            os.remove('test_preset.pt')
+        if os.path.exists("test_preset.pt"):
+            os.remove("test_preset.pt")
 
     def test_a2c(self):
         self.validate(a2c)
@@ -60,7 +64,7 @@ def test_vqn(self):
         self.validate(vqn)
 
     def validate(self, builder):
-        preset = builder.device('cpu').env(self.env).build()
+        preset = builder.device("cpu").env(self.env).build()
         if isinstance(preset, ParallelPreset):
             return self.validate_parallel_preset(preset)
         return self.validate_standard_preset(preset)
@@ -73,8 +77,8 @@ def validate_standard_preset(self, preset):
         test_agent = preset.test_agent()
         test_agent.act(self.env.state)
         # test save/load
-        preset.save('test_preset.pt')
-        preset = torch.load('test_preset.pt')
+        preset.save("test_preset.pt")
+        preset = torch.load("test_preset.pt")
         test_agent = preset.test_agent()
         test_agent.act(self.env.state)
 
@@ -89,8 +93,8 @@ def validate_parallel_preset(self, preset):
         parallel_test_agent = preset.test_agent()
         parallel_test_agent.act(self.parallel_env.state_array)
         # test save/load
-        preset.save('test_preset.pt')
-        preset = torch.load('test_preset.pt')
+        preset.save("test_preset.pt")
+        preset = torch.load("test_preset.pt")
         test_agent = preset.test_agent()
         test_agent.act(self.env.state)
 
diff --git a/all/presets/continuous_test.py b/all/presets/continuous_test.py
index 4d003ae8..5f3278a1 100644
--- a/all/presets/continuous_test.py
+++ b/all/presets/continuous_test.py
@@ -1,30 +1,29 @@
 import os
 import unittest
+
 import torch
-from all.core import State
-from all.environments import GymEnvironment, DuplicateEnvironment
+
+from all.environments import DuplicateEnvironment, GymEnvironment
 from all.logging import DummyLogger
-from all.presets import Preset, ParallelPreset
-from all.presets.continuous import (
-    ddpg,
-    ppo,
-    sac,
-)
+from all.presets import ParallelPreset
+from all.presets.continuous import ddpg, ppo, sac
 
 
 class TestContinuousPresets(unittest.TestCase):
     def setUp(self):
-        self.env = GymEnvironment('LunarLanderContinuous-v2')
+        self.env = GymEnvironment("LunarLanderContinuous-v2")
         self.env.reset()
-        self.parallel_env = DuplicateEnvironment([
-            GymEnvironment('LunarLanderContinuous-v2'),
-            GymEnvironment('LunarLanderContinuous-v2'),
-        ])
+        self.parallel_env = DuplicateEnvironment(
+            [
+                GymEnvironment("LunarLanderContinuous-v2"),
+                GymEnvironment("LunarLanderContinuous-v2"),
+            ]
+        )
         self.parallel_env.reset()
 
     def tearDown(self):
-        if os.path.exists('test_preset.pt'):
-            os.remove('test_preset.pt')
+        if os.path.exists("test_preset.pt"):
+            os.remove("test_preset.pt")
 
     def test_ddpg(self):
         self.validate(ddpg)
@@ -36,7 +35,7 @@ def test_sac(self):
         self.validate(sac)
 
     def validate(self, builder):
-        preset = builder.device('cpu').env(self.env).build()
+        preset = builder.device("cpu").env(self.env).build()
         if isinstance(preset, ParallelPreset):
             return self.validate_parallel_preset(preset)
         return self.validate_standard_preset(preset)
@@ -49,8 +48,8 @@ def validate_standard_preset(self, preset):
         test_agent = preset.test_agent()
         test_agent.act(self.env.state)
         # test save/load
-        preset.save('test_preset.pt')
-        preset = torch.load('test_preset.pt')
+        preset.save("test_preset.pt")
+        preset = torch.load("test_preset.pt")
         test_agent = preset.test_agent()
         test_agent.act(self.env.state)
 
@@ -65,8 +64,8 @@ def validate_parallel_preset(self, preset):
         parallel_test_agent = preset.test_agent()
         parallel_test_agent.act(self.parallel_env.state_array)
         # test save/load
-        preset.save('test_preset.pt')
-        preset = torch.load('test_preset.pt')
+        preset.save("test_preset.pt")
+        preset = torch.load("test_preset.pt")
         test_agent = preset.test_agent()
         test_agent.act(self.env.state)
 
diff --git a/all/presets/multiagent_atari_test.py b/all/presets/multiagent_atari_test.py
index 4f752da5..4173cbbc 100644
--- a/all/presets/multiagent_atari_test.py
+++ b/all/presets/multiagent_atari_test.py
@@ -1,27 +1,31 @@
 import os
 import unittest
+
 import torch
+
 from all.environments import MultiagentAtariEnv
 from all.logging import DummyLogger
-from all.presets.atari import dqn
 from all.presets import IndependentMultiagentPreset
+from all.presets.atari import dqn
 
 
 class TestMultiagentAtariPresets(unittest.TestCase):
     def setUp(self):
-        self.env = MultiagentAtariEnv('pong_v3', device='cpu')
+        self.env = MultiagentAtariEnv("pong_v3", device="cpu")
         self.env.reset()
 
     def tearDown(self):
-        if os.path.exists('test_preset.pt'):
-            os.remove('test_preset.pt')
+        if os.path.exists("test_preset.pt"):
+            os.remove("test_preset.pt")
 
     def test_independent(self):
         presets = {
-            agent_id: dqn.device('cpu').env(self.env.subenvs[agent_id]).build()
+            agent_id: dqn.device("cpu").env(self.env.subenvs[agent_id]).build()
             for agent_id in self.env.agents
         }
-        self.validate_preset(IndependentMultiagentPreset('independent', 'cpu', presets), self.env)
+        self.validate_preset(
+            IndependentMultiagentPreset("independent", "cpu", presets), self.env
+        )
 
     def validate_preset(self, preset, env):
         # normal agent
@@ -31,8 +35,8 @@ def validate_preset(self, preset, env):
         test_agent = preset.test_agent()
         test_agent.act(self.env.last())
         # test save/load
-        preset.save('test_preset.pt')
-        preset = torch.load('test_preset.pt')
+        preset.save("test_preset.pt")
+        preset = torch.load("test_preset.pt")
         test_agent = preset.test_agent()
         test_agent.act(self.env.last())
 
diff --git a/integration/atari_test.py b/integration/atari_test.py
index 7ebad9f4..8ffc9743 100644
--- a/integration/atari_test.py
+++ b/integration/atari_test.py
@@ -1,20 +1,10 @@
 import unittest
+
 import torch
-from all.environments import AtariEnvironment
-from all.presets.atari import (
-    a2c,
-    c51,
-    ddqn,
-    dqn,
-    ppo,
-    rainbow,
-    vac,
-    vpg,
-    vsarsa,
-    vqn
-)
 from validate_agent import validate_agent
 
+from all.environments import AtariEnvironment
+from all.presets.atari import a2c, c51, ddqn, dqn, ppo, rainbow, vac, vpg, vqn, vsarsa
 
 CPU = torch.device("cpu")
 if torch.cuda.is_available():
@@ -66,10 +56,16 @@ def test_dqn_cuda(self):
         )
 
     def test_ppo(self):
-        validate_agent(ppo.device(CPU).hyperparameters(n_envs=4), AtariEnvironment("Breakout", device=CPU))
+        validate_agent(
+            ppo.device(CPU).hyperparameters(n_envs=4),
+            AtariEnvironment("Breakout", device=CPU),
+        )
 
     def test_ppo_cuda(self):
-        validate_agent(ppo.device(CUDA).hyperparameters(n_envs=4), AtariEnvironment("Breakout", device=CUDA))
+        validate_agent(
+            ppo.device(CUDA).hyperparameters(n_envs=4),
+            AtariEnvironment("Breakout", device=CUDA),
+        )
 
     def test_rainbow(self):
         validate_agent(
@@ -84,35 +80,45 @@ def test_rainbow_cuda(self):
         )
 
     def test_vac(self):
-        validate_agent(vac.device(CPU).hyperparameters(n_envs=4), AtariEnvironment("Breakout", device=CPU))
+        validate_agent(
+            vac.device(CPU).hyperparameters(n_envs=4),
+            AtariEnvironment("Breakout", device=CPU),
+        )
 
     def test_vac_cuda(self):
         validate_agent(
-            vac.device(CUDA).hyperparameters(n_envs=4), AtariEnvironment("Breakout", device=CUDA)
+            vac.device(CUDA).hyperparameters(n_envs=4),
+            AtariEnvironment("Breakout", device=CUDA),
         )
 
     def test_vpg(self):
         validate_agent(vpg.device(CPU), AtariEnvironment("Breakout", device=CPU))
 
     def test_vpg_cuda(self):
-        validate_agent(
-            vpg.device(CUDA), AtariEnvironment("Breakout", device=CUDA)
-        )
+        validate_agent(vpg.device(CUDA), AtariEnvironment("Breakout", device=CUDA))
 
     def test_vsarsa(self):
-        validate_agent(vsarsa.device(CPU).hyperparameters(n_envs=4), AtariEnvironment("Breakout", device=CPU))
+        validate_agent(
+            vsarsa.device(CPU).hyperparameters(n_envs=4),
+            AtariEnvironment("Breakout", device=CPU),
+        )
 
     def test_vsarsa_cuda(self):
         validate_agent(
-            vsarsa.device(CUDA).hyperparameters(n_envs=4), AtariEnvironment("Breakout", device=CUDA)
+            vsarsa.device(CUDA).hyperparameters(n_envs=4),
+            AtariEnvironment("Breakout", device=CUDA),
         )
 
     def test_vqn(self):
-        validate_agent(vqn.device(CPU).hyperparameters(n_envs=4), AtariEnvironment("Breakout", device=CPU))
+        validate_agent(
+            vqn.device(CPU).hyperparameters(n_envs=4),
+            AtariEnvironment("Breakout", device=CPU),
+        )
 
     def test_vqn_cuda(self):
         validate_agent(
-            vqn.device(CUDA).hyperparameters(n_envs=4), AtariEnvironment("Breakout", device=CUDA)
+            vqn.device(CUDA).hyperparameters(n_envs=4),
+            AtariEnvironment("Breakout", device=CUDA),
         )
 
 
diff --git a/integration/classic_control_test.py b/integration/classic_control_test.py
index a1b49748..67bd72e8 100644
--- a/integration/classic_control_test.py
+++ b/integration/classic_control_test.py
@@ -1,4 +1,7 @@
 import unittest
+
+from validate_agent import validate_agent
+
 from all.environments import GymEnvironment
 from all.presets.classic_control import (
     a2c,
@@ -12,7 +15,6 @@
     vqn,
     vsarsa,
 )
-from validate_agent import validate_agent
 
 
 class TestClassicControlPresets(unittest.TestCase):
@@ -47,7 +49,7 @@ def test_vqn(self):
         self.validate(vqn)
 
     def validate(self, builder):
-        validate_agent(builder.device('cpu'), GymEnvironment("CartPole-v0"))
+        validate_agent(builder.device("cpu"), GymEnvironment("CartPole-v0"))
 
 
 if __name__ == "__main__":
diff --git a/integration/continuous_test.py b/integration/continuous_test.py
index f3e528a9..10414703 100644
--- a/integration/continuous_test.py
+++ b/integration/continuous_test.py
@@ -1,8 +1,10 @@
 import unittest
-from all.environments import GymEnvironment, PybulletEnvironment, MujocoEnvironment
-from all.presets.continuous import ddpg, ppo, sac
+
 from validate_agent import validate_agent
 
+from all.environments import GymEnvironment, MujocoEnvironment, PybulletEnvironment
+from all.presets.continuous import ddpg, ppo, sac
+
 
 class TestContinuousPresets(unittest.TestCase):
     def test_ddpg(self):
diff --git a/integration/multiagent_atari_test.py b/integration/multiagent_atari_test.py
index 9e0b54cc..f57e465d 100644
--- a/integration/multiagent_atari_test.py
+++ b/integration/multiagent_atari_test.py
@@ -1,10 +1,11 @@
 import unittest
+
 import torch
+from validate_agent import validate_multiagent
+
 from all.environments import MultiagentAtariEnv
 from all.presets import IndependentMultiagentPreset
 from all.presets.atari import dqn
-from validate_agent import validate_multiagent
-
 
 CPU = torch.device("cpu")
 if torch.cuda.is_available():
@@ -20,20 +21,24 @@
 
 class TestMultiagentAtariPresets(unittest.TestCase):
     def test_independent(self):
-        env = MultiagentAtariEnv('pong_v3', max_cycles=1000, device=CPU)
+        env = MultiagentAtariEnv("pong_v3", max_cycles=1000, device=CPU)
         presets = {
             agent_id: dqn.device(CPU).env(env.subenvs[agent_id]).build()
             for agent_id in env.agents
         }
-        validate_multiagent(IndependentMultiagentPreset('independent', CPU, presets), env)
+        validate_multiagent(
+            IndependentMultiagentPreset("independent", CPU, presets), env
+        )
 
     def test_independent_cuda(self):
-        env = MultiagentAtariEnv('pong_v3', max_cycles=1000, device=CUDA)
+        env = MultiagentAtariEnv("pong_v3", max_cycles=1000, device=CUDA)
         presets = {
             agent_id: dqn.device(CUDA).env(env.subenvs[agent_id]).build()
             for agent_id in env.agents
         }
-        validate_multiagent(IndependentMultiagentPreset('independent', CUDA, presets), env)
+        validate_multiagent(
+            IndependentMultiagentPreset("independent", CUDA, presets), env
+        )
 
 
 if __name__ == "__main__":