diff --git a/rl_coach/level_manager.py b/rl_coach/level_manager.py
index 945ab6c..40cb948 100644
--- a/rl_coach/level_manager.py
+++ b/rl_coach/level_manager.py
@@ -17,6 +17,7 @@ import copy
 from typing import Union, Dict
 
 from rl_coach.agents.composite_agent import CompositeAgent
+from rl_coach.agents.agent_interface import AgentInterface
 from rl_coach.core_types import EnvResponse, ActionInfo, RunPhase, ActionType, EnvironmentSteps, Transition
 from rl_coach.environments.environment import Environment
 from rl_coach.environments.environment_interface import EnvironmentInterface
@@ -38,7 +39,7 @@ class LevelManager(EnvironmentInterface):
     """
     def __init__(self,
                  name: str,
-                 agents: Union['Agent', CompositeAgent, Dict[str, Union['Agent', CompositeAgent]]],
+                 agents: Union[AgentInterface, Dict[str, AgentInterface]],
                  environment: Union['LevelManager', Environment],
                  real_environment: Environment = None,
                  steps_limit: EnvironmentSteps = EnvironmentSteps(1),
@@ -48,7 +49,7 @@ class LevelManager(EnvironmentInterface):
         """
         A level manager controls a single or multiple composite agents and a single environment.
         The environment can be either a real environment or another level manager behaving as an environment.
-        :param agents: a list of agents or composite agents to control
+        :param agents: a single agent or a dictionary of agents or composite agents to control
         :param environment: an environment or level manager to control
         :param real_environment: the real environment that is is acted upon. if this is None (which it should be for
          the most bottom level), it will be replaced by the environment parameter. For simple RL schemes, where there