[RLlib] Allow rllib rollout to run distributed via evaluation workers. (#13718)

rllib/BUILD

@@ -517,7 +517,7 @@ py_test(
 py_test(
     name = "test_marwil",
     tags = ["agents_dir"],
-    size = "medium",
+    size = "large",
     # Include the json data file.
     data = ["tests/data/cartpole/large.json"],
     srcs = ["agents/marwil/tests/test_marwil.py"]
@@ -527,7 +527,7 @@ py_test(
 py_test(
     name = "test_bc",
     tags = ["agents_dir"],
-    size = "medium",
+    size = "large",
     # Include the json data file.
     data = ["tests/data/cartpole/large.json"],
     srcs = ["agents/marwil/tests/test_bc.py"]
@@ -1753,7 +1753,7 @@ py_test(
     name = "examples/custom_eval_tf",
     main = "examples/custom_eval.py",
     tags = ["examples", "examples_C"],
-    size = "small",
+    size = "medium",
     srcs = ["examples/custom_eval.py"],
     args = ["--num-cpus=4", "--as-test"]
 )
@@ -1762,7 +1762,7 @@ py_test(
     name = "examples/custom_eval_torch",
     main = "examples/custom_eval.py",
     tags = ["examples", "examples_C"],
-    size = "small",
+    size = "medium",
     srcs = ["examples/custom_eval.py"],
     args = ["--num-cpus=4", "--as-test", "--torch"]
 )

rllib/__init__.py

@@ -27,12 +27,12 @@ def _setup_logger():
 
 def _register_all():
     from ray.rllib.agents.trainer import Trainer, with_common_config
-    from ray.rllib.agents.registry import ALGORITHMS, get_agent_class
+    from ray.rllib.agents.registry import ALGORITHMS, get_trainer_class
     from ray.rllib.contrib.registry import CONTRIBUTED_ALGORITHMS
 
     for key in list(ALGORITHMS.keys()) + list(CONTRIBUTED_ALGORITHMS.keys(
     )) + ["__fake", "__sigmoid_fake_data", "__parameter_tuning"]:
-        register_trainable(key, get_agent_class(key))
+        register_trainable(key, get_trainer_class(key))
 
     def _see_contrib(name):
         """Returns dummy agent class warning algo is in contrib/."""

rllib/agents/mock.py

@@ -118,14 +118,14 @@ class _ParameterTuningTrainer(_MockTrainer):
             info={})
 
 
-def _agent_import_failed(trace):
+def _trainer_import_failed(trace):
     """Returns dummy agent class for if PyTorch etc. is not installed."""
 
-    class _AgentImportFailed(Trainer):
-        _name = "AgentImportFailed"
+    class _TrainerImportFailed(Trainer):
+        _name = "TrainerImportFailed"
         _default_config = with_common_config({})
 
         def setup(self, config):
             raise ImportError(trace)
 
-    return _AgentImportFailed
+    return _TrainerImportFailed

rllib/agents/registry.py

@@ -3,126 +3,127 @@
 import traceback
 
 from ray.rllib.contrib.registry import CONTRIBUTED_ALGORITHMS
+from ray.rllib.utils.deprecation import deprecation_warning
 
 
 def _import_a2c():
     from ray.rllib.agents import a3c
-    return a3c.A2CTrainer
+    return a3c.A2CTrainer, a3c.a2c.A2C_DEFAULT_CONFIG
 
 
 def _import_a3c():
     from ray.rllib.agents import a3c
-    return a3c.A3CTrainer
+    return a3c.A3CTrainer, a3c.DEFAULT_CONFIG
 
 
 def _import_apex():
     from ray.rllib.agents import dqn
-    return dqn.ApexTrainer
+    return dqn.ApexTrainer, dqn.apex.APEX_DEFAULT_CONFIG
 
 
 def _import_apex_ddpg():
     from ray.rllib.agents import ddpg
-    return ddpg.ApexDDPGTrainer
+    return ddpg.ApexDDPGTrainer, ddpg.apex.APEX_DDPG_DEFAULT_CONFIG
 
 
 def _import_appo():
     from ray.rllib.agents import ppo
-    return ppo.APPOTrainer
+    return ppo.APPOTrainer, ppo.appo.DEFAULT_CONFIG
 
 
 def _import_ars():
     from ray.rllib.agents import ars
-    return ars.ARSTrainer
+    return ars.ARSTrainer, ars.DEFAULT_CONFIG
 
 
 def _import_bc():
     from ray.rllib.agents import marwil
-    return marwil.BCTrainer
+    return marwil.BCTrainer, marwil.DEFAULT_CONFIG
 
 
 def _import_cql():
     from ray.rllib.agents import cql
-    return cql.CQLTrainer
+    return cql.CQLTrainer, cql.CQL_DEFAULT_CONFIG
 
 
 def _import_ddpg():
     from ray.rllib.agents import ddpg
-    return ddpg.DDPGTrainer
+    return ddpg.DDPGTrainer, ddpg.DEFAULT_CONFIG
 
 
 def _import_ddppo():
     from ray.rllib.agents import ppo
-    return ppo.DDPPOTrainer
+    return ppo.DDPPOTrainer, ppo.DEFAULT_CONFIG
 
 
 def _import_dqn():
     from ray.rllib.agents import dqn
-    return dqn.DQNTrainer
+    return dqn.DQNTrainer, dqn.DEFAULT_CONFIG
 
 
 def _import_dreamer():
     from ray.rllib.agents import dreamer
-    return dreamer.DREAMERTrainer
+    return dreamer.DREAMERTrainer, dreamer.DEFAULT_CONFIG
 
 
 def _import_es():
     from ray.rllib.agents import es
-    return es.ESTrainer
+    return es.ESTrainer, es.DEFAULT_CONFIG
 
 
 def _import_impala():
     from ray.rllib.agents import impala
-    return impala.ImpalaTrainer
+    return impala.ImpalaTrainer, impala.DEFAULT_CONFIG
 
 
 def _import_maml():
     from ray.rllib.agents import maml
-    return maml.MAMLTrainer
+    return maml.MAMLTrainer, maml.DEFAULT_CONFIG
 
 
 def _import_marwil():
     from ray.rllib.agents import marwil
-    return marwil.MARWILTrainer
+    return marwil.MARWILTrainer, marwil.DEFAULT_CONFIG
 
 
 def _import_mbmpo():
     from ray.rllib.agents import mbmpo
-    return mbmpo.MBMPOTrainer
+    return mbmpo.MBMPOTrainer, mbmpo.DEFAULT_CONFIG
 
 
 def _import_pg():
     from ray.rllib.agents import pg
-    return pg.PGTrainer
+    return pg.PGTrainer, pg.DEFAULT_CONFIG
 
 
 def _import_ppo():
     from ray.rllib.agents import ppo
-    return ppo.PPOTrainer
+    return ppo.PPOTrainer, ppo.DEFAULT_CONFIG
 
 
 def _import_qmix():
     from ray.rllib.agents import qmix
-    return qmix.QMixTrainer
+    return qmix.QMixTrainer, qmix.DEFAULT_CONFIG
 
 
 def _import_sac():
     from ray.rllib.agents import sac
-    return sac.SACTrainer
+    return sac.SACTrainer, sac.DEFAULT_CONFIG
 
 
 def _import_simple_q():
     from ray.rllib.agents import dqn
-    return dqn.SimpleQTrainer
+    return dqn.SimpleQTrainer, dqn.simple_q.DEFAULT_CONFIG
 
 
 def _import_slate_q():
     from ray.rllib.agents import slateq
-    return slateq.SlateQTrainer
+    return slateq.SlateQTrainer, slateq.DEFAULT_CONFIG
 
 
 def _import_td3():
     from ray.rllib.agents import ddpg
-    return ddpg.TD3Trainer
+    return ddpg.TD3Trainer, ddpg.td3.TD3_DEFAULT_CONFIG
 
 
 ALGORITHMS = {
@@ -153,32 +154,47 @@ ALGORITHMS = {
 }
 
 
-def get_agent_class(alg: str) -> type:
-    """Returns the class of a known agent given its name."""
+def get_trainer_class(alg: str, return_config=False) -> type:
+    """Returns the class of a known Trainer given its name."""
 
     try:
-        return _get_agent_class(alg)
+        return _get_trainer_class(alg, return_config=return_config)
     except ImportError:
-        from ray.rllib.agents.mock import _agent_import_failed
-        return _agent_import_failed(traceback.format_exc())
+        from ray.rllib.agents.mock import _trainer_import_failed
+        class_ = _trainer_import_failed(traceback.format_exc())
+        config = class_._default_config
+        if return_config:
+            return class_, config
+        return class_
 
 
-def _get_agent_class(alg: str) -> type:
+# Deprecated: Use `get_trainer_class` instead.
+def get_agent_class(alg: str) -> type:
+    deprecation_warning("get_agent_class", "get_trainer_class", error=False)
+    return get_trainer_class(alg)
+
+
+def _get_trainer_class(alg: str, return_config=False) -> type:
     if alg in ALGORITHMS:
-        return ALGORITHMS[alg]()
+        class_, config = ALGORITHMS[alg]()
     elif alg in CONTRIBUTED_ALGORITHMS:
-        return CONTRIBUTED_ALGORITHMS[alg]()
+        class_, config = CONTRIBUTED_ALGORITHMS[alg]()
     elif alg == "script":
         from ray.tune import script_runner
-        return script_runner.ScriptRunner
+        class_, config = script_runner.ScriptRunner, {}
     elif alg == "__fake":
         from ray.rllib.agents.mock import _MockTrainer
-        return _MockTrainer
+        class_, config = _MockTrainer, _MockTrainer._default_config
     elif alg == "__sigmoid_fake_data":
         from ray.rllib.agents.mock import _SigmoidFakeData
-        return _SigmoidFakeData
+        class_, config = _SigmoidFakeData, _SigmoidFakeData._default_config
     elif alg == "__parameter_tuning":
         from ray.rllib.agents.mock import _ParameterTuningTrainer
-        return _ParameterTuningTrainer
+        class_, config = _ParameterTuningTrainer, \
+            _ParameterTuningTrainer._default_config
     else:
         raise Exception(("Unknown algorithm {}.").format(alg))
+
+    if return_config:
+        return class_, config
+    return class_

rllib/agents/trainer.py

@@ -52,7 +52,7 @@ COMMON_CONFIG: TrainerConfigDict = {
     # Number of rollout worker actors to create for parallel sampling. Setting
     # this to 0 will force rollouts to be done in the trainer actor.
     "num_workers": 2,
-    # Number of environments to evaluate vectorwise per worker. This enables
+    # Number of environments to evaluate vector-wise per worker. This enables
     # model inference batching, which can improve performance for inference
     # bottlenecked workloads.
     "num_envs_per_worker": 1,
@@ -120,10 +120,18 @@ COMMON_CONFIG: TrainerConfigDict = {
     # set this if soft_horizon=True, unless your env is actually running
     # forever without returning done=True.
     "no_done_at_end": False,
-    # Arguments to pass to the env creator.
-    "env_config": {},
     # Environment name can also be passed via config.
     "env": None,
+    # Arguments to pass to the env creator.
+    "env_config": {},
+    # If True, try to render the environment on the local worker or on worker
+    # 1 (if num_workers > 0). For vectorized envs, this usually means that only
+    # the first sub-environment will be rendered.
+    "render_env": False,
+    # If True, store evaluation videos in the output dir.
+    # Alternatively, provide a path (str) to a directory here, where the env
+    # recordings should be stored instead.
+    "record_env": False,
     # Unsquash actions to the upper and lower bounds of env's action space
     "normalize_actions": False,
     # Whether to clip rewards during Policy's postprocessing.
@@ -213,9 +221,10 @@ COMMON_CONFIG: TrainerConfigDict = {
     },
     # Number of parallel workers to use for evaluation. Note that this is set
     # to zero by default, which means evaluation will be run in the trainer
-    # process. If you increase this, it will increase the Ray resource usage
-    # of the trainer since evaluation workers are created separately from
-    # rollout workers.
+    # process (only if evaluation_interval is not None). If you increase this,
+    # it will increase the Ray resource usage of the trainer since evaluation
+    # workers are created separately from rollout workers (used to sample data
+    # for training).
     "evaluation_num_workers": 0,
     # Customize the evaluation method. This must be a function of signature
     # (trainer: Trainer, eval_workers: WorkerSet) -> metrics: dict. See the
@@ -662,7 +671,6 @@ class Trainer(Trainable):
                     extra_config["in_evaluation"] is True
                 extra_config.update({
                     "batch_mode": "complete_episodes",
-                    "rollout_fragment_length": 1,
                     "in_evaluation": True,
                 })
                 logger.debug(

rllib/contrib/registry.py

@@ -3,28 +3,29 @@
 
 def _import_random_agent():
     from ray.rllib.contrib.random_agent.random_agent import RandomAgent
-    return RandomAgent
+    return RandomAgent, RandomAgent._default_config
 
 
 def _import_maddpg():
     from ray.rllib.contrib import maddpg
-    return maddpg.MADDPGTrainer
+    return maddpg.MADDPGTrainer, maddpg.DEFAULT_CONFIG
 
 
 def _import_alphazero():
     from ray.rllib.contrib.alpha_zero.core.alpha_zero_trainer import\
-        AlphaZeroTrainer
-    return AlphaZeroTrainer
+        AlphaZeroTrainer, DEFAULT_CONFIG
+    return AlphaZeroTrainer, DEFAULT_CONFIG
 
 
 def _import_bandit_lints():
-    from ray.rllib.contrib.bandits.agents.lin_ts import LinTSTrainer
-    return LinTSTrainer
+    from ray.rllib.contrib.bandits.agents.lin_ts import LinTSTrainer, TS_CONFIG
+    return LinTSTrainer, TS_CONFIG
 
 
 def _import_bandit_linucb():
-    from ray.rllib.contrib.bandits.agents.lin_ucb import LinUCBTrainer
-    return LinUCBTrainer
+    from ray.rllib.contrib.bandits.agents.lin_ucb import LinUCBTrainer, \
+        UCB_CONFIG
+    return LinUCBTrainer, UCB_CONFIG
 
 
 CONTRIBUTED_ALGORITHMS = {

rllib/env/base_env.py

@@ -5,8 +5,8 @@ from ray.rllib.env.external_multi_agent_env import ExternalMultiAgentEnv
 from ray.rllib.env.multi_agent_env import MultiAgentEnv
 from ray.rllib.env.vector_env import VectorEnv
 from ray.rllib.utils.annotations import override, PublicAPI
-from ray.rllib.utils.typing import EnvType, MultiEnvDict, EnvID, \
-    AgentID, MultiAgentDict
+from ray.rllib.utils.typing import AgentID, EnvID, EnvType, MultiAgentDict, \
+    MultiEnvDict, PartialTrainerConfigDict
 
 if TYPE_CHECKING:
     from ray.rllib.models.preprocessors import Preprocessor
@@ -80,11 +80,14 @@ class BaseEnv:
     """
 
     @staticmethod
-    def to_base_env(env: EnvType,
-                    make_env: Callable[[int], EnvType] = None,
-                    num_envs: int = 1,
-                    remote_envs: bool = False,
-                    remote_env_batch_wait_ms: int = 0) -> "BaseEnv":
+    def to_base_env(
+            env: EnvType,
+            make_env: Callable[[int], EnvType] = None,
+            num_envs: int = 1,
+            remote_envs: bool = False,
+            remote_env_batch_wait_ms: int = 0,
+            policy_config: PartialTrainerConfigDict = None,
+    ) -> "BaseEnv":
         """Wraps any env type as needed to expose the async interface."""
 
         from ray.rllib.env.remote_vector_env import RemoteVectorEnv
@@ -129,7 +132,9 @@ class BaseEnv:
                         existing_envs=[env],
                         num_envs=num_envs,
                         action_space=env.action_space,
-                        observation_space=env.observation_space)
+                        observation_space=env.observation_space,
+                        policy_config=policy_config,
+                    )
                     env = _VectorEnvToBaseEnv(env)
         assert isinstance(env, BaseEnv), env
         return env
@@ -205,6 +210,18 @@ class BaseEnv:
             if hasattr(env, "close"):
                 env.close()
 
+    # Experimental method.
+    def try_render(self, env_id: Optional[EnvID] = None) -> None:
+        """Tries to render the environment.
+
+        Args:
+            env_id (Optional[int]): The sub-env ID if applicable. If None,
+                renders the entire Env (i.e. all sub-envs).
+        """
+
+        # By default, do nothing.
+        pass
+
 
 # Fixed agent identifier when there is only the single agent in the env
 _DUMMY_AGENT_ID = "agent0"
@@ -346,14 +363,19 @@ class _VectorEnvToBaseEnv(BaseEnv):
             self.vector_env.vector_step(action_vector)
 
     @override(BaseEnv)
-    def try_reset(self,
-                  env_id: Optional[EnvID] = None) -> Optional[MultiAgentDict]:
+    def try_reset(self, env_id: Optional[EnvID] = None) -> MultiAgentDict:
+        assert env_id is None or isinstance(env_id, int)
         return {_DUMMY_AGENT_ID: self.vector_env.reset_at(env_id)}
 
     @override(BaseEnv)
     def get_unwrapped(self) -> List[EnvType]:
         return self.vector_env.get_unwrapped()
 
+    @override(BaseEnv)
+    def try_render(self, env_id: Optional[EnvID] = None) -> None:
+        assert env_id is None or isinstance(env_id, int)
+        return self.vector_env.try_render_at(env_id)
+
 
 class _MultiAgentEnvToBaseEnv(BaseEnv):
     """Internal adapter of MultiAgentEnv to BaseEnv.

rllib/env/vector_env.py

@@ -1,11 +1,12 @@
 import logging
 import gym
+from gym import wrappers as gym_wrappers
 import numpy as np
-from typing import Callable, List, Tuple
+from typing import Callable, List, Optional, Tuple
 
 from ray.rllib.utils.annotations import override, PublicAPI
-from ray.rllib.utils.typing import EnvType, EnvConfigDict, EnvObsType, \
-    EnvInfoDict, EnvActionType
+from ray.rllib.utils.typing import EnvActionType, EnvConfigDict, EnvInfoDict, \
+    EnvObsType, EnvType, PartialTrainerConfigDict
 
 logger = logging.getLogger(__name__)
 
@@ -30,19 +31,22 @@ class VectorEnv:
         self.num_envs = num_envs
 
     @staticmethod
-    def wrap(make_env: Callable[[int], EnvType] = None,
-             existing_envs: List[gym.Env] = None,
+    def wrap(make_env: Optional[Callable[[int], EnvType]] = None,
+             existing_envs: Optional[List[gym.Env]] = None,
              num_envs: int = 1,
-             action_space: gym.Space = None,
-             observation_space: gym.Space = None,
-             env_config: EnvConfigDict = None):
+             action_space: Optional[gym.Space] = None,
+             observation_space: Optional[gym.Space] = None,
+             env_config: Optional[EnvConfigDict] = None,
+             policy_config: Optional[PartialTrainerConfigDict] = None):
         return _VectorizedGymEnv(
             make_env=make_env,
             existing_envs=existing_envs or [],
             num_envs=num_envs,
             observation_space=observation_space,
             action_space=action_space,
-            env_config=env_config)
+            env_config=env_config,
+            policy_config=policy_config,
+        )
 
     @PublicAPI
     def vector_reset(self) -> List[EnvObsType]:
@@ -54,9 +58,12 @@ class VectorEnv:
         raise NotImplementedError
 
     @PublicAPI
-    def reset_at(self, index: int) -> EnvObsType:
+    def reset_at(self, index: Optional[int] = None) -> EnvObsType:
         """Resets a single environment.
 
+        Args:
+            index (Optional[int]): An optional sub-env index to reset.
+
         Returns:
             obs (obj): Observations from the reset sub environment.
         """
@@ -88,19 +95,31 @@ class VectorEnv:
         """
         raise NotImplementedError
 
+    # Experimental method.
+    def try_render_at(self, index: Optional[int] = None) -> None:
+        """Renders a single environment.
+
+        Args:
+            index (Optional[int]): An optional sub-env index to render.
+        """
+        pass
+
 
 class _VectorizedGymEnv(VectorEnv):
     """Internal wrapper to translate any gym envs into a VectorEnv object.
     """
 
-    def __init__(self,
-                 make_env=None,
-                 existing_envs=None,
-                 num_envs=1,
-                 *,
-                 observation_space=None,
-                 action_space=None,
-                 env_config=None):
+    def __init__(
+            self,
+            make_env=None,
+            existing_envs=None,
+            num_envs=1,
+            *,
+            observation_space=None,
+            action_space=None,
+            env_config=None,
+            policy_config=None,
+    ):
         """Initializes a _VectorizedGymEnv object.
 
         Args:
@@ -116,11 +135,27 @@ class _VectorizedGymEnv(VectorEnv):
                 If None, use existing_envs[0]'s action space.
             env_config (Optional[dict]): Additional sub env config to pass to
                 make_env as first arg.
+            policy_config (Optional[PartialTrainerConfigDict]): An optional
+                trainer/policy config dict.
         """
-        self.make_env = make_env
         self.envs = existing_envs
+
+        # Fill up missing envs (so we have exactly num_envs sub-envs in this
+        # VectorEnv.
         while len(self.envs) < num_envs:
-            self.envs.append(self.make_env(len(self.envs)))
+            self.envs.append(make_env(len(self.envs)))
+
+        # Wrap all envs with video recorder if necessary.
+        if policy_config is not None and policy_config.get("record_env"):
+
+            def wrapper_(env):
+                return gym_wrappers.Monitor(
+                    env=env,
+                    directory=policy_config["record_env"],
+                    video_callable=lambda _: True,
+                    force=True)
+
+            self.envs = [wrapper_(e) for e in self.envs]
 
         super().__init__(
             observation_space=observation_space
@@ -133,7 +168,9 @@ class _VectorizedGymEnv(VectorEnv):
         return [e.reset() for e in self.envs]
 
     @override(VectorEnv)
-    def reset_at(self, index):
+    def reset_at(self, index: Optional[int] = None) -> EnvObsType:
+        if index is None:
+            index = 0
         return self.envs[index].reset()
 
     @override(VectorEnv)
@@ -157,3 +194,9 @@ class _VectorizedGymEnv(VectorEnv):
     @override(VectorEnv)
     def get_unwrapped(self):
         return self.envs
+
+    @override(VectorEnv)
+    def try_render_at(self, index: Optional[int] = None):
+        if index is None:
+            index = 0
+        return self.envs[index].render()

rllib/evaluation/rollout_worker.py

@@ -546,7 +546,9 @@ class RolloutWorker(ParallelIteratorWorker):
                 make_env=make_env,
                 num_envs=num_envs,
                 remote_envs=remote_worker_envs,
-                remote_env_batch_wait_ms=remote_env_batch_wait_ms)
+                remote_env_batch_wait_ms=remote_env_batch_wait_ms,
+                policy_config=policy_config,
+            )
 
         # `truncate_episodes`: Allow a batch to contain more than one episode
         # (fragments) and always make the batch `rollout_fragment_length`
@@ -583,6 +585,11 @@ class RolloutWorker(ParallelIteratorWorker):
                 raise ValueError(
                     "Unknown evaluation method: {}".format(method))
 
+        render = False
+        if policy_config.get("render_env") is True and \
+                (num_workers == 0 or worker_index == 1):
+            render = True
+
         if self.env is None:
             self.sampler = None
         elif sample_async:
@@ -608,6 +615,7 @@ class RolloutWorker(ParallelIteratorWorker):
                 _use_trajectory_view_api=_use_trajectory_view_api,
                 sample_collector_class=policy_config.get(
                     "sample_collector_class"),
+                render=render,
             )
             # Start the Sampler thread.
             self.sampler.start()
@@ -633,6 +641,7 @@ class RolloutWorker(ParallelIteratorWorker):
                 _use_trajectory_view_api=_use_trajectory_view_api,
                 sample_collector_class=policy_config.get(
                     "sample_collector_class"),
+                render=render,
             )
 
         self.input_reader: InputReader = input_creator(self.io_context)

rllib/evaluation/sampler.py

@@ -65,17 +65,16 @@ class _PerfStats:
 
     def __init__(self):
         self.iters = 0
-        self.env_wait_time = 0.0
         self.raw_obs_processing_time = 0.0
         self.inference_time = 0.0
         self.action_processing_time = 0.0
+        self.env_wait_time = 0.0
+        self.env_render_time = 0.0
 
     def get(self):
         # Mean multiplicator (1000 = ms -> sec).
         factor = 1000 / self.iters
         return {
-            # Waiting for environment (during poll).
-            "mean_env_wait_ms": self.env_wait_time * factor,
             # Raw observation preprocessing.
             "mean_raw_obs_processing_ms": self.raw_obs_processing_time *
             factor,
@@ -83,6 +82,10 @@ class _PerfStats:
             "mean_inference_ms": self.inference_time * factor,
             # Processing actions (to be sent to env, e.g. clipping).
             "mean_action_processing_ms": self.action_processing_time * factor,
+            # Waiting for environment (during poll).
+            "mean_env_wait_ms": self.env_wait_time * factor,
+            # Environment rendering (False by default).
+            "mean_env_render_ms": self.env_render_time * factor,
         }
 
 
@@ -141,7 +144,9 @@ class SyncSampler(SamplerInput):
             no_done_at_end: bool = False,
             observation_fn: "ObservationFunction" = None,
             _use_trajectory_view_api: bool = False,
-            sample_collector_class: Optional[Type[SampleCollector]] = None):
+            sample_collector_class: Optional[Type[SampleCollector]] = None,
+            render: bool = False,
+    ):
         """Initializes a SyncSampler object.
 
         Args:
@@ -184,6 +189,8 @@ class SyncSampler(SamplerInput):
             sample_collector_class (Optional[Type[SampleCollector]]): An
                 optional Samplecollector sub-class to use to collect, store,
                 and retrieve environment-, model-, and sampler data.
+            render (bool): Whether to try to render the environment after each
+                step.
         """
 
         self.base_env = BaseEnv.to_base_env(env)
@@ -207,6 +214,7 @@ class SyncSampler(SamplerInput):
                 count_steps_by=count_steps_by)
         else:
             self.sample_collector = None
+        self.render = render
 
         # Create the rollout generator to use for calls to `get_data()`.
         self.rollout_provider = _env_runner(
@@ -215,7 +223,7 @@ class SyncSampler(SamplerInput):
             self.preprocessors, self.obs_filters, clip_rewards, clip_actions,
             multiple_episodes_in_batch, callbacks, tf_sess, self.perf_stats,
             soft_horizon, no_done_at_end, observation_fn,
-            _use_trajectory_view_api, self.sample_collector)
+            _use_trajectory_view_api, self.sample_collector, self.render)
         self.metrics_queue = queue.Queue()
 
     @override(SamplerInput)
@@ -280,6 +288,7 @@ class AsyncSampler(threading.Thread, SamplerInput):
             observation_fn: "ObservationFunction" = None,
             _use_trajectory_view_api: bool = False,
             sample_collector_class: Optional[Type[SampleCollector]] = None,
+            render: bool = False,
     ):
         """Initializes a AsyncSampler object.
 
@@ -327,6 +336,8 @@ class AsyncSampler(threading.Thread, SamplerInput):
             sample_collector_class (Optional[Type[SampleCollector]]): An
                 optional Samplecollector sub-class to use to collect, store,
                 and retrieve environment-, model-, and sampler data.
+            render (bool): Whether to try to render the environment after each
+                step.
         """
         for _, f in obs_filters.items():
             assert getattr(f, "is_concurrent", False), \
@@ -356,6 +367,7 @@ class AsyncSampler(threading.Thread, SamplerInput):
         self.shutdown = False
         self.observation_fn = observation_fn
         self._use_trajectory_view_api = _use_trajectory_view_api
+        self.render = render
         if _use_trajectory_view_api:
             if not sample_collector_class:
                 sample_collector_class = SimpleListCollector
@@ -392,7 +404,7 @@ class AsyncSampler(threading.Thread, SamplerInput):
             self.clip_actions, self.multiple_episodes_in_batch, self.callbacks,
             self.tf_sess, self.perf_stats, self.soft_horizon,
             self.no_done_at_end, self.observation_fn,
-            self._use_trajectory_view_api, self.sample_collector)
+            self._use_trajectory_view_api, self.sample_collector, self.render)
         while not self.shutdown:
             # The timeout variable exists because apparently, if one worker
             # dies, the other workers won't die with it, unless the timeout is
@@ -458,6 +470,7 @@ def _env_runner(
         observation_fn: "ObservationFunction",
         _use_trajectory_view_api: bool = False,
         sample_collector: Optional[SampleCollector] = None,
+        render: bool = None,
 ) -> Iterable[SampleBatchType]:
     """This implements the common experience collection logic.
 
@@ -497,7 +510,9 @@ def _env_runner(
             `_use_trajectory_view_api` to make generic trajectory views
             available to Models. Default: False.
         sample_collector (Optional[SampleCollector]): An optional
-            SampleCollector object to use
+            SampleCollector object to use.
+        render (bool): Whether to try to render the environment after each
+            step.
 
     Yields:
         rollout (SampleBatch): Object containing state, action, reward,
@@ -686,6 +701,12 @@ def _env_runner(
         base_env.send_actions(actions_to_send)
         perf_stats.env_wait_time += time.time() - t4
 
+        # Try to render the env, if required.
+        if render:
+            t5 = time.time()
+            base_env.try_render()
+            perf_stats.env_render_time += time.time() - t5
+
 
 def _process_observations(
         *,

rllib/examples/export/cartpole_dqn_export.py

@@ -3,7 +3,7 @@
 import os
 import ray
 
-from ray.rllib.agents.registry import get_agent_class
+from ray.rllib.agents.registry import get_trainer_class
 from ray.rllib.utils.framework import try_import_tf
 
 tf1, tf, tfv = try_import_tf()
@@ -12,7 +12,7 @@ ray.init(num_cpus=10)
 
 
 def train_and_export(algo_name, num_steps, model_dir, ckpt_dir, prefix):
-    cls = get_agent_class(algo_name)
+    cls = get_trainer_class(algo_name)
     alg = cls(config={}, env="CartPole-v0")
     for _ in range(num_steps):
         alg.train()

rllib/examples/pettingzoo_env.py

@@ -4,7 +4,7 @@ import os
 from supersuit import normalize_obs_v0, dtype_v0, color_reduction_v0
 
 import ray
-from ray.rllib.agents.registry import get_agent_class
+from ray.rllib.agents.registry import get_trainer_class
 from ray.rllib.env import PettingZooEnv
 from pettingzoo.butterfly import pistonball_v1
 
@@ -33,7 +33,7 @@ if __name__ == "__main__":
     num_rollouts = 2
 
     # 1. Gets default training configuration and specifies the POMgame to load.
-    config = deepcopy(get_agent_class(alg_name)._default_config)
+    config = deepcopy(get_trainer_class(alg_name)._default_config)
 
     # 2. Set environment config. This will be passed to
     # the env_creator function via the register env lambda below.
@@ -76,7 +76,7 @@ if __name__ == "__main__":
 
     # 6. Initialize ray and trainer object
     ray.init(num_cpus=num_cpus + 1)
-    trainer = get_agent_class(alg_name)(env="pistonball", config=config)
+    trainer = get_trainer_class(alg_name)(env="pistonball", config=config)
 
     # 7. Train once
     trainer.train()

rllib/examples/rock_paper_scissors_multiagent.py

@@ -14,7 +14,7 @@ import random
 
 from ray import tune
 from ray.rllib.agents.pg import PGTrainer, PGTFPolicy, PGTorchPolicy
-from ray.rllib.agents.registry import get_agent_class
+from ray.rllib.agents.registry import get_trainer_class
 from ray.rllib.examples.env.rock_paper_scissors import RockPaperScissors
 from ray.rllib.examples.policy.rock_paper_scissors_dummies import \
     BeatLastHeuristic, AlwaysSameHeuristic
@@ -87,7 +87,7 @@ def run_heuristic_vs_learned(args, use_lstm=False, trainer="PG"):
         },
         "framework": "torch" if args.torch else "tf",
     }
-    cls = get_agent_class(trainer) if isinstance(trainer, str) else trainer
+    cls = get_trainer_class(trainer) if isinstance(trainer, str) else trainer
     trainer_obj = cls(config=config)
     env = trainer_obj.workers.local_worker().env
     for _ in range(args.stop_iters):

rllib/execution/learner_thread.py

@@ -1,8 +1,7 @@
-from typing import Dict
-import threading
 import copy
-
 from six.moves import queue
+import threading
+from typing import Dict
 
 from ray.rllib.evaluation.metrics import get_learner_stats
 from ray.rllib.execution.minibatch_buffer import MinibatchBuffer
@@ -69,7 +68,10 @@ class LearnerThread(threading.Thread):
 
     def step(self) -> None:
         with self.queue_timer:
-            batch, _ = self.minibatch_buffer.get()
+            try:
+                batch, _ = self.minibatch_buffer.get()
+            except queue.Empty:
+                return
 
         with self.grad_timer:
             fetches = self.local_worker.learn_on_batch(batch)

rllib/rollout.py

@@ -12,24 +12,27 @@ import shelve
 
 import ray
 import ray.cloudpickle as cloudpickle
+from ray.rllib.agents.registry import get_trainer_class
 from ray.rllib.env import MultiAgentEnv
 from ray.rllib.env.base_env import _DUMMY_AGENT_ID
 from ray.rllib.env.env_context import EnvContext
 from ray.rllib.evaluation.worker_set import WorkerSet
 from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
-from ray.rllib.utils.deprecation import deprecation_warning
 from ray.rllib.utils.spaces.space_utils import flatten_to_single_ndarray
 from ray.tune.utils import merge_dicts
 from ray.tune.registry import get_trainable_cls, _global_registry, ENV_CREATOR
 
 EXAMPLE_USAGE = """
-Example Usage via RLlib CLI:
+Example usage via RLlib CLI:
     rllib rollout /tmp/ray/checkpoint_dir/checkpoint-0 --run DQN
     --env CartPole-v0 --steps 1000000 --out rollouts.pkl
 
-Example Usage via executable:
+Example usage via executable:
     ./rollout.py /tmp/ray/checkpoint_dir/checkpoint-0 --run DQN
     --env CartPole-v0 --steps 1000000 --out rollouts.pkl
+
+Example usage w/o checkpoint (for testing purposes):
+    ./rollout.py --run PPO --env CartPole-v0 --episodes 500
 """
 
 # Note: if you use any custom models or envs, register them here first, e.g.:
@@ -42,6 +45,94 @@ Example Usage via executable:
 # register_env("pa_cartpole", lambda _: ParametricActionsCartPole(10))
 
 
+def create_parser(parser_creator=None):
+    parser_creator = parser_creator or argparse.ArgumentParser
+    parser = parser_creator(
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        description="Roll out a reinforcement learning agent "
+        "given a checkpoint.",
+        epilog=EXAMPLE_USAGE)
+
+    parser.add_argument(
+        "checkpoint",
+        type=str,
+        nargs="?",
+        help="(Optional) checkpoint from which to roll out. "
+        "If none given, will use an initial (untrained) Trainer.")
+
+    required_named = parser.add_argument_group("required named arguments")
+    required_named.add_argument(
+        "--run",
+        type=str,
+        required=True,
+        help="The algorithm or model to train. This may refer to the name "
+        "of a built-on algorithm (e.g. RLLib's `DQN` or `PPO`), or a "
+        "user-defined trainable function or class registered in the "
+        "tune registry.")
+    required_named.add_argument(
+        "--env",
+        type=str,
+        help="The environment specifier to use. This could be an openAI gym "
+        "specifier (e.g. `CartPole-v0`) or a full class-path (e.g. "
+        "`ray.rllib.examples.env.simple_corridor.SimpleCorridor`).")
+    parser.add_argument(
+        "--local-mode",
+        action="store_true",
+        help="Run ray in local mode for easier debugging.")
+    parser.add_argument(
+        "--no-render",
+        default=False,
+        action="store_const",
+        const=True,
+        help="Suppress rendering of the environment.")
+    parser.add_argument(
+        "--video-dir",
+        type=str,
+        default=None,
+        help="Specifies the directory into which videos of all episode "
+        "rollouts will be stored.")
+    parser.add_argument(
+        "--steps",
+        default=10000,
+        help="Number of timesteps to roll out. Rollout will also stop if "
+        "`--episodes` limit is reached first. A value of 0 means no "
+        "limitation on the number of timesteps run.")
+    parser.add_argument(
+        "--episodes",
+        default=0,
+        help="Number of complete episodes to roll out. Rollout will also stop "
+        "if `--steps` (timesteps) limit is reached first. A value of 0 means "
+        "no limitation on the number of episodes run.")
+    parser.add_argument("--out", default=None, help="Output filename.")
+    parser.add_argument(
+        "--config",
+        default="{}",
+        type=json.loads,
+        help="Algorithm-specific configuration (e.g. env, hyperparams). "
+        "Gets merged with loaded configuration from checkpoint file and "
+        "`evaluation_config` settings therein.")
+    parser.add_argument(
+        "--save-info",
+        default=False,
+        action="store_true",
+        help="Save the info field generated by the step() method, "
+        "as well as the action, observations, rewards and done fields.")
+    parser.add_argument(
+        "--use-shelve",
+        default=False,
+        action="store_true",
+        help="Save rollouts into a python shelf file (will save each episode "
+        "as it is generated). An output filename must be set using --out.")
+    parser.add_argument(
+        "--track-progress",
+        default=False,
+        action="store_true",
+        help="Write progress to a temporary file (updated "
+        "after each episode). An output filename must be set using --out; "
+        "the progress file will live in the same folder.")
+    return parser
+
+
 class RolloutSaver:
     """Utility class for storing rollouts.
 
@@ -165,108 +256,31 @@ class RolloutSaver:
         self._total_steps += 1
 
 
-def create_parser(parser_creator=None):
-    parser_creator = parser_creator or argparse.ArgumentParser
-    parser = parser_creator(
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        description="Roll out a reinforcement learning agent "
-        "given a checkpoint.",
-        epilog=EXAMPLE_USAGE)
-
-    parser.add_argument(
-        "checkpoint", type=str, help="Checkpoint from which to roll out.")
-    required_named = parser.add_argument_group("required named arguments")
-    required_named.add_argument(
-        "--run",
-        type=str,
-        required=True,
-        help="The algorithm or model to train. This may refer to the name "
-        "of a built-on algorithm (e.g. RLLib's DQN or PPO), or a "
-        "user-defined trainable function or class registered in the "
-        "tune registry.")
-    required_named.add_argument(
-        "--env", type=str, help="The gym environment to use.")
-    parser.add_argument(
-        "--no-render",
-        default=False,
-        action="store_const",
-        const=True,
-        help="Suppress rendering of the environment.")
-    parser.add_argument(
-        "--monitor",
-        default=False,
-        action="store_true",
-        help="Wrap environment in gym Monitor to record video. NOTE: This "
-        "option is deprecated: Use `--video-dir [some dir]` instead.")
-    parser.add_argument(
-        "--video-dir",
-        type=str,
-        default=None,
-        help="Specifies the directory into which videos of all episode "
-        "rollouts will be stored.")
-    parser.add_argument(
-        "--steps",
-        default=10000,
-        help="Number of timesteps to roll out (overwritten by --episodes).")
-    parser.add_argument(
-        "--episodes",
-        default=0,
-        help="Number of complete episodes to roll out (overrides --steps).")
-    parser.add_argument("--out", default=None, help="Output filename.")
-    parser.add_argument(
-        "--config",
-        default="{}",
-        type=json.loads,
-        help="Algorithm-specific configuration (e.g. env, hyperparams). "
-        "Gets merged with loaded configuration from checkpoint file and "
-        "`evaluation_config` settings therein.")
-    parser.add_argument(
-        "--save-info",
-        default=False,
-        action="store_true",
-        help="Save the info field generated by the step() method, "
-        "as well as the action, observations, rewards and done fields.")
-    parser.add_argument(
-        "--use-shelve",
-        default=False,
-        action="store_true",
-        help="Save rollouts into a python shelf file (will save each episode "
-        "as it is generated). An output filename must be set using --out.")
-    parser.add_argument(
-        "--track-progress",
-        default=False,
-        action="store_true",
-        help="Write progress to a temporary file (updated "
-        "after each episode). An output filename must be set using --out; "
-        "the progress file will live in the same folder.")
-    return parser
-
-
 def run(args, parser):
     # Load configuration from checkpoint file.
-    config_dir = os.path.dirname(args.checkpoint)
-    config_path = os.path.join(config_dir, "params.pkl")
-    # Try parent directory.
-    if not os.path.exists(config_path):
-        config_path = os.path.join(config_dir, "../params.pkl")
-
-    # If no pkl file found, require command line `--config`.
-    if not os.path.exists(config_path):
-        if not args.config:
-            raise ValueError(
-                "Could not find params.pkl in either the checkpoint dir or "
-                "its parent directory AND no config given on command line!")
-        else:
-            config = args.config
+    config_path = ""
+    if args.checkpoint:
+        config_dir = os.path.dirname(args.checkpoint)
+        config_path = os.path.join(config_dir, "params.pkl")
+        # Try parent directory.
+        if not os.path.exists(config_path):
+            config_path = os.path.join(config_dir, "../params.pkl")
 
     # Load the config from pickled.
-    else:
+    if os.path.exists(config_path):
         with open(config_path, "rb") as f:
             config = cloudpickle.load(f)
+    # If no pkl file found, require command line `--config`.
+    else:
+        # If no config in given checkpoint -> Error.
+        if args.checkpoint:
+            raise ValueError(
+                "Could not find params.pkl in either the checkpoint dir or "
+                "its parent directory AND no `--config` given on command "
+                "line!")
 
-    # Set num_workers to be at least 2.
-    if "num_workers" in config:
-        config["num_workers"] = min(2, config["num_workers"])
+        # Use default config for given agent.
+        _, config = get_trainer_class(args.run, return_config=True)
 
     # Make sure worker 0 has an Env.
     config["create_env_on_driver"] = True
@@ -285,25 +299,31 @@ def run(args, parser):
             parser.error("the following arguments are required: --env")
         args.env = config.get("env")
 
-    ray.init()
+    # Make sure we have evaluation workers.
+    if not config.get("evaluation_num_workers"):
+        config["evaluation_num_workers"] = config.get("num_workers", 0)
+    if not config.get("evaluation_num_episodes"):
+        config["evaluation_num_episodes"] = 1
+    config["render_env"] = not args.no_render
+    config["record_env"] = args.video_dir
+
+    ray.init(local_mode=args.local_mode)
 
     # Create the Trainer from config.
     cls = get_trainable_cls(args.run)
     agent = cls(env=args.env, config=config)
-    # Load state from checkpoint.
-    agent.restore(args.checkpoint)
+
+    # Load state from checkpoint, if provided.
+    if args.checkpoint:
+        agent.restore(args.checkpoint)
+
     num_steps = int(args.steps)
     num_episodes = int(args.episodes)
 
     # Determine the video output directory.
-    # Deprecated way: Use (--out|~/ray_results) + "/monitor" as dir.
     video_dir = None
-    if args.monitor:
-        video_dir = os.path.join(
-            os.path.dirname(args.out or "")
-            or os.path.expanduser("~/ray_results/"), "monitor")
-    # New way: Allow user to specify a video output path.
-    elif args.video_dir:
+    # Allow user to specify a video output path.
+    if args.video_dir:
         video_dir = os.path.expanduser(args.video_dir)
 
     # Do the actual rollout.
@@ -333,13 +353,13 @@ def default_policy_agent_mapping(unused_agent_id):
 
 def keep_going(steps, num_steps, episodes, num_episodes):
     """Determine whether we've collected enough data"""
-    # if num_episodes is set, this overrides num_steps
-    if num_episodes:
-        return episodes < num_episodes
-    # if num_steps is set, continue until we reach the limit
-    if num_steps:
-        return steps < num_steps
-    # otherwise keep going forever
+    # If num_episodes is set, stop if limit reached.
+    if num_episodes and episodes >= num_episodes:
+        return False
+    # If num_steps is set, stop if limit reached.
+    elif num_steps and steps >= num_steps:
+        return False
+    # Otherwise, keep going.
     return True
 
 
@@ -355,16 +375,36 @@ def rollout(agent,
     if saver is None:
         saver = RolloutSaver()
 
-    if hasattr(agent, "workers") and isinstance(agent.workers, WorkerSet):
+    # Normal case: Agent was setup correctly with an evaluation WorkerSet,
+    # which we will now use to rollout.
+    if hasattr(agent, "evaluation_workers") and isinstance(
+            agent.evaluation_workers, WorkerSet):
+        steps = 0
+        episodes = 0
+        while keep_going(steps, num_steps, episodes, num_episodes):
+            saver.begin_rollout()
+            eval_result = agent._evaluate()["evaluation"]
+            # Increase timestep and episode counters.
+            eps = agent.config["evaluation_num_episodes"]
+            episodes += eps
+            steps += eps * eval_result["episode_len_mean"]
+            # Print out results and continue.
+            print("Episode #{}: reward: {}".format(
+                episodes, eval_result["episode_reward_mean"]))
+            saver.end_rollout()
+        return
+
+    # Agent has no evaluation workers, but RolloutWorkers.
+    elif hasattr(agent, "workers") and isinstance(agent.workers, WorkerSet):
         env = agent.workers.local_worker().env
         multiagent = isinstance(env, MultiAgentEnv)
         if agent.workers.local_worker().multiagent:
             policy_agent_mapping = agent.config["multiagent"][
                 "policy_mapping_fn"]
-
         policy_map = agent.workers.local_worker().policy_map
         state_init = {p: m.get_initial_state() for p, m in policy_map.items()}
         use_lstm = {p: len(s) > 0 for p, s in state_init.items()}
+    # Agent has neither evaluation- nor rollout workers.
     else:
         from gym import envs
         if envs.registry.env_specs.get(agent.config["env"]):
@@ -397,7 +437,7 @@ def rollout(agent,
         env = gym_wrappers.Monitor(
             env=env,
             directory=video_dir,
-            video_callable=lambda x: True,
+            video_callable=lambda _: True,
             force=True)
 
     steps = 0
@@ -470,15 +510,6 @@ if __name__ == "__main__":
     parser = create_parser()
     args = parser.parse_args()
 
-    # Old option: monitor, use video-dir instead.
-    if args.monitor:
-        deprecation_warning("--monitor", "--video-dir=[some dir]")
-    # User tries to record videos, but no-render is set: Error.
-    if (args.monitor or args.video_dir) and args.no_render:
-        raise ValueError(
-            "You have --no-render set, but are trying to record rollout videos"
-            " (via options --video-dir/--monitor)! "
-            "Either unset --no-render or do not use --video-dir/--monitor.")
     # --use_shelve w/o --out option.
     if args.use_shelve and not args.out:
         raise ValueError(

rllib/tests/test_checkpoint_restore.py

@@ -4,7 +4,7 @@ import numpy as np
 import unittest
 
 import ray
-from ray.rllib.agents.registry import get_agent_class
+from ray.rllib.agents.registry import get_trainer_class
 from ray.rllib.utils.test_utils import check, framework_iterator
 
 
@@ -69,7 +69,7 @@ def ckpt_restore_test(alg_name, tfe=False):
     for fw in framework_iterator(config, frameworks=frameworks):
         for use_object_store in [False, True]:
             print("use_object_store={}".format(use_object_store))
-            cls = get_agent_class(alg_name)
+            cls = get_trainer_class(alg_name)
             if "DDPG" in alg_name or "SAC" in alg_name:
                 alg1 = cls(config=config, env="Pendulum-v0")
                 alg2 = cls(config=config, env="Pendulum-v0")

rllib/tests/test_eager_support.py

@@ -2,7 +2,7 @@ import unittest
 
 import ray
 from ray import tune
-from ray.rllib.agents.registry import get_agent_class
+from ray.rllib.agents.registry import get_trainer_class
 from ray.rllib.utils.framework import try_import_tf
 
 tf1, tf, tfv = try_import_tf()
@@ -23,7 +23,7 @@ def check_support(alg, config, test_eager=False, test_trace=True):
         else:
             config["env"] = "CartPole-v0"
 
-        a = get_agent_class(alg)
+        a = get_trainer_class(alg)
         if test_eager:
             print("tf-eager: alg={} cont.act={}".format(alg, cont))
             config["eager_tracing"] = False

rllib/tests/test_export.py

@@ -5,7 +5,7 @@ import shutil
 import unittest
 
 import ray
-from ray.rllib.agents.registry import get_agent_class
+from ray.rllib.agents.registry import get_trainer_class
 from ray.tune.trial import ExportFormat
 
 CONFIGS = {
@@ -74,7 +74,7 @@ def export_test(alg_name, failures):
             and os.path.exists(os.path.join(checkpoint_dir, "model.index")) \
             and os.path.exists(os.path.join(checkpoint_dir, "checkpoint"))
 
-    cls = get_agent_class(alg_name)
+    cls = get_trainer_class(alg_name)
     if "DDPG" in alg_name or "SAC" in alg_name:
         algo = cls(config=CONFIGS[alg_name], env="Pendulum-v0")
     else:

rllib/tests/test_ignore_worker_failure.py

@@ -3,7 +3,7 @@ import unittest
 
 import ray
 from ray.rllib import _register_all
-from ray.rllib.agents.registry import get_agent_class
+from ray.rllib.agents.registry import get_trainer_class
 from ray.rllib.utils.test_utils import framework_iterator
 from ray.tune.registry import register_env
 
@@ -37,7 +37,7 @@ class IgnoresWorkerFailure(unittest.TestCase):
 
     def _do_test_fault_recover(self, alg, config):
         register_env("fault_env", lambda c: FaultInjectEnv(c))
-        agent_cls = get_agent_class(alg)
+        agent_cls = get_trainer_class(alg)
 
         # Test fault handling
         config["num_workers"] = 2
@@ -51,7 +51,7 @@ class IgnoresWorkerFailure(unittest.TestCase):
 
     def _do_test_fault_fatal(self, alg, config):
         register_env("fault_env", lambda c: FaultInjectEnv(c))
-        agent_cls = get_agent_class(alg)
+        agent_cls = get_trainer_class(alg)
         # Test raises real error when out of workers
         config["num_workers"] = 2
         config["ignore_worker_failures"] = True

rllib/tests/test_model_imports.py

@@ -6,7 +6,7 @@ from pathlib import Path
 import unittest
 
 import ray
-from ray.rllib.agents.registry import get_agent_class
+from ray.rllib.agents.registry import get_trainer_class
 from ray.rllib.models.catalog import ModelCatalog
 from ray.rllib.models.tf.misc import normc_initializer
 from ray.rllib.models.tf.tf_modelv2 import TFModelV2
@@ -127,7 +127,7 @@ def model_import_test(algo, config, env):
     rllib_dir = Path(__file__).parent.parent
     import_file = str(rllib_dir) + "/tests/data/model_weights/weights.h5"
 
-    agent_cls = get_agent_class(algo)
+    agent_cls = get_trainer_class(algo)
 
     for fw in framework_iterator(config, ["tf", "torch"]):
         config["model"]["custom_model"] = "keras_model" if fw != "torch" else \

rllib/tests/test_pettingzoo_env.py

@@ -4,7 +4,7 @@ from copy import deepcopy
 import ray
 from ray.tune.registry import register_env
 from ray.rllib.env import PettingZooEnv
-from ray.rllib.agents.registry import get_agent_class
+from ray.rllib.agents.registry import get_trainer_class
 
 from pettingzoo.mpe import simple_spread_v2
 
@@ -20,7 +20,7 @@ class TestPettingZooEnv(unittest.TestCase):
         register_env("simple_spread",
                      lambda _: PettingZooEnv(simple_spread_v2.env()))
 
-        agent_class = get_agent_class("PPO")
+        agent_class = get_trainer_class("PPO")
 
         config = deepcopy(agent_class._default_config)
 

rllib/tests/test_supported_multi_agent.py

@@ -1,7 +1,7 @@
 import unittest
 
 import ray
-from ray.rllib.agents.registry import get_agent_class
+from ray.rllib.agents.registry import get_trainer_class
 from ray.rllib.examples.env.multi_agent import MultiAgentCartPole, \
     MultiAgentMountainCar
 from ray.rllib.utils.test_utils import framework_iterator
@@ -19,10 +19,11 @@ def check_support_multiagent(alg, config):
                 alg in ["A3C", "APEX", "APEX_DDPG", "IMPALA"]:
             continue
         if alg in ["DDPG", "APEX_DDPG", "SAC"]:
-            a = get_agent_class(alg)(
+            a = get_trainer_class(alg)(
                 config=config, env="multi_agent_mountaincar")
         else:
-            a = get_agent_class(alg)(config=config, env="multi_agent_cartpole")
+            a = get_trainer_class(alg)(
+                config=config, env="multi_agent_cartpole")
 
         print(a.train())
         a.stop()

rllib/tests/test_supported_spaces.py

@@ -3,7 +3,7 @@ import numpy as np
 import unittest
 
 import ray
-from ray.rllib.agents.registry import get_agent_class
+from ray.rllib.agents.registry import get_trainer_class
 from ray.rllib.examples.env.random_env import RandomEnv
 from ray.rllib.models.tf.fcnet import FullyConnectedNetwork as FCNetV2
 from ray.rllib.models.tf.visionnet import VisionNetwork as VisionNetV2
@@ -65,7 +65,7 @@ def check_support(alg, config, train=True, check_bounds=False, tfe=False):
         stat = "ok"
 
         try:
-            a = get_agent_class(alg)(config=config, env=RandomEnv)
+            a = get_trainer_class(alg)(config=config, env=RandomEnv)
         except UnsupportedSpaceException:
             stat = "unsupported"
         else:

rllib/train.py

@@ -60,8 +60,7 @@ def create_parser(parser_creator=None):
     parser.add_argument(
         "--local-mode",
         action="store_true",
-        help="Whether to run ray with `local_mode=True`. "
-        "Only if --ray-num-nodes is not used.")
+        help="Run ray in local mode for easier debugging.")
     parser.add_argument(
         "--ray-num-cpus",
         default=None,