Repository: MathisFederico/LearnRL (language: python)

Issue:
Adding time in Logger
It would be nice to add ETA in the Logger.
It has to manage the different levels of verbose and use the metrics at hand to give the best an most usefull time estimate.
Tests for dt_step and dt_episodes
We need tests for dt_step and dt_episodes : 
- Use dt_step only on steps
- Use dt_episodes for episodes and episodes_cycles
- Assert dt_episode averaging over episodes_cycles


Discussion / hints:

Fix patch:
diff --git a/.coveragerc b/.coveragerc
new file mode 100644
index 0000000..284ad64
--- /dev/null
+++ b/.coveragerc
@@ -0,0 +1,6 @@
+[report]
+show_missing = True
+omit =
+    learnrl/callbacks/logger.py
+    learnrl/callbacks/tensorboard.py
+    learnrl/callbacks/wandb.py
\ No newline at end of file
diff --git a/VERSION b/VERSION
index 0d91a54..afaf360 100644
--- a/VERSION
+++ b/VERSION
@@ -1,1 +1,1 @@
-0.3.0
+1.0.0
\ No newline at end of file
diff --git a/learnrl/__init__.py b/learnrl/__init__.py
index b6bc176..3931c64 100644
--- a/learnrl/__init__.py
+++ b/learnrl/__init__.py
@@ -1,9 +1,9 @@
 # LearnRL a python library to learn and use reinforcement learning
 # Copyright (C) 2020 Mathïs FEDERICO <https://www.gnu.org/licenses/>
 
-"""LearnRL a python library to learn and use reinforcement learning"""
+"""LearnRL a python library to learn and log reinforcement learning"""
 
-__version__ = "0.3.0"
+__version__ = "1.0.0"
 
 from learnrl.agent import Agent
 from learnrl.envs import TurnEnv
diff --git a/learnrl/agent.py b/learnrl/agent.py
index 6203fef..4cf353c 100644
--- a/learnrl/agent.py
+++ b/learnrl/agent.py
@@ -22,6 +22,7 @@ def act(self, observation, greedy: bool=False) -> Union[int, float, np.ndarray]:
             greedy: If True, act greedely (without exploration).
 
         """
+        raise NotImplementedError('Agent.act must be user-defined by subclassing.')
 
     def learn(self) -> Dict[str, Union[int, float, np.ndarray]]:
         """How the :ref:`Agent` learns from his experiences.
diff --git a/learnrl/callbacks/__init__.py b/learnrl/callbacks/__init__.py
index 794de9c..f59df2f 100644
--- a/learnrl/callbacks/__init__.py
+++ b/learnrl/callbacks/__init__.py
@@ -3,7 +3,7 @@
 
 """Callbacks to perform actions while running a playground"""
 
-import importlib
+import importlib.util
 
 from learnrl.callbacks.callback import Callback, CallbackList
 from learnrl.callbacks.logging_callback import LoggingCallback, MetricList, Metric
diff --git a/learnrl/callbacks/callback.py b/learnrl/callbacks/callback.py
index 60788f1..c6a54f4 100644
--- a/learnrl/callbacks/callback.py
+++ b/learnrl/callbacks/callback.py
@@ -148,7 +148,7 @@ def set_playground(self, playground):
             callback.set_playground(playground)
 
     def _call_key_hook(self, key, hook, value=None, logs=None):
-        """ Helper func for {step|episode|cycle|run}_{begin|end} methods. """
+        """ Helper func for {step|episode|steps_cycle|episodes_cycle|run}_{begin|end} methods. """
 
         if len(self.callbacks) == 0:
             return
diff --git a/learnrl/callbacks/logger.py b/learnrl/callbacks/logger.py
index a27e5c5..0f2f838 100644
--- a/learnrl/callbacks/logger.py
+++ b/learnrl/callbacks/logger.py
@@ -3,7 +3,7 @@
 
 """Default Logger to show progress in console"""
 
-from typing import List
+from typing import List, Union
 import numpy as np
 
 from learnrl.callbacks import LoggingCallback, MetricList, Metric
@@ -34,9 +34,9 @@ class Logger(LoggingCallback):
     """
 
     def __init__(self,
-            metrics: List[str]=(('reward', {'steps': 'sum', 'episode': 'sum'})),
-            detailed_step_metrics: List[str]=('observation', 'action', 'next_observation'),
-            episode_only_metrics: List[str]=('dt_episode~'),
+            metrics: List[Union[str, tuple]]=None,
+            detailed_step_metrics: List[str]=None,
+            episode_only_metrics: List[str]=None,
             titles_on_top: bool=True):
 
         """Default logger in every :class:`~learnrl.playground.Playground` run.
@@ -49,6 +49,13 @@ def __init__(self,
 
         """
 
+        if metrics is None:
+            metrics = [('reward', {'steps_cycles': 'sum', 'episode': 'sum'})]
+        if detailed_step_metrics is None:
+            detailed_step_metrics = ['observation', 'action', 'next_observation']
+        if episode_only_metrics is None:
+            episode_only_metrics = ('dt_episode~')
+
         super().__init__(
             metrics=metrics,
             detailed_step_metrics=detailed_step_metrics,
@@ -288,7 +295,7 @@ def _get_time_text(elapsed_time, unit):
         if unit == 'episode':
             unit = 'eps'
         if elapsed_time < 1e-9:
-            return 'N/A        '
+            return '< 1ns      '
         if elapsed_time < 1e-6:
             time_display = f'{elapsed_time/1e-9:.01f}'
             time_unit = 'ns'
diff --git a/learnrl/callbacks/logging_callback.py b/learnrl/callbacks/logging_callback.py
index 16838a9..e9dd884 100644
--- a/learnrl/callbacks/logging_callback.py
+++ b/learnrl/callbacks/logging_callback.py
@@ -36,6 +36,8 @@ class MetricList():
     """ An helper object to represent a list of metrics """
 
     def __init__(self, metric_list: List[Metric]):
+        if metric_list is None:
+            metric_list = []
         self.metric_list = [self.to_metric(m) for m in metric_list]
         self.metric_names = [m.name for m in self.metric_list]
 
@@ -68,9 +70,9 @@ class LoggingCallback(Callback):
     """ Generic class for tracking metrics """
 
     def __init__(self,
-            metrics: List[str]=(('reward', {'steps': 'sum', 'episode': 'sum'})),
-            detailed_step_metrics: List[str]=('observation', 'action', 'next_observation'),
-            episode_only_metrics: List[str]=('dt_episode~')
+            metrics: List[Union[str, tuple]]=None,
+            detailed_step_metrics: List[str]=None,
+            episode_only_metrics: List[str]=None
         ):
         super().__init__()
         self.detailed_step_metrics = MetricList(detailed_step_metrics)
@@ -94,20 +96,25 @@ def _extract_lists(metrics):
         episode_metrics = []
         episodes_cycle_metrics = []
 
-        if not isinstance(metrics, (tuple, list)):
-            raise ValueError('Wrong metrics format. See metrics codes in documentation.')
-
-        for metric in metrics:
-            if isinstance(metric, (tuple, list)):
-                metric_name, ops = metric
-            elif isinstance(metric, str):
-                metric_name = metric
-                ops = {}
-
-            step_metrics.append(metric_name)
-            steps_cycle_metrics.append('.'.join((metric_name, ops.get('steps', 'avg'))))
-            episode_metrics.append('.'.join((metric_name, ops.get('episode', 'avg'))))
-            episodes_cycle_metrics.append('.'.join((metric_name, ops.get('episodes', 'avg'))))
+        if metrics is not None:
+            if not isinstance(metrics, (tuple, list)):
+                raise ValueError('Wrong metrics format. See metrics codes in documentation.')
+
+            for metric in metrics:
+                if isinstance(metric, (tuple, list)):
+                    metric_name, ops = metric
+                elif isinstance(metric, str):
+                    metric_name = metric
+                    if '.' in metric_name:
+                        metric_name, opertor_for_all = metric_name.split('.')
+                        ops = {key:opertor_for_all for key in ('steps', 'episode', 'episodes')}
+                    else:
+                        ops = {}
+
+                step_metrics.append(metric_name)
+                steps_cycle_metrics.append('.'.join((metric_name, ops.get('steps', 'avg'))))
+                episode_metrics.append('.'.join((metric_name, ops.get('episode', 'avg'))))
+                episodes_cycle_metrics.append('.'.join((metric_name, ops.get('episodes', 'avg'))))
         return step_metrics, steps_cycle_metrics, episode_metrics, episodes_cycle_metrics
 
     def _reset_attr(self, attr_name, operator):
@@ -225,32 +232,29 @@ def _update_metrics_all_agents(self,
             **kwargs
         ):
         """Update the logger attributes of all agents based on a metric list"""
-        if self.n_agents > 1:
-            for agent_id in range(self.n_agents):
-                self._update_metrics(metric_list, target_prefix, agent_id=agent_id, **kwargs)
-        else:
-            self._update_metrics(metric_list, target_prefix, agent_id=0, **kwargs)
+        for agent_id in range(self.n_agents):
+            self._update_metrics(metric_list, target_prefix, agent_id=agent_id, **kwargs)
 
-    def on_step_end(self, step, logs=None):
-        agent_id = logs.get('agent_id') if self.n_agents > 1 else None
-        self._update_metrics(self.episode_metrics, 'episode', logs=logs, agent_id=agent_id)
-        self._update_metrics_all_agents(self.steps_cycle_metrics, 'steps_cycle', logs=logs)
+    def on_run_begin(self, logs=None):
+        self.n_agents = len(self.playground.agents)
 
-    def on_steps_cycle_begin(self, step, logs=None):
-        self._update_metrics_all_agents(self.steps_cycle_metrics,
-            'steps_cycle', logs=logs, reset=True)
+    def on_episodes_cycle_begin(self, episode, logs=None):
+        self._update_metrics_all_agents(self.episodes_cycle_metrics,
+            'episodes_cycle', source_prefix='episode', logs=logs, reset=True)
 
     def on_episode_begin(self, episode, logs=None):
         self._update_metrics_all_agents(self.episode_metrics,
             'episode', logs=logs, reset=True)
 
+    def on_steps_cycle_begin(self, step, logs=None):
+        self._update_metrics_all_agents(self.steps_cycle_metrics,
+            'steps_cycle', logs=logs, reset=True)
+
+    def on_step_end(self, step, logs=None):
+        agent_id = logs.get('agent_id') if self.n_agents > 1 else None
+        self._update_metrics(self.episode_metrics, 'episode', logs=logs, agent_id=agent_id)
+        self._update_metrics_all_agents(self.steps_cycle_metrics, 'steps_cycle', logs=logs)
+
     def on_episode_end(self, episode, logs=None):
         self._update_metrics_all_agents(self.episodes_cycle_metrics,
             'episodes_cycle', source_prefix='episode', logs=logs)
-
-    def on_episodes_cycle_begin(self, episode, logs=None):
-        self._update_metrics_all_agents(self.episodes_cycle_metrics,
-            'episodes_cycle', source_prefix='episode', logs=logs, reset=True)
-
-    def on_run_begin(self, logs=None):
-        self.n_agents = len(self.playground.agents)
diff --git a/learnrl/callbacks/tensorboard.py b/learnrl/callbacks/tensorboard.py
index 50638f3..fb62bf8 100644
--- a/learnrl/callbacks/tensorboard.py
+++ b/learnrl/callbacks/tensorboard.py
@@ -8,7 +8,7 @@
 """
 
 import os, datetime
-from typing import List, Tuple, Dict, Any
+from typing import List, Tuple, Dict, Any, Union
 
 import tensorflow as tf
 
@@ -22,9 +22,9 @@ class TensorboardCallback(LoggingCallback):
     def __init__(self,
             log_dir: str,
             run_name: str=None,
-            metrics: List[Tuple[str, Dict[str, str]]]=(('reward', {'steps': 'sum', 'episode': 'sum'})),
-            detailed_step_metrics: List[str]=['observation', 'action', 'next_observation'],
-            episode_only_metrics: List[str]=['dt_episode~']
+            metrics: List[Union[str, tuple]]=None,
+            detailed_step_metrics: List[str]=None,
+            episode_only_metrics: List[str]=None
         ):
         
         """Tensorboard logger will log metrics in tensorboard.
diff --git a/learnrl/callbacks/wandb.py b/learnrl/callbacks/wandb.py
index 2377557..e371526 100644
--- a/learnrl/callbacks/wandb.py
+++ b/learnrl/callbacks/wandb.py
@@ -7,7 +7,7 @@
 
 """
 
-from typing import List
+from typing import List, Union
 
 import wandb
 from learnrl.callbacks.logging_callback import LoggingCallback, MetricList
@@ -74,9 +74,9 @@ class WandbCallback(LoggingCallback):
     """ WandbCallback will log metrics to wandb."""
 
     def __init__(self, run,
-            metrics: List[str]=(('reward', {'steps': 'sum', 'episode': 'sum'})),
-            detailed_step_metrics: List[str]=('observation', 'action', 'next_observation'),
-            episode_only_metrics: List[str]=('dt_episode~')
+            metrics: List[Union[str, tuple]]=None,
+            detailed_step_metrics: List[str]=None,
+            episode_only_metrics: List[str]=None
         ):
 
         """ WandbCallback will log metrics to wandb.
diff --git a/learnrl/envs.py b/learnrl/envs.py
index 22b98fa..63384a3 100644
--- a/learnrl/envs.py
+++ b/learnrl/envs.py
@@ -44,6 +44,7 @@ def step(self, action):
             info (dict): Additional informations given by the |gym.Env|.
 
         """
+        raise NotImplementedError('TurnEnv.step must be user-defined by subclassing.')
 
     @abstractmethod
     def turn(self, state) -> int:
@@ -60,6 +61,7 @@ def turn(self, state) -> int:
             agent_id (int): The next player id
 
         """
+        raise NotImplementedError('TurnEnv.turn must be user-defined by subclassing.')
 
     @abstractmethod
     def reset(self):
@@ -69,3 +71,4 @@ def reset(self):
             observation: The observation for the first :class:`Agent` to play
 
         """
+        raise NotImplementedError('TurnEnv.reset must be user-defined by subclassing.')
diff --git a/learnrl/playground.py b/learnrl/playground.py
index 1bd0007..1a22dbd 100644
--- a/learnrl/playground.py
+++ b/learnrl/playground.py
@@ -5,8 +5,9 @@
 
 import warnings
 
-from typing import List, Union, Callable
+from typing import List, Union, Callable, Any
 from abc import abstractmethod
+from numbers import Number
 
 import numpy as np
 from gym import Env
@@ -27,7 +28,8 @@ class DoneHandler():
     """
 
     @abstractmethod
-    def done(self, previous_observation, action, reward, done, info, observation) -> bool:
+    def done(self, previous_observation, action, reward,
+            done, info, observation, logs) -> bool:
         """Replace the environment done.
 
         Often used to make episodes shorter when the agent is stuck for example.
@@ -52,7 +54,8 @@ def reset(self):
 
     def _done(self, *args) -> bool:
         done = self.done(*args)
-        if not isinstance(done, (bool, np.bool, np.bool_)):
+        if not isinstance(done, bool) and \
+           not (isinstance(done, np.ndarray) and done.dtype == bool):
             raise ValueError(f"Done should be bool, got {done} of type {type(done)} instead")
         return done
 
@@ -73,7 +76,7 @@ class RewardHandler():
 
     @abstractmethod
     def reward(self, previous_observation, action, reward,
-            done, info, observation)-> float:
+            done, info, observation, logs)-> float:
         """Replace the environment reward.
 
         Often used to scale rewards or to do reward shaping.
@@ -98,9 +101,10 @@ def reset(self):
 
     def _reward(self, *args) -> float:
         reward = self.reward(*args)
-        if not isinstance(reward, (int, float)):
+        if not isinstance(reward, Number) and \
+           not (isinstance(reward, np.ndarray) and np.issubdtype(reward.dtype, np.floating)):
             raise ValueError(
-                f"Rewards should be scalars, got {reward} of type {type(reward)} instead"
+                f"Rewards should be a float, got {reward} of type {type(reward)} instead"
             )
         return float(reward)
 
@@ -150,9 +154,125 @@ def _get_episodes_cycle_len(episodes_cycle_len, episodes):
 
         return episodes_cycle_len
 
+    def _build_callbacks(self, callbacks, logger, params):
+        callbacks = callbacks if callbacks is not None else []
+        if logger is not None:
+            callbacks += [logger]
+
+        callbacks = CallbackList(callbacks)
+        callbacks.set_params(params)
+        callbacks.set_playground(self)
+        return callbacks
+
+    def _reset(self, reward_handler, done_handler):
+        previous = [{'observation':None, 'action':None, 'reward':None, 'done':None, 'info':None}
+            for _ in range(len(self.agents))]
+        if isinstance(reward_handler, RewardHandler):
+            reward_handler.reset()
+        if isinstance(done_handler, DoneHandler):
+            done_handler.reset()
+        return self.env.reset(), 0, False, previous
+
+    def _get_next_agent(self, observation):
+        turn_id = self.env.turn(observation) if isinstance(self.env, TurnEnv) else 0
+        try:
+            agent_id = self.agents_order[turn_id]
+            agent = self.agents[agent_id]
+        except IndexError as index_error:
+            error_msg = f'Not enough agents to play environement {self.env}'
+            raise ValueError(error_msg) from index_error
+        return agent, agent_id
+
+    @staticmethod
+    def _call_handlers(reward, done, experience,
+            reward_handler=None, done_handler=None, logs=None):
+        logs = {} if logs is None else logs
+
+        logs.update({'reward': reward})
+        if reward_handler is not None:
+            handled_reward = reward_handler(*experience, logs)
+            experience['reward'] = handled_reward
+            logs.update({'handled_reward': handled_reward})
+
+        logs.update({'done': done})
+        if done_handler is not None:
+            handled_done = done_handler(*experience, logs)
+            experience['done'] = handled_done
+            logs.update({'handled_done': handled_done})
+
+    def _run_step(self,
+            observation: Any,
+            previous: list,
+            logs: dict,
+            learn=False,
+            render=False,
+            render_mode='human',
+            reward_handler=None,
+            done_handler=None,
+        ):
+        """Run a single step"""
+        # Render the environment
+        if render:
+            self.env.render(render_mode)
+
+        agent, agent_id = self._get_next_agent(observation)
+
+        # If the agent has played before, perform a learning step
+        prev = previous[agent_id]
+        if learn and prev['observation'] is not None:
+            agent.remember(
+                prev['observation'], prev['action'], prev['reward'],
+                prev['done'], observation, prev['info']
+            )
+            agent_logs = agent.learn()
+            logs.update({f'agent_{agent_id}': agent_logs})
+
+        # Adds step informations to logs
+        logs.update({'agent_id':agent_id, 'observation':observation})
+
+        # Ask action to agent
+        action = agent.act(observation, greedy=not learn)
+        # Perform environment step
+        next_observation, reward, done, info = self.env.step(action)
+        logs.update(info)
+
+        # Compact experience
+        experience = {
+            'observation': observation,
+            'action': action,
+            'reward': reward,
+            'done': done,
+            'next_observation': next_observation,
+            'info': info
+        }
+
+        # Use Handlers
+        self._call_handlers(reward, done, experience, reward_handler, done_handler, logs)
+
+        # Store experience in prev
+        if learn:
+            for key, value in zip(prev, [observation, action, reward, done, info]):
+                prev[key] = value
+
+            # Perform a last learning step if done
+            if done:
+                agent.remember(**experience)
+                agent_logs = agent.learn()
+                logs.update({f'agent_{agent_id}': agent_logs})
+
+        # Do a last rendering if done
+        if done and render:
+            self.env.render(render_mode)
+
+        # Add experience to logs
+        logs.update(experience)
+
+        return next_observation, done
+
     def run(self,
             episodes: int,
             render: bool=True,
+            render_mode: str='human',
             learn: bool=True,
             steps_cycle_len: int=10,
             episodes_cycle_len: Union[int, float]=0.05,
@@ -170,7 +290,9 @@ def run(self,
 
         Args:
             episodes: Number of episodes to run.
-            render: If True, call :meth:`TurnEnv.render` every step.
+            render: If True, call |gym.render| every step.
+            render_mode: Rendering mode.
+                One of {'human', 'rgb_array', 'ansi'} (see |gym.render|).
             learn: If True, call :meth:`Agent.learn` every step.
             steps_cycle_len: Number of steps that compose a cycle.
             episode_cycle_len: Number of episodes that compose a cycle.
@@ -184,11 +306,8 @@ def run(self,
                 If None use the default :class:`~learnrl.callbacks.logger.Logger`.
 
         """
-
-        # Get episode cycle lenght
         episodes_cycle_len = self._get_episodes_cycle_len(episodes_cycle_len, episodes)
 
-        # Build callback list
         params = {
             'episodes': episodes,
             'episodes_cycle_len': episodes_cycle_len,
@@ -199,14 +318,12 @@ def run(self,
         }
 
         logger = logger if logger else Logger(**kwargs)
-        callbacks = callbacks if callbacks is not None else []
-        callbacks = CallbackList(callbacks + [logger])
-        callbacks.set_params(params)
-        callbacks.set_playground(self)
+        callbacks = self._build_callbacks(callbacks, logger, params)
 
         # Start the run
         logs = {}
         logs.update(params)
+
         callbacks.on_run_begin(logs)
 
         for episode in range(episodes):
@@ -214,24 +331,9 @@ def run(self,
             if episode % episodes_cycle_len == 0:
                 callbacks.on_episodes_cycle_begin(episode, logs)
 
-            # Reset environment and Handlers
-            observation = self.env.reset()
-            if isinstance(reward_handler, RewardHandler):
-                reward_handler.reset()
-            if isinstance(done_handler, DoneHandler):
-                done_handler.reset()
+            observation, step, done, previous = self._reset(reward_handler, done_handler)
 
-            # Store previous for correct attribution in multi_agent settings (TurnEnv)
-            previous = [
-                {'observation':None, 'action':None, 'reward':None, 'done':None, 'info':None}
-                for _ in range(len(self.agents))
-            ]
-
-            # Initialize episode variables
-            done = False
-            step = 0
-
-            logs.update({'episode':episode})
+            logs.update({'episode': episode})
             callbacks.on_episode_begin(episode, logs)
 
             while not done:
@@ -239,103 +341,33 @@ def run(self,
                 if step % steps_cycle_len == 0:
                     callbacks.on_steps_cycle_begin(step, logs)
 
+                logs.update({'step': step})
                 callbacks.on_step_begin(step, logs)
 
-                next_observation, logs = self._run_step(
-                    step=step,
+                observation, done = self._run_step(
                     observation=observation,
                     previous=previous,
+                    logs=logs,
                     learn=learn,
                     render=render,
+                    render_mode=render_mode,
                     reward_handler=reward_handler,
                     done_handler=done_handler,
-                    logs=logs
                 )
-
                 callbacks.on_step_end(step, logs)
 
-                step += 1
-                observation = next_observation
-
                 if (step + 1) % steps_cycle_len == 0 or done:
                     callbacks.on_steps_cycle_end(step, logs)
 
-                # Do a last rendering if done
-                if done and render:
-                    self.env.render()
+                step += 1
 
             callbacks.on_episode_end(episode, logs)
 
-            if (episode + 1) % episodes_cycle_len == 0 or episode == episodes - 1:
+            if (episode + 1) % episodes_cycle_len == 0 or episode + 1 == episodes:
                 callbacks.on_episodes_cycle_end(episode, logs)
 
         callbacks.on_run_end(logs)
 
-    def _run_step(self, step, observation, previous,
-        learn, render, reward_handler, done_handler, logs):
-        """Run a single step"""
-        # Render the environment
-        if render:
-            self.env.render()
-
-        # Get playing agent (TurnEnv)
-        turn_id = self.env.turn(observation) if isinstance(self.env, TurnEnv) else 0
-        agent_id = self.agents_order[turn_id]
-        try:
-            agent = self.agents[agent_id]
-        except IndexError as error:
-            error_msg = f'Not enough agents to play environement {self.env}'
-            raise ValueError(error_msg) from error
-
-        # If the agent has played before, perform a learning step
-        prev = previous[agent_id]
-        if learn and prev['observation'] is not None:
-            agent.remember(
-                prev['observation'], prev['action'], prev['reward'],
-                prev['done'], observation, prev['info']
-            )
-            agent_logs = agent.learn()
-            logs.update({f'agent_{agent_id}': agent_logs})
-
-        # Adds step informations to logs
-        logs.update({'step':step, 'agent_id':agent_id, 'observation':observation})
-
-        # Ask action to agent
-        action = agent.act(observation, greedy=not learn)
-        # Perform environment step
-        next_observation, reward, done, info = self.env.step(action)
-
-        # Perform reward handling
-        logs.update({'reward': reward})
-        if reward_handler is not None:
-            reward = reward_handler(
-                observation, action, reward, done, info, next_observation
-            )
-            logs.update({'handled_reward': reward})
-
-        # Perform done handling
-        logs.update({'done': done})
-        if done_handler is not None:
-            done = done_handler(
-                observation, action, reward, done, info, next_observation
-            )
-            logs.update({'handled_done': done})
-
-        # Perform a learning step
-        if learn:
-            for key, value in zip(prev, [observation, action, reward, done, info]):
-                prev[key] = value
-            if done:
-                agent.remember(
-                    observation, action, reward, done, next_observation, info
-                )
-                agent_logs = agent.learn()
-                logs.update({f'agent_{agent_id}': agent_logs})
-
-        logs.update({'action': action, 'next_observation': next_observation})
-        logs.update(info)
-        return next_observation, logs
-
 
     def fit(self, episodes, **kwargs):
         """Train the agent(s) on the environement for a number of episodes."""


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index fe6cdf3..4b2bc08 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -2,6 +2,7 @@
 
 # LearnRL a python library to learn and use reinforcement learning
 # Copyright (C) 2020 Mathïs FEDERICO <https://www.gnu.org/licenses/>
+# pylint: disable=unused-argument
 
 import importlib
 import pytest
diff --git a/tests/integrations/callbacks/test_logging.py b/tests/integrations/callbacks/test_logging.py
index ffcdb77..7c0080b 100644
--- a/tests/integrations/callbacks/test_logging.py
+++ b/tests/integrations/callbacks/test_logging.py
@@ -1,5 +1,6 @@
 # LearnRL a python library to learn and use reinforcement learning
 # Copyright (C) 2020 Mathïs FEDERICO <https://www.gnu.org/licenses/>
+# pylint: disable=too-many-branches, too-many-statements, missing-function-docstring
 
 """ Test metrics logging behavior """
 
@@ -108,9 +109,11 @@ def run(self, callbacks, eps_end_func=None, verbose=0):
 
                 if (step + 1) % steps_cycle_len == 0 or step == n_steps - 1:
                     for metric_name in ('reward', 'loss'):
+                        sum_value = steps_cycle_tracker[metric_name]
+                        avg_value = steps_cycle_tracker[metric_name] / steps_cycle_seen
                         logs.update({
-                            f'{metric_name}_steps_sum': steps_cycle_tracker[metric_name],
-                            f'{metric_name}_steps_avg': steps_cycle_tracker[metric_name] / steps_cycle_seen,
+                            f'{metric_name}_steps_sum': sum_value,
+                            f'{metric_name}_steps_avg': avg_value,
                         })
                     callbacks.on_steps_cycle_end(step, logs)
 
@@ -167,6 +170,7 @@ def run(self, callbacks, eps_end_func=None, verbose=0):
 def test_logging_steps_operators_(cycle_operator, metric_name):
 
     print(cycle_operator, metric_name)
+
     logging_callback = LoggingCallback(
         metrics=[('reward', {'steps':cycle_operator}), ('loss', {'steps':cycle_operator})],
     )
diff --git a/tests/unit/callbacks/test_callback.py b/tests/unit/callbacks/test_callback.py
new file mode 100644
index 0000000..a788d5b
--- /dev/null
+++ b/tests/unit/callbacks/test_callback.py
@@ -0,0 +1,141 @@
+# LearnRL a python library to learn and use reinforcement learning
+# Copyright (C) 2020 Mathïs FEDERICO <https://www.gnu.org/licenses/>
+# pylint: disable=protected-access, attribute-defined-outside-init, missing-function-docstring
+
+""" Test Callback abstract classes """
+
+import pytest
+import pytest_check as check
+from learnrl.callbacks.callback import Callback, CallbackList
+
+class TestCallback:
+
+    """ Callback """
+
+    def test_init(self):
+        """ should instanciate correctly. """
+        callback = Callback()
+        check.equal(callback.params, {})
+        check.is_none(callback.playground)
+
+    def test_set_params(self):
+        """ should update params correctly. """
+        callback = Callback()
+        expected_params = {"param": 1}
+        callback.set_params(expected_params)
+        check.equal(callback.params, expected_params)
+
+    def test_set_playground(self):
+        """ should update playground correctly. """
+        callback = Callback()
+        expected_playground = "playground"
+        callback.set_playground(expected_playground)
+        check.equal(callback.playground, expected_playground)
+
+    @pytest.mark.parametrize('hook', ['begin', 'end'])
+    @pytest.mark.parametrize('timescale', ['step', 'steps_cycle',
+        'episode', 'episodes_cycle', 'run'])
+    def test_callback_on_(self, timescale, hook):
+        callback = Callback()
+        callback_hook = getattr(callback, f'on_{timescale}_{hook}')
+        if timescale in ['step', 'steps_cycle']:
+            callback_hook(step=7, logs={'log': 'info'})
+        elif timescale in ['episode', 'episodes_cycle']:
+            callback_hook(episode=7, logs={'log': 'info'})
+        else:
+            callback_hook(logs={'log': 'info'})
+
+
+class TestCallbackList:
+
+    """ CallbackList """
+
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        self.callback1 = Callback()
+        self.callback2 = Callback()
+        self.callbacks = [self.callback1, self.callback2]
+        self.callback_path = 'learnrl.callbacks.callback.Callback'
+
+    def test_init(self):
+        """ should instanciate correctly. """
+        callbacks = CallbackList(self.callbacks)
+        check.equal(callbacks.params, {})
+        check.is_none(callbacks.playground)
+
+    def test_set_params(self, mocker):
+        """ should update params of all callbacks correctly. """
+        mocker.patch(self.callback_path + '.set_params')
+
+        callbacks = CallbackList(self.callbacks)
+        expected_params = {"param": 1}
+        callbacks.set_params(expected_params)
+
+        check.equal(callbacks.params, expected_params)
+        for callback in self.callbacks:
+            args, _ = callback.set_params.call_args
+            check.equal(args[0], expected_params)
+
+    def test_set_playground(self, mocker):
+        """ should update playground of all callbacks correctly. """
+        mocker.patch(self.callback_path + '.set_playground')
+
+        callbacks = CallbackList(self.callbacks)
+        expected_playground = "playground"
+        callbacks.set_playground(expected_playground)
+
+        check.equal(callbacks.playground, expected_playground)
+        for callback in self.callbacks:
+            args, _ = callback.set_playground.call_args
+            check.equal(args[0], expected_playground)
+
+    def test_call_key_hook_no_callback(self):
+        """ _call_key_hook should do nothing if no callbacks. """
+        callbacks = CallbackList([])
+        callbacks._call_key_hook('timescale', 'at')
+
+    @pytest.mark.parametrize('hook', ['begin', 'end'])
+    @pytest.mark.parametrize('timescale', ['step', 'steps_cycle',
+        'episode', 'episodes_cycle', 'run'])
+    def test_call_key_hook(self, mocker, timescale, hook):
+        time = 0.1 if hook == 'begin' else 0.123
+        value = 17 if timescale != 'run' else None
+        logs = {'log': 7}
+        t_begin_name = f't_{timescale}_begin'
+        dt_name = f'dt_{timescale}'
+        hook_name = f'on_{timescale}_{hook}'
+
+        mocker.patch('time.time', return_value=time)
+        mocker.patch('learnrl.callbacks.callback.Callback.' + hook_name)
+
+        callbacks = CallbackList(self.callbacks)
+        if hook == 'end':
+            setattr(callbacks, t_begin_name, 0.1)
+
+        callbacks._call_key_hook(timescale, hook, value=value, logs=logs)
+
+        if hook == 'end':
+            check.almost_equal(logs.get(dt_name), 0.023)
+
+        for callback in self.callbacks:
+            check.is_true(getattr(callback, hook_name).called)
+
+
+    @pytest.mark.parametrize('hook', ['begin', 'end'])
+    @pytest.mark.parametrize('timescale', ['step', 'steps_cycle',
+        'episode', 'episodes_cycle', 'run'])
+    def test_calls_all_callbacks_on_(self, timescale, hook, mocker):
+        mocker.patch('learnrl.callbacks.callback.CallbackList._call_key_hook')
+
+        callbacks = CallbackList(self.callbacks)
+        callback_hook = getattr(callbacks, f'on_{timescale}_{hook}')
+        if timescale in ['step', 'steps_cycle']:
+            callback_hook(step=7, logs={'log': 'info'})
+        elif timescale in ['episode', 'episodes_cycle']:
+            callback_hook(episode=7, logs={'log': 'info'})
+        else:
+            callback_hook(logs={'log': 'info'})
+
+        args, _ = callbacks._call_key_hook.call_args
+        check.equal(args[0], timescale)
+        check.equal(args[1], hook)
diff --git a/tests/unit/callbacks/test_logging.py b/tests/unit/callbacks/test_logging.py
index 9fd37c1..4bf1b01 100644
--- a/tests/unit/callbacks/test_logging.py
+++ b/tests/unit/callbacks/test_logging.py
@@ -1,6 +1,6 @@
 # LearnRL a python library to learn and use reinforcement learning
 # Copyright (C) 2020 Mathïs FEDERICO <https://www.gnu.org/licenses/>
-# pylint: disable=protected-access, attribute-defined-outside-init, unused-argument
+# pylint: disable=protected-access, attribute-defined-outside-init, unused-argument, missing-function-docstring
 
 """ Test metrics logging behavior """
 
@@ -167,6 +167,11 @@ def test_add_metriclist(self):
 class TestLoggingCallback:
     """ LoggingCallback """
 
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        """ Setup fixtures. """
+        self.logging_callback_path = "learnrl.callbacks.logging_callback.LoggingCallback"
+
     def test_init(self):
         """ should instanciate correctly. """
         LoggingCallback()
@@ -224,6 +229,91 @@ def test_update_attr_raise(self):
             logging_callback.reward = 'N/A'
             logging_callback._update_attr('reward', 1, 'x')
 
+    def test_on_run_begin(self, mocker):
+        """ should set n_agents on_run_begin. """
+        class DummyPlayground():
+            """DummyPlaygrounf"""
+            def __init__(self, n_agents):
+                self.agents = [None] * n_agents
+
+        n_agents = 5
+        mocker.patch(self.logging_callback_path + '', return_value=n_agents)
+        logging_callback = LoggingCallback()
+        logging_callback.playground = DummyPlayground(n_agents)
+        check.is_none(logging_callback.n_agents)
+        logging_callback.on_run_begin()
+        check.equal(logging_callback.n_agents, n_agents)
+
+    def test_on_episodes_cycle_begin(self, mocker):
+        """ should reset episodes_cycle_metrics on_episodes_cycle_begin. """
+        mocker.patch(self.logging_callback_path + '._update_metrics_all_agents')
+        logging_callback = LoggingCallback()
+        logging_callback.on_episodes_cycle_begin(episode=7)
+        args, kwargs = logging_callback._update_metrics_all_agents.call_args
+        check.equal(args[1], 'episodes_cycle')
+        check.is_true(kwargs.get('reset'))
+
+    def test_on_episode_begin(self, mocker):
+        """ should reset episode_metrics on_episode_begin. """
+        mocker.patch(self.logging_callback_path + '._update_metrics_all_agents')
+        logging_callback = LoggingCallback()
+        logging_callback.on_episode_begin(episode=7)
+        args, kwargs = logging_callback._update_metrics_all_agents.call_args
+        check.equal(args[1], 'episode')
+        check.is_true(kwargs.get('reset'))
+
+    def test_on_step_cycle_begin(self, mocker):
+        """ should reset steps_cycle_metrics on step_cycle_begin. """
+        mocker.patch(self.logging_callback_path + '._update_metrics_all_agents')
+        logging_callback = LoggingCallback()
+        logging_callback.on_steps_cycle_begin(step=7)
+        args, kwargs = logging_callback._update_metrics_all_agents.call_args
+        check.equal(args[1], 'steps_cycle')
+        check.is_true(kwargs.get('reset'))
+
+    def test_on_episode_end(self, mocker):
+        """ should update episodes_cycle_metrics on_episode_end from episode metrics. """
+        mocker.patch(self.logging_callback_path + '._update_metrics_all_agents')
+        logging_callback = LoggingCallback()
+        logging_callback.on_episode_end(episode=7)
+        args, kwargs = logging_callback._update_metrics_all_agents.call_args
+        check.equal(args[1], 'episodes_cycle')
+        check.is_false(kwargs.get('reset'))
+        check.equal(kwargs.get('source_prefix'), 'episode')
+
+    def test_on_step_end(self, mocker):
+        """ should update episode_metrics for the current agent and
+            steps_cycle_metrics for all agents on_step_end. """
+        mocker.patch(self.logging_callback_path + '._update_metrics_all_agents')
+        mocker.patch(self.logging_callback_path + '._update_metrics')
+
+        logging_callback = LoggingCallback()
+        logging_callback.n_agents = 1
+        logging_callback.on_step_end(step=7)
+
+        args, kwargs = logging_callback._update_metrics.call_args
+        check.equal(args[1], 'episode')
+        check.is_false(kwargs.get('reset'))
+        check.is_none(kwargs.get('agent_id'))
+
+        args, kwargs = logging_callback._update_metrics_all_agents.call_args
+        check.equal(args[1], 'steps_cycle')
+        check.is_false(kwargs.get('reset'))
+
+        # Multi agents
+        logging_callback = LoggingCallback()
+        logging_callback.n_agents = 5
+        logging_callback.on_step_end(step=7, logs={'agent_id': 3})
+
+        args, kwargs = logging_callback._update_metrics.call_args
+        check.equal(args[1], 'episode')
+        check.is_false(kwargs.get('reset'))
+        check.equal(kwargs.get('agent_id'), 3)
+
+        args, kwargs = logging_callback._update_metrics_all_agents.call_args
+        check.equal(args[1], 'steps_cycle')
+        check.is_false(kwargs.get('reset'))
+
 
 class TestGetAttrName:
     """ LoggingCallback._get_attr_name """
@@ -254,7 +344,6 @@ class TestExtractFromLogs:
     @pytest.fixture(autouse=True)
     def setup(self):
         """ Setup fake logs and retrieve function to test. """
-
         self.extract_from_logs = LoggingCallback._extract_metric_from_logs
         self.logs = {
             'value': 0,
@@ -270,26 +359,21 @@ def setup(self):
 
     def test_find_value(self):
         """ should find value in logs. """
-
         value = self.extract_from_logs('value', self.logs)
         check.equal(value, 0)
 
     def test_find_nothing(self):
         """ should return N/A when there is no value. """
-
         nothing = self.extract_from_logs('nothing', self.logs)
         check.equal(nothing, 'N/A')
 
-
     def test_find_any_specific_value(self):
         """ should find any specific value in agent when no agent is specified. """
-
         specific_value = self.extract_from_logs('specific_value', self.logs)
         check.equal(specific_value, 42)
 
     def test_find_agent_specific_value(self):
         """ should find specific agent values. """
-
         value_0 = self.extract_from_logs('value', self.logs, agent_id=0)
         value_1 = self.extract_from_logs('value', self.logs, agent_id=1)
 
@@ -298,7 +382,6 @@ def test_find_agent_specific_value(self):
 
     def test_return_outer_value_when_no_specific_agent_value(self):
         """ should return outer value if no specific value is found when an agent is specified. """
-
         step = self.extract_from_logs('step', self.logs, agent_id=0)
         check.equal(step, 2)
 
@@ -309,7 +392,6 @@ class TestLoggingCallbackExtractLists:
     @pytest.fixture(autouse=True)
     def setup_logs(self):
         """ Setup fake logs and retrieve function to test. """
-
         self.extract_lists = LoggingCallback._extract_lists
 
     def test_extract_lists_use_case(self):
@@ -318,16 +400,17 @@ def test_extract_lists_use_case(self):
         metrics = [
             ('reward~rwd', {'steps': 'sum', 'episode': 'sum'}),
             ('loss', {'episodes': 'last'}),
-            'exploration~exp'
+            'exploration~exp.last',
+            'decay'
         ]
 
         metric_lists = self.extract_lists(metrics)
 
         expected_metric_lists = [
-            ['reward~rwd', 'loss', 'exploration~exp'],
-            ['reward~rwd.sum', 'loss.avg', 'exploration~exp.avg'],
-            ['reward~rwd.sum', 'loss.avg', 'exploration~exp.avg'],
-            ['reward~rwd.avg', 'loss.last', 'exploration~exp.avg']
+            ['reward~rwd', 'loss', 'exploration~exp', 'decay'],
+            ['reward~rwd.sum', 'loss.avg', 'exploration~exp.last', 'decay.avg'],
+            ['reward~rwd.sum', 'loss.avg', 'exploration~exp.last', 'decay.avg'],
+            ['reward~rwd.avg', 'loss.last', 'exploration~exp.last', 'decay.avg']
         ]
 
         metric_lists_names = [
@@ -346,7 +429,6 @@ def test_extract_lists_use_case(self):
 
     def test_extract_lists_wrong_format(self):
         """ should raise ValueError on wrong metric nested format. """
-
         metrics = ('reward')
 
         with pytest.raises(ValueError, match=r".*metrics format.*"):
@@ -403,3 +485,125 @@ def test_get_value_no_prefix_no_logs(self, mocker):
         metric = Metric('attribute')
         value = logging_callback._get_value(metric)
         check.equal(value, expected_value)
+
+
+class TestLoggingCallbackUpdateMetrics:
+    """ LoggingCallback._update_metrics """
+
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        self.metric_list = MetricList([
+            Metric("metric_1"),
+            Metric("metric_2"),
+            Metric("metric_2")
+        ])
+
+        self.logging_callback = LoggingCallback()
+
+    def test_na_value(self, mocker):
+        """ should not update attr if value is N/A. """
+        logging_callback_path = 'learnrl.callbacks.logging_callback.LoggingCallback'
+        mocker.patch(logging_callback_path + '._get_attr_name', return_value="target_name")
+        mocker.patch(logging_callback_path + '._reset_attr')
+        mocker.patch(logging_callback_path + '._get_value', return_value='N/A')
+        mocker.patch(logging_callback_path + '._update_attr')
+
+        self.logging_callback._update_metrics(
+            self.metric_list,
+            'target_prefix',
+            source_prefix='source_prefix',
+            logs='logs',
+            agent_id='agent_id',
+            reset=False
+        )
+
+        for args, _ in self.logging_callback._get_attr_name.call_args_list:
+            check.equal(args[0], 'target_prefix')
+            check.equal(args[2], 'agent_id')
+
+        for args, _ in self.logging_callback._get_value.call_args_list:
+            check.equal(args[1], 'source_prefix')
+            check.equal(args[2], 'agent_id')
+            check.equal(args[3], 'logs')
+
+        check.is_false(self.logging_callback._reset_attr.called)
+        check.is_false(self.logging_callback._update_attr.called)
+
+    def test_update(self, mocker):
+        """ should update attr if value is not N/A. """
+        logging_callback_path = 'learnrl.callbacks.logging_callback.LoggingCallback'
+        mocker.patch(logging_callback_path + '._get_attr_name', return_value="target_name")
+        mocker.patch(logging_callback_path + '._reset_attr')
+        mocker.patch(logging_callback_path + '._get_value', return_value="value")
+        mocker.patch(logging_callback_path + '._update_attr')
+
+        self.logging_callback._update_metrics(
+            self.metric_list,
+            'target_prefix',
+            source_prefix='source_prefix',
+            logs='logs',
+            agent_id='agent_id',
+            reset=False
+        )
+
+        for args, _ in self.logging_callback._get_attr_name.call_args_list:
+            check.equal(args[0], 'target_prefix')
+            check.equal(args[2], 'agent_id')
+
+        for args, _ in self.logging_callback._get_value.call_args_list:
+            check.equal(args[1], 'source_prefix')
+            check.equal(args[2], 'agent_id')
+            check.equal(args[3], 'logs')
+
+        check.is_false(self.logging_callback._reset_attr.called)
+
+        for args, _ in self.logging_callback._update_attr.call_args_list:
+            check.equal(args[0], 'target_name')
+            check.equal(args[1], 'value')
+
+    def test_reset(self, mocker):
+        """ should reset attr if reset is True. """
+        logging_callback_path = 'learnrl.callbacks.logging_callback.LoggingCallback'
+        mocker.patch(logging_callback_path + '._get_attr_name', return_value="target_name")
+        mocker.patch(logging_callback_path + '._reset_attr')
+        mocker.patch(logging_callback_path + '._get_value', return_value="value")
+        mocker.patch(logging_callback_path + '._update_attr')
+
+        self.logging_callback._update_metrics(
+            self.metric_list,
+            'target_prefix',
+            source_prefix='source_prefix',
+            logs='logs',
+            agent_id='agent_id',
+            reset=True
+        )
+
+        for args, _ in self.logging_callback._get_attr_name.call_args_list:
+            check.equal(args[0], 'target_prefix')
+            check.equal(args[2], 'agent_id')
+
+        for args, _ in self.logging_callback._get_value.call_args_list:
+            check.equal(args[1], 'source_prefix')
+            check.equal(args[2], 'agent_id')
+            check.equal(args[3], 'logs')
+
+        for args, _ in self.logging_callback._reset_attr.call_args_list:
+            check.equal(args[0], 'target_name')
+
+        check.is_false(self.logging_callback._update_attr.called)
+
+    def test_all_agents(self, mocker):
+        """ should update all agents indepentently. """
+        logging_callback_path = 'learnrl.callbacks.logging_callback.LoggingCallback'
+        mocker.patch(logging_callback_path + '._update_metrics')
+
+        n_agents = 5
+        self.logging_callback.n_agents = n_agents
+        self.logging_callback._update_metrics_all_agents(
+            self.metric_list,
+            'target_prefix'
+        )
+
+        for i in range(n_agents):
+            _, kwargs = self.logging_callback._update_metrics.call_args_list[i]
+            check.equal(kwargs.get('agent_id'), i)
diff --git a/tests/unit/callbacks/test_tensorboard.py b/tests/unit/callbacks/test_tensorboard.py
index 503ead9..a17de68 100644
--- a/tests/unit/callbacks/test_tensorboard.py
+++ b/tests/unit/callbacks/test_tensorboard.py
@@ -1,9 +1,10 @@
-""" Test Tensorboard callback """
-
 # LearnRL a python library to learn and use reinforcement learning
 # Copyright (C) 2020 Mathïs FEDERICO <https://www.gnu.org/licenses/>
+# pylint: disable=unused-argument
+
+""" TensorboardCallback """
 
-import importlib
+import importlib.util
 import pytest
 
 import learnrl
diff --git a/tests/unit/callbacks/test_wandb.py b/tests/unit/callbacks/test_wandb.py
index 39117e9..1d14ffd 100644
--- a/tests/unit/callbacks/test_wandb.py
+++ b/tests/unit/callbacks/test_wandb.py
@@ -1,9 +1,10 @@
-""" Test wandb callback """
-
 # LearnRL a python library to learn and use reinforcement learning
 # Copyright (C) 2020 Mathïs FEDERICO <https://www.gnu.org/licenses/>
+# pylint: disable=unused-argument
+
+""" Test WandbCallback """
 
-import importlib
+import importlib.util
 import pytest
 
 import learnrl
diff --git a/tests/unit/test_agent.py b/tests/unit/test_agent.py
new file mode 100644
index 0000000..e9b1d12
--- /dev/null
+++ b/tests/unit/test_agent.py
@@ -0,0 +1,38 @@
+# LearnRL a python library to learn and use reinforcement learning
+# Copyright (C) 2020 Mathïs FEDERICO <https://www.gnu.org/licenses/>
+# pylint: disable=protected-access, attribute-defined-outside-init, unused-argument
+
+""" Test agent.py """
+
+import pytest
+import pytest_check as check
+
+from learnrl.agent import Agent
+
+class TestAgent:
+
+    """Agent"""
+
+    def test_act(self):
+        """ act should raise NotImplementedError. """
+        agent = Agent()
+        with pytest.raises(NotImplementedError):
+            agent.act('obs')
+
+    def test_learn(self):
+        """ learn should return an empty dictionary by default. """
+        agent = Agent()
+        logs = agent.learn()
+        check.equal(logs, {})
+
+    def test_remember(self):
+        """ remember should pass by default. """
+        agent = Agent()
+        agent.remember(
+            'observation',
+            'action',
+            'reward',
+            'done',
+            next_observation='next_observation',
+            info={'info': 'info'},
+        )
diff --git a/tests/unit/test_envs.py b/tests/unit/test_envs.py
new file mode 100644
index 0000000..a00758c
--- /dev/null
+++ b/tests/unit/test_envs.py
@@ -0,0 +1,32 @@
+# LearnRL a python library to learn and use reinforcement learning
+# Copyright (C) 2020 Mathïs FEDERICO <https://www.gnu.org/licenses/>
+# pylint: disable=protected-access, attribute-defined-outside-init, unused-argument
+
+""" Test agent.py """
+
+import pytest
+
+from learnrl.envs import TurnEnv
+
+
+class TestTurnEnv:
+
+    """ TurnEnv """
+
+    def test_step(self):
+        """ step should raise NotImplementedError. """
+        env = TurnEnv()
+        with pytest.raises(NotImplementedError):
+            env.step('action')
+
+    def test_turn(self):
+        """ turn should raise NotImplementedError. """
+        env = TurnEnv()
+        with pytest.raises(NotImplementedError):
+            env.turn('observation')
+
+    def test_reset(self):
+        """ reset should raise NotImplementedError. """
+        env = TurnEnv()
+        with pytest.raises(NotImplementedError):
+            env.reset()
diff --git a/tests/unit/test_playground.py b/tests/unit/test_playground.py
index dfc8673..0283b80 100644
--- a/tests/unit/test_playground.py
+++ b/tests/unit/test_playground.py
@@ -1,46 +1,687 @@
+# LearnRL a python library to learn and use reinforcement learning
+# Copyright (C) 2020 Mathïs FEDERICO <https://www.gnu.org/licenses/>
+# pylint: disable=protected-access, attribute-defined-outside-init, unused-argument, too-many-instance-attributes
+
+""" Test playground.py """
+
 import pytest
+import pytest_check as check
+
+import numpy as np
 
 from gym import Env
-from learnrl import Playground, Agent
+from learnrl.agent import Agent
+from learnrl.envs import TurnEnv
+from learnrl.playground import Playground, RewardHandler, DoneHandler
+from learnrl.callbacks.callback import CallbackList, Callback
+
+class TestPlayground:
+
+    """ Playground """
+
+    @pytest.fixture(autouse=True)
+    def setup_playground(self):
+        """Setup of used fixtures"""
+        self.env = Env()
+        self.n_agents = 5
+        self.agents = [Agent() for _ in range(self.n_agents)]
+
+    def test_single_agent_argument(self):
+        " should transform a single agent in a list containing itself. "
+        single_agent = Agent()
+        playground = Playground(self.env, single_agent)
+        check.equal(playground.agents, [single_agent])
+
+    def test_env_typeerror(self):
+        " should raise a TypeError if the environment isn't a subclass of gym. "
+        with pytest.raises(TypeError, match=r"environement.*gym.Env"):
+            Playground("env", self.agents)
+
+    def test_agent_typeerror(self):
+        " should raise a TypeError if any agent isn't a subclass of learnrl.Agent. "
+        with pytest.raises(TypeError, match=r"agent.*learnrl.Agent"):
+            Playground(self.env, [Agent(), 'agent'])
+
+    def test_run(self, mocker):
+        """ run should call callbacks at the right time and in the right order. """
+
+        steps_outputs = [
+            (f'obs_{i+1}', False)
+            for i in range(9)
+        ]
+        steps_outputs += [(f'obs_{9}', True)]
+        steps_outputs = steps_outputs*10
+        steps_outputs = steps_outputs[::-1]
+
+        def dummy_run_step(*args, **kwargs):
+            return steps_outputs.pop()
+
+        class RegisterCallback(Callback):
+
+            """Dummy Callback to register calls"""
+
+            def __init__(self):
+                super().__init__()
+                self.stored_key = ""
+
+            def on_run_begin(self, logs=None):
+                self.stored_key += "|-"
+
+            def on_episodes_cycle_begin(self, episode, logs=None):
+                self.stored_key += "["
+
+            def on_episode_begin(self, episode, logs=None):
+                self.stored_key += "("
+
+            def on_steps_cycle_begin(self, step, logs=None):
+                self.stored_key += "<"
+
+            def on_step_begin(self, step, logs=None):
+                self.stored_key += ","
+
+            def on_step_end(self, step, logs=None):
+                self.stored_key += "."
+
+            def on_steps_cycle_end(self, step, logs=None):
+                self.stored_key += ">"
+
+            def on_episode_end(self, episode, logs=None):
+                self.stored_key += ")"
+
+            def on_episodes_cycle_end(self, episode, logs=None):
+                self.stored_key += "]"
+
+            def on_run_end(self, logs=None):
+                self.stored_key += "-|"
+
+        mocker.patch(
+            'learnrl.playground.Playground._get_episodes_cycle_len',
+            return_value=3
+        )
+        mocker.patch(
+            'learnrl.playground.Playground._reset',
+            lambda *args: ('obs_0', 0, False, {})
+        )
+        mocker.patch(
+            'learnrl.playground.Playground._build_callbacks',
+            lambda self, callbacks, logger, params: callbacks[0]
+        )
+        mocker.patch(
+            'learnrl.playground.Playground._run_step',
+            dummy_run_step
+        )
+
+        playground = Playground(self.env, self.agents)
+        register_callback = RegisterCallback()
+        playground.run(episodes=10, callbacks=[register_callback], steps_cycle_len=3)
+        episode_key = "(<,.,.,.><,.,.,.><,.,.,.><,.>)"
+        expected_key = "|-[" + episode_key*3 + "][" + episode_key*3 + "][" + \
+            episode_key*3  + "][" + episode_key + "]-|"
+        check.equal(register_callback.stored_key, expected_key)
+
+    def test_fit(self, mocker):
+        """ fit should call run with learn=True and render=False. """
+        mocker.patch('learnrl.playground.Playground.run')
+        playground = Playground(self.env, self.agents)
+        playground.fit(10)
+        _, kwargs = playground.run.call_args
+        check.is_true(kwargs.get('learn'))
+        check.is_false(kwargs.get('render'))
+
+    def test_fit_warn_learn(self, mocker):
+        """ fit should warn a UserWarning if learn=False. """
+        mocker.patch('learnrl.playground.Playground.run')
+        playground = Playground(self.env, self.agents)
+        with pytest.warns(UserWarning, match=r".*agents will not improve.*"):
+            playground.fit(10, learn=False)
+
+    def test_fit_warn_render(self, mocker):
+        """ fit should warn a RuntimeWarning if render=True. """
+        mocker.patch('learnrl.playground.Playground.run')
+        playground = Playground(self.env, self.agents)
+        with pytest.warns(RuntimeWarning, match=r".*computation speed.*"):
+            playground.fit(10, render=True)
+
+    def test_test(self, mocker):
+        """ test should call run with learn=False and render=True. """
+        mocker.patch('learnrl.playground.Playground.run')
+        playground = Playground(self.env, self.agents)
+        playground.test(10)
+        _, kwargs = playground.run.call_args
+        check.is_false(kwargs.get('learn'))
+        check.is_true(kwargs.get('render'))
+
+    def test_test_warn_learn(self, mocker):
+        """ test should warn a UserWarning if learn=True. """
+        mocker.patch('learnrl.playground.Playground.run')
+        playground = Playground(self.env, self.agents)
+        with pytest.warns(UserWarning, match=r".*not act greedy.*"):
+            playground.test(10, learn=True)
+
+    def test_test_warn_render(self, mocker):
+        """ test should warn a UserWarning if render=False. """
+        mocker.patch('learnrl.playground.Playground.run')
+        playground = Playground(self.env, self.agents)
+        with pytest.warns(UserWarning, match=r".*render=True.*"):
+            playground.test(10, render=False, verbose=0)
+
+class TestPlaygroundGetEpisodeCycleLen:
+
+    """ Playground._get_episodes_cycle_len """
+
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        """Setup of used fixtures"""
+        self.get_episodes_cycle_len = Playground._get_episodes_cycle_len
+
+    def test_over_one(self):
+        " should be the floor if > 1. "
+        value = self.get_episodes_cycle_len(10, 100)
+        check.equal(value, 10)
+
+        value = self.get_episodes_cycle_len(12.4, 100)
+        check.equal(value, 12)
+
+    def test_between_zero_one(self):
+        " should be a proportion of total episodes if between 0 and 1. "
+        value = self.get_episodes_cycle_len(0.42, 100)
+        check.equal(value, 42)
+
+    def test_bellow_zero(self):
+        " should raise a ValueError if < 0. "
+        with pytest.raises(ValueError, match=r".*> 0.*"):
+            self.get_episodes_cycle_len(0, 100)
+
+        with pytest.raises(ValueError, match=r".*> 0.*"):
+            self.get_episodes_cycle_len(-0.2, 100)
+
+
+class TestPlaygroundBuildCallbacks:
+    """ Playground._build_callbacks """
+
+    @pytest.fixture(autouse=True)
+    def setup_playground(self):
+        """Setup of used fixtures"""
+        self.env = Env()
+        self.n_agents = 5
+        self.agents = [Agent() for _ in range(self.n_agents)]
+        self.playground = Playground(self.env, self.agents)
+
+    def test_build_callback(self, mocker):
+        """ should build a CallbackList from given callbacks and logger
+        and set their params and playground. """
+        callback_path = 'learnrl.callbacks.callback.Callback'
+        mocker.patch(callback_path + '.set_params')
+        mocker.patch(callback_path + '.set_playground')
+
+        callbacks = [Callback(), Callback()]
+        logger = Callback()
+        params = {'param123': 123}
+        callbacklist = self.playground._build_callbacks(callbacks, logger, params)
+        check.is_instance(callbacklist, CallbackList)
+        for callback in callbacklist.callbacks:
+            args, _ = callback.set_params.call_args
+            check.equal(args[0], params)
+            args, _ = callback.set_playground.call_args
+            check.equal(args[0], self.playground)
+        check.is_in(logger, callbacklist.callbacks)
+
+    def test_builb_callback_no_logger(self):
+        """ should work if logger is None. """
+        callbacks = [Callback(), Callback()]
+        logger = None
+        params = {'param123': 123}
+        callbacklist = self.playground._build_callbacks(callbacks, logger, params)
+        check.is_instance(callbacklist, CallbackList)
+        for callback in callbacklist.callbacks:
+            check.equal(callback.params, params)
+            check.equal(callback.playground, self.playground)
+
+
+class TestPlaygroundReset:
+
+    """ Playground._reset """
+
+    @pytest.fixture(autouse=True)
+    def setup_playground(self):
+        """Setup of used fixtures"""
+        self.env = Env()
+        self.n_agents = 5
+        self.agents = [Agent() for _ in range(self.n_agents)]
+        self.playground = Playground(self.env, self.agents)
+
+    def test_no_handlers(self, mocker):
+        """ should reset the environment with no handlers. """
+
+        mocker.patch('gym.Env.reset', lambda self:'obs')
+
+        observation, step, done, previous = self.playground._reset(None, None)
+        check.equal(observation, 'obs')
+        check.equal(step, 0)
+        check.equal(done, False)
+        expected_previous = [
+            {'observation':None, 'action':None, 'reward':None, 'done':None, 'info':None}
+            for _ in range(len(self.agents))
+        ]
+        check.equal(previous, expected_previous)
+
+    def test_callable_handlers(self, mocker):
+        """ should reset the environment with callable handlers. """
+
+        def done_handler(**kwargs):
+            return not kwargs.get('done')
+
+        def reward_handler(**kwargs):
+            return 2 * kwargs.get('reward')
+
+        mocker.patch('gym.Env.reset', lambda self:'obs')
+
+        observation, step, done, previous = self.playground._reset(reward_handler, done_handler)
+        check.equal(observation, 'obs')
+        check.equal(step, 0)
+        check.equal(done, False)
+        expected_previous = [
+            {'observation':None, 'action':None, 'reward':None, 'done':None, 'info':None}
+            for _ in range(len(self.agents))
+        ]
+        check.equal(previous, expected_previous)
+
+    def test_with_handlers(self, mocker):
+        """ should reset the environment and handlers with true handlers. """
+
+        mocker.patch('gym.Env.reset', lambda self:'obs')
+
+        mocker.patch('learnrl.playground.RewardHandler.reset')
+        reward_handler = RewardHandler()
+
+        mocker.patch('learnrl.playground.DoneHandler.reset')
+        done_handler = DoneHandler()
+
+        observation, step, done, previous = self.playground._reset(reward_handler, done_handler)
+
+        check.equal(observation, 'obs')
+        check.equal(step, 0)
+        check.equal(done, False)
+        expected_previous = [
+            {'observation':None, 'action':None, 'reward':None, 'done':None, 'info':None}
+            for _ in range(len(self.agents))
+        ]
+        check.equal(previous, expected_previous)
+
+        check.is_true(reward_handler.reset.called)
+        check.is_true(done_handler.reset.called)
+
+
+class TestPlaygroundAgentOrder:
+
+    """ Playground.set_agent_order """
+
+    @pytest.fixture(autouse=True)
+    def setup_playground(self):
+        """Setup of used fixtures"""
+        self.env = Env()
+        self.n_agents = 5
+        self.agents = [Agent() for _ in range(self.n_agents)]
+        self.playground = Playground(self.env, self.agents)
+
+    def test_default(self):
+        """ should have a correct default agent order. """
+        check.equal(self.playground.agents_order, list(range(self.n_agents)),
+            f"Default agents_order shoud be {list(range(self.n_agents))}"
+            f"but was {self.playground.agents_order}"
+        )
+
+    def test_custom_at_init(self):
+        """ should be able to have custom order at initialization. """
+        custom_order = [4, 3, 1, 2, 0]
+        playground = Playground(self.env, self.agents, agents_order=custom_order)
+        check.equal(playground.agents_order, custom_order,
+            f"Custom agents_order shoud be {custom_order}"
+            f"but was {playground.agents_order}"
+        )
+
+    def test_custom_after_init(self):
+        """ should be able to set custom order after initialization. """
+        new_custom_order = [3, 4, 1, 2, 0]
+        self.playground.set_agents_order(new_custom_order)
+
+        check.equal(self.playground.agents_order, new_custom_order,
+            f"Custom agents_order shoud be {new_custom_order}"
+            f"but was {self.playground.agents_order}"
+        )
+
+    def test_not_enought_indexes(self):
+        """ should raise ValueError if not enough indexes in custom order. """
+        with pytest.raises(ValueError, match=r"Not every agents*"):
+            Playground(self.env, self.agents, agents_order=[4, 3, 1, 2])
+
+    def test_missing_indexes(self):
+        """ should raise ValueError if missing indexes in custom order. """
+        with pytest.raises(ValueError, match=r".*not taking every index*"):
+            Playground(self.env, self.agents, agents_order=[4, 6, 1, 2, 0])
+
+
+class TestPlaygroundGetNextAgent:
+    """ Playground._get_next_agent """
+
+    @pytest.fixture(autouse=True)
+    def setup_playground(self):
+        """Setup of used fixtures"""
+        self.n_agents = 5
+        self.agents = [Agent() for _ in range(self.n_agents)]
+
+    def test_no_turnenv(self):
+        """ should return the first agent if env is not a TurnEnv. """
+        playground = Playground(Env(), self.agents)
+        _, agent_id = playground._get_next_agent('observation')
+        check.equal(agent_id, 0)
+
+    def test_turnenv(self, mocker):
+        """ should return the agent designed by agent_order and turn if env is a TurnEnv. """
+        mocker.patch('learnrl.envs.TurnEnv.turn', return_value=2)
+        playground = Playground(TurnEnv(), self.agents)
+        playground.agents_order = [3, 2, 1, 0, 4]
+        _, agent_id = playground._get_next_agent('observation')
+        check.equal(agent_id, 1)
+
+    def test_turnenv_indexerror(self, mocker):
+        """ should raise ValueError if turn result is out of agent_order. """
+        mocker.patch('learnrl.envs.TurnEnv.turn', return_value=10)
+        playground = Playground(TurnEnv(), self.agents)
+        with pytest.raises(ValueError, match=r'Not enough agents.*'):
+            playground._get_next_agent('observation')
+
+
+class TestPlaygroundCallHandlers:
+    """Playground._call_handlers"""
+
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        """Setup of used fixtures"""
+        self.call_handlers = Playground._call_handlers
+        self.done = False
+        self.reward = 1.2
+        self.experience = {
+            'observation': 'observation',
+            'action': 'action',
+            'reward': self.reward,
+            'done': self.done,
+            'next_observation': 'next_observation',
+            'info': {'info': 'info'}
+        }
+
+
+    def test_no_handlers(self):
+        """ should not log or change experience if no handler is given. """
+        logs = {}
+        self.call_handlers(
+            self.reward,
+            self.done,
+            self.experience,
+            reward_handler=None,
+            done_handler=None,
+            logs=logs
+        )
+        check.equal(self.experience['done'], self.done)
+        check.equal(logs['done'], self.done)
+
+        check.equal(self.experience['reward'], self.reward)
+        check.equal(logs['reward'], self.reward)
+
+    def test_done_only(self):
+        """ should log and change experience correctly with a done_handler only. """
+        logs = {}
+        handled_done = not self.done
+        self.call_handlers(
+            self.reward,
+            self.done,
+            self.experience,
+            reward_handler=None,
+            done_handler=lambda *args: handled_done,
+            logs=logs
+        )
+        check.equal(self.experience['done'], handled_done)
+        check.equal(logs['done'], self.done)
+        check.equal(logs['handled_done'], handled_done)
+
+        check.equal(self.experience['reward'], self.reward)
+        check.equal(logs['reward'], self.reward)
+
+    def test_reward_only(self):
+        """ should log and change experience correctly with a reward_handler only. """
+        logs = {}
+        handled_reward = 1.7
+        self.call_handlers(
+            self.reward,
+            self.done,
+            self.experience,
+            reward_handler=lambda *args:handled_reward,
+            done_handler=None,
+            logs=logs
+        )
+        check.equal(self.experience['done'], self.done)
+        check.equal(logs['done'], self.done)
+
+        check.equal(self.experience['reward'], handled_reward)
+        check.equal(logs['reward'], self.reward)
+        check.equal(logs['handled_reward'], handled_reward)
+
+    def test_both_handlers(self):
+        """ should log and change experience correctly with a both handlers. """
+        logs = {}
+        handled_done = not self.done
+        handled_reward = 1.7
+        self.call_handlers(
+            self.reward,
+            self.done,
+            self.experience,
+            reward_handler=lambda *args: handled_reward,
+            done_handler=lambda *args: handled_done,
+            logs=logs
+        )
+        check.equal(self.experience['done'], handled_done)
+        check.equal(logs['done'], self.done)
+        check.equal(logs['handled_done'], handled_done)
+
+        check.equal(self.experience['reward'], handled_reward)
+        check.equal(logs['reward'], self.reward)
+        check.equal(logs['handled_reward'], handled_reward)
+
+class TestPlaygroundRunStep:
+    """Playground._run_step"""
+
+    @pytest.fixture(autouse=True)
+    def setup_playground(self, mocker):
+        """Setup of used fixtures"""
+
+        self.observation = 'observation'
+        self.next_observation = 'next_observation'
+        self.reward = 1.2
+        self.done = False
+        self.info = {'env_info': 'env_info'}
+
+        mocker.patch('gym.Env.render')
+        mocker.patch('gym.Env.step',
+            return_value=(self.next_observation, self.reward,
+                self.done, self.info)
+        )
+        self.env = Env()
+
+        mocker.patch('learnrl.agent.Agent.remember')
+        mocker.patch('learnrl.agent.Agent.learn')
+        self.action = 3
+        mocker.patch('learnrl.agent.Agent.act', return_value=self.action)
+        self.n_agents = 5
+        self.agents = [Agent() for _ in range(self.n_agents)]
+
+        self.agent_id = 0
+        mocker.patch('learnrl.playground.Playground._get_next_agent',
+            return_value=(self.agents[self.agent_id], self.agent_id))
+        mocker.patch('learnrl.playground.Playground._call_handlers')
+        self.playground = Playground(self.env, self.agents)
+
+        self.previous = [
+            {'observation':None,'action':None,
+            'reward':None, 'done':None, 'info':None}
+            for _ in range(self.n_agents)
+        ]
+
+    def test_run_step(self):
+        """ should update the observation, done and logs correcty. """
+        logs = {}
+        done = not self.done
+        observation, done = self.playground._run_step(
+            self.observation,
+            self.previous,
+            logs=logs,
+        )
+        check.equal(observation, self.next_observation)
+        check.equal(done, self.done)
+
+        for log_name in ['reward', 'observation', 'next_observation',
+                            'info', 'done', 'agent_id', 'action']:
+            expected = getattr(self, log_name)
+            check.equal(logs[log_name], expected)
+
+    def test_render_not_done(self, mocker):
+        """ should render at each step beginning if not done. """
+        render_mode = 'render_mode'
+        self.playground._run_step(
+            self.observation,
+            self.previous,
+            logs={},
+            render=True,
+            render_mode=render_mode
+        )
+
+        check.equal(len(self.env.render.call_args_list), 1)
+        render_args, _ = self.env.render.call_args
+        check.equal(render_args[0], render_mode)
+
+    def test_render_done(self, mocker):
+        """ should render at step beginning and end if done. """
+        render_mode = 'render_mode'
+        self.env.step.return_value = (
+            self.next_observation, self.reward, True, self.info
+        )
+        self.playground._run_step(
+            self.observation,
+            self.previous,
+            logs={},
+            render=True,
+            render_mode=render_mode
+        )
+
+        check.equal(len(self.env.render.call_args_list), 2)
+        render_args, _ = self.env.render.call_args_list[-1]
+        check.equal(render_args[0], render_mode)
+
+    def test_learn_without_prev_not_done(self, mocker):
+        """ should store experience without learn or remember for first experience. """
+        previous = self.previous
+        self.playground._run_step(
+            self.observation,
+            previous,
+            logs={},
+            learn=True,
+        )
+        check.equal(previous[self.agent_id],
+            {name: getattr(self, name)
+            for name in ['observation', 'action', 'reward', 'done', 'info']}
+        )
+        check.is_false(Agent.learn.called)
+        check.is_false(Agent.remember.called)
+
+    def test_learn_not_done(self, mocker):
+        """ should call learn and remember once at the beginning of the step if not done. """
+        logs = {}
+        previous = self.previous
+        previous[0] = {
+            name: getattr(self, name)
+            for name in ['observation', 'action', 'reward', 'done', 'info']
+        }
+        self.playground._run_step(
+            self.observation,
+            previous,
+            logs=logs,
+            learn=True,
+        )
+
+        check.equal(len(Agent.learn.call_args_list), 1)
+        check.equal(len(Agent.remember.call_args_list), 1)
+        check.is_in('agent_0', logs)
+
+    def test_learn_done(self, mocker):
+        """ should call learn and remember once at the beginning of the step if not done. """
+        logs = {}
+        previous = self.previous
+        self.env.step.return_value = (
+            self.next_observation, self.reward, True, self.info
+        )
+        self.playground._run_step(
+            self.observation,
+            previous,
+            logs=logs,
+            learn=True,
+        )
+
+        check.equal(len(Agent.learn.call_args_list), 1)
+        check.equal(len(Agent.remember.call_args_list), 1)
+        check.is_in('agent_0', logs)
+
+class TestDoneHandler:
+
+    """DoneHandler"""
+
+    def test_done(self, mocker):
+        """ should use `done` if the output of `done` is a bool or a bool numpy array. """
+        mocker.patch("learnrl.playground.DoneHandler.done", lambda *args: True)
+        handler = DoneHandler()
+        check.is_true(handler._done())
 
+        mocker.patch("learnrl.playground.DoneHandler.done", lambda *args: np.array(True))
+        handler = DoneHandler()
+        check.is_true(handler._done())
 
-def test_single_agent():
-    env = Env()
-    agents = Agent()
-    pg = Playground(env, agents)
-    assert pg.agents == [agents], \
-        "A single agent should be transformed in a list containing itself"
+    def test_not_bool_done(self, mocker):
+        """ should raise ValueError if the output of `done` is not a bool. """
+        mocker.patch("learnrl.playground.DoneHandler.done", lambda *args: 'True')
+        handler = DoneHandler()
+        with pytest.raises(ValueError, match=r"Done should be bool.*"):
+            handler._done()
 
-def test_agent_order():
-    env = Env()
-    n_agents = 5
-    agents = [Agent() for _ in range(n_agents)]
+    def test_call(self, mocker):
+        """ should call done on class call. """
+        mocker.patch("learnrl.playground.DoneHandler._done", return_value=True)
+        handler = DoneHandler()
+        check.is_true(handler())
+        check.is_true(handler._done.called)
 
-    # Default order
-    pg = Playground(env, agents)
-    assert pg.agents_order == list(range(n_agents)), \
-        f"Default agents_order shoud be {list(range(n_agents))} but was {pg.agents_order}"
 
-    # Custom order at init
-    custom_order = [4, 3, 1, 2, 0]
-    pg = Playground(env, agents, agents_order=custom_order)
-    assert pg.agents_order == custom_order, \
-        f"Custom agents_order shoud be {custom_order} but was {pg.agents_order}"
+class TestRewardHandler:
 
-    # Not enough indexes in custom order
-    with pytest.raises(ValueError, match=r"Not every agents*"):
-        pg = Playground(env, agents, agents_order=[4, 3, 1, 2])
+    """RewardHandler"""
 
-    # Missing indexes in custom order
-    with pytest.raises(ValueError, match=r".*not taking every index*"):
-        pg = Playground(env, agents, agents_order=[4, 6, 1, 2, 0])
+    def test_reward(self, mocker):
+        """ should use `reward` if the output of `reward` is a float or floating numpy array. """
+        mocker.patch("learnrl.playground.RewardHandler.reward", lambda *args: 1.2)
+        handler = RewardHandler()
+        check.equal(handler._reward(), 1.2)
 
-    # Custom order after init
-    custom_order = [4, 3, 1, 2, 0]
-    pg = Playground(env, agents, agents_order=custom_order)
+        mocker.patch("learnrl.playground.RewardHandler.reward", lambda *args: np.array(1.2))
+        handler = RewardHandler()
+        check.equal(handler._reward(), 1.2)
 
-    new_custom_order = [3, 4, 1, 2, 0]
-    pg.set_agents_order(new_custom_order)
+    def test_not_scalar_reward(self, mocker):
+        """ should raise ValueError if the output of `reward` is not a float. """
+        mocker.patch("learnrl.playground.RewardHandler.reward", lambda *args: '1.2')
+        handler = RewardHandler()
+        with pytest.raises(ValueError, match=r"Rewards should be a float.*"):
+            handler._reward()
 
-    assert pg.agents_order == new_custom_order, \
-        f"Custom agents_order shoud be {new_custom_order} but was {pg.agents_order}"
+    def test_call(self, mocker):
+        """ should call reward on class call. """
+        mocker.patch("learnrl.playground.RewardHandler._reward", return_value=True)
+        handler = RewardHandler()
+        check.is_true(handler())
+        check.is_true(handler._reward.called)