From c009c178f6b8adaff9bb8e9cca2d0797d3d14425 Mon Sep 17 00:00:00 2001 From: danuo <66017297+danuo@users.noreply.github.com> Date: Wed, 25 Nov 2020 08:43:17 +0100 Subject: [PATCH] [RLlib] Closes #11924: Add support for custom/ray environments in rollouts.py for agents without workers (#11926) * Closes #11924 Formerly, rollout.py would only load environments from gym (with gym.make() ) , if an agent without workers is employed (such as ES or ARS). This will result in an error, if a custom environment is used. This PR adds the possibility to load environments from the ray registry, while maintaining the support for gym environments. * Update rllib/rollout.py Co-authored-by: Sven Mika --- rllib/rollout.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/rllib/rollout.py b/rllib/rollout.py index fecf65b87..415cd869d 100755 --- a/rllib/rollout.py +++ b/rllib/rollout.py @@ -14,12 +14,13 @@ import ray import ray.cloudpickle as cloudpickle 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 +from ray.tune.registry import get_trainable_cls, _global_registry, ENV_CREATOR EXAMPLE_USAGE = """ Example Usage via RLlib CLI: @@ -365,7 +366,17 @@ def rollout(agent, 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()} else: - env = gym.make(env_name) + from gym import envs + if envs.registry.env_specs.get(agent.config["env"]): + # if environment is gym environment, load from gym + env = gym.make(agent.config["env"]) + else: + # if environment registered ray environment, load from ray + env_creator = _global_registry.get(ENV_CREATOR, + agent.config["env"]) + env_context = EnvContext( + agent.config["env_config"] or {}, worker_index=0) + env = env_creator(env_context) multiagent = False try: policy_map = {DEFAULT_POLICY_ID: agent.policy}