diff --git a/.travis.yml b/.travis.yml index 78175b73e..86934a895 100644 --- a/.travis.yml +++ b/.travis.yml @@ -247,7 +247,7 @@ matrix: - . ./ci/travis/ci.sh preload - ./ci/keep_alive bazel test --config=ci --build_tests_only --test_tag_filters=examples_A,examples_B rllib/... - ./ci/keep_alive bazel test --config=ci --build_tests_only --test_tag_filters=examples_C rllib/... - - ./ci/keep_alive bazel test --config=ci --build_tests_only --test_tag_filters=examples_E,examples_L,examples_M,examples_P rllib/... + - ./ci/keep_alive bazel test --config=ci --build_tests_only --test_tag_filters=examples_E,examples_L,examples_M,examples_N,examples_P rllib/... - ./ci/keep_alive bazel test --config=ci --build_tests_only --test_tag_filters=examples_U,examples_R,examples_S,examples_T rllib/... # RLlib: tests_dir: Everything in rllib/tests/ directory (A-I). diff --git a/rllib/BUILD b/rllib/BUILD index de1064d0a..894f48720 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -1136,10 +1136,11 @@ py_test( ) py_test( - name = "tests/test_nested_spaces", + name = "tests/test_nested_observation_spaces", + main = "tests/test_nested_observation_spaces.py", tags = ["tests_dir", "tests_dir_N"], size = "small", - srcs = ["tests/test_nested_spaces.py"] + srcs = ["tests/test_nested_observation_spaces.py"] ) py_test( @@ -1401,6 +1402,15 @@ py_test( args = ["--num-iters=2"] ) +py_test( + name = "examples/nested_action_spaces_ppo", + main = "examples/nested_action_spaces.py", + tags = ["examples", "examples_N"], + size = "small", + srcs = ["examples/nested_action_spaces.py"], + args = ["--stop=-500", "--run=PPO"] +) + py_test( name = "examples/parametric_action_cartpole_pg", main="examples/parametric_action_cartpole.py", tags = ["examples", "examples_P"], diff --git a/rllib/agents/ars/ars_tf_policy.py b/rllib/agents/ars/ars_tf_policy.py index b09d8b081..3f8292cdf 100644 --- a/rllib/agents/ars/ars_tf_policy.py +++ b/rllib/agents/ars/ars_tf_policy.py @@ -7,7 +7,7 @@ import numpy as np import ray import ray.experimental.tf_utils from ray.rllib.agents.es.es_tf_policy import make_session -from ray.rllib.evaluation.sampler import _unbatch_tuple_actions +from ray.rllib.evaluation.sampler import unbatch_actions from ray.rllib.models import ModelCatalog from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.utils.filter import get_filter @@ -56,7 +56,7 @@ class ARSTFPolicy: observation = self.observation_filter(observation[None], update=update) action = self.sess.run( self.sampler, feed_dict={self.inputs: observation}) - action = _unbatch_tuple_actions(action) + action = unbatch_actions(action) if add_noise and isinstance(self.action_space, gym.spaces.Box): action += np.random.randn(*action.shape) * self.action_noise_std return action diff --git a/rllib/agents/es/es_tf_policy.py b/rllib/agents/es/es_tf_policy.py index 03a5297bf..40e33f3c5 100644 --- a/rllib/agents/es/es_tf_policy.py +++ b/rllib/agents/es/es_tf_policy.py @@ -6,13 +6,16 @@ import numpy as np import ray import ray.experimental.tf_utils -from ray.rllib.evaluation.sampler import _unbatch_tuple_actions +from ray.rllib.evaluation.sampler import unbatch_actions from ray.rllib.models import ModelCatalog from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils import try_import_tree from ray.rllib.utils.filter import get_filter from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.space_utils import get_base_struct_from_space tf = try_import_tf() +tree = try_import_tree() def rollout(policy, env, timestep_limit=None, add_noise=False, offset=0.0): @@ -66,6 +69,7 @@ def make_session(single_threaded): class ESTFPolicy: def __init__(self, obs_space, action_space, config): self.action_space = action_space + self.action_space_struct = get_base_struct_from_space(action_space) self.action_noise_std = config["action_noise_std"] self.preprocessor = ModelCatalog.get_preprocessor_for_space(obs_space) self.observation_filter = get_filter(config["observation_filter"], @@ -95,12 +99,22 @@ class ESTFPolicy: def compute_actions(self, observation, add_noise=False, update=True): observation = self.preprocessor.transform(observation) observation = self.observation_filter(observation[None], update=update) - action = self.sess.run( + # `actions` is a list of (component) batches. + actions = self.sess.run( self.sampler, feed_dict={self.inputs: observation}) - action = _unbatch_tuple_actions(action) - if add_noise and isinstance(self.action_space, gym.spaces.Box): - action += np.random.randn(*action.shape) * self.action_noise_std - return action + if add_noise: + actions = tree.map_structure(self._add_noise, actions, + self.action_space_struct) + # Convert `flat_actions` to a list of lists of action components + # (list of single actions). + actions = unbatch_actions(actions) + return actions + + def _add_noise(self, single_action, single_action_space): + if isinstance(single_action_space, gym.spaces.Box): + single_action += np.random.randn(*single_action.shape) * \ + self.action_noise_std + return single_action def set_flat_weights(self, x): self.variables.set_flat(x) diff --git a/rllib/agents/es/es_torch_policy.py b/rllib/agents/es/es_torch_policy.py index 00f01365c..34bcd10bc 100644 --- a/rllib/agents/es/es_torch_policy.py +++ b/rllib/agents/es/es_torch_policy.py @@ -5,7 +5,7 @@ import gym import numpy as np import ray -from ray.rllib.evaluation.sampler import _unbatch_tuple_actions +from ray.rllib.evaluation.sampler import unbatch_actions from ray.rllib.models import ModelCatalog from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.policy.torch_policy_template import build_torch_policy @@ -61,7 +61,7 @@ def before_init(policy, observation_space, action_space, config): }, [], None) dist = policy.dist_class(dist_inputs, policy.model) action = dist.sample().detach().numpy() - action = _unbatch_tuple_actions(action) + action = unbatch_actions(action) if add_noise and isinstance(policy.action_space, gym.spaces.Box): action += np.random.randn(*action.shape) * policy.action_noise_std return action diff --git a/rllib/agents/qmix/qmix_policy.py b/rllib/agents/qmix/qmix_policy.py index 13430a88b..d8c254ef0 100644 --- a/rllib/agents/qmix/qmix_policy.py +++ b/rllib/agents/qmix/qmix_policy.py @@ -14,7 +14,6 @@ from ray.rllib.models.model import _unpack_obs from ray.rllib.env.constants import GROUP_REWARDS from ray.rllib.utils.framework import try_import_torch from ray.rllib.utils.annotations import override -from ray.rllib.utils.tuple_actions import TupleActions # Torch must be installed. torch, nn = try_import_torch(error=True) @@ -289,7 +288,7 @@ class QMixTorchPolicy(Policy): actions = actions.cpu().numpy() hiddens = [s.cpu().numpy() for s in hiddens] - return TupleActions(list(actions.transpose([1, 0]))), hiddens, {} + return tuple(actions.transpose([1, 0])), hiddens, {} @override(Policy) def compute_log_likelihoods(self, diff --git a/rllib/evaluation/episode.py b/rllib/evaluation/episode.py index 22fb405c1..2f3d37379 100644 --- a/rllib/evaluation/episode.py +++ b/rllib/evaluation/episode.py @@ -1,12 +1,14 @@ from collections import defaultdict +import numpy as np import random -import numpy as np - from ray.rllib.env.base_env import _DUMMY_AGENT_ID +from ray.rllib.utils import try_import_tree from ray.rllib.utils.annotations import DeveloperAPI from ray.rllib.utils.space_utils import flatten_to_single_ndarray +tree = try_import_tree() + @DeveloperAPI class MultiAgentEpisode: diff --git a/rllib/evaluation/sampler.py b/rllib/evaluation/sampler.py index 373f61551..117888605 100644 --- a/rllib/evaluation/sampler.py +++ b/rllib/evaluation/sampler.py @@ -15,11 +15,13 @@ from ray.rllib.policy.tf_policy import TFPolicy from ray.rllib.env.base_env import BaseEnv, ASYNC_RESET_RETURN from ray.rllib.env.atari_wrappers import get_wrapper_by_cls, MonitorEnv from ray.rllib.offline import InputReader +from ray.rllib.utils import try_import_tree from ray.rllib.utils.annotations import override from ray.rllib.utils.debug import summarize -from ray.rllib.utils.tuple_actions import TupleActions -from ray.rllib.utils.space_utils import flatten_to_single_ndarray from ray.rllib.utils.tf_run_builder import TFRunBuilder +from ray.rllib.utils.space_utils import flatten_to_single_ndarray + +tree = try_import_tree() logger = logging.getLogger(__name__) @@ -581,9 +583,7 @@ def _do_policy_eval(tf_sess, to_eval, policies, active_episodes): obs_batch = [t.obs for t in eval_data] state_batches = _to_column_format(rnn_in) - # TODO(ekl): how can we make info batch available to TF code? - obs_batch = [t.obs for t in eval_data] prev_action_batch = [t.prev_action for t in eval_data] prev_reward_batch = [t.prev_reward for t in eval_data] @@ -642,6 +642,11 @@ def _process_policy_eval_results(to_eval, eval_results, active_episodes, rnn_out_cols = eval_results[policy_id][1] pi_info_cols = eval_results[policy_id][2] + # In case actions is a list (representing the 0th dim of a batch of + # primitive actions), try to convert it first. + if isinstance(actions, list): + actions = np.array(actions) + if len(rnn_in_cols) != len(rnn_out_cols): raise ValueError("Length of RNN in did not match RNN out, got: " "{} vs {}".format(rnn_in_cols, rnn_out_cols)) @@ -650,17 +655,17 @@ def _process_policy_eval_results(to_eval, eval_results, active_episodes, pi_info_cols["state_in_{}".format(f_i)] = column for f_i, column in enumerate(rnn_out_cols): pi_info_cols["state_out_{}".format(f_i)] = column - # Save output rows - actions = _unbatch_tuple_actions(actions) + policy = _get_or_raise(policies, policy_id) + # Clip if necessary (while action components are still batched). + if clip_actions: + actions = clip_action(actions, policy.action_space_struct) + # Split action-component batches into single action rows. + actions = unbatch_actions(actions) for i, action in enumerate(actions): env_id = eval_data[i].env_id agent_id = eval_data[i].agent_id - if clip_actions: - actions_to_send[env_id][agent_id] = clip_action( - action, policy.action_space) - else: - actions_to_send[env_id][agent_id] = action + actions_to_send[env_id][agent_id] = action episode = active_episodes[env_id] episode._set_rnn_state(agent_id, [c[i] for c in rnn_out_cols]) episode._set_last_pi_info( @@ -694,17 +699,41 @@ def _fetch_atari_metrics(base_env): return atari_out -def _unbatch_tuple_actions(action_batch): - # convert list of batches -> batch of lists - if isinstance(action_batch, TupleActions): - out = [] - for j in range(len(action_batch.batches[0])): - out.append([ - action_batch.batches[i][j] - for i in range(len(action_batch.batches)) - ]) - return out - return action_batch +def unbatch_actions(action_batches): + """Converts action_batches from list of batches to batch of lists. + + Input: Struct of batches: + {"a": [1, 2, 3], "b": ([4, 5, 6], [7.0, 8.0, 9.0])} + Output: Batch (list) of structs (each of these structs representing a + single action): + [ + {"a": 1, "b": (4, 7.0)}, <- action 1 + {"a": 2, "b": (5, 8.0)}, <- action 2 + {"a": 3, "b": (6, 9.0)}, <- action 3 + ] + + Args: + action_batches (any): The list of action-component batches. Each item + in this list represents the batch for a single action component + (in case action is Tuple/Dict), meaning the list is already + flattened. + Alternatively, `action_batches` may also simply be a batch of + primitive actions (non Tuple/Dict). + + Returns: + List[List[action-components]]: The list of action rows. Each item + in the returned list represents a single (maybe complex) action. + """ + flat_action_batches = tree.flatten(action_batches) + + out = [] + for batch_pos in range(len(flat_action_batches[0])): + out.append( + tree.unflatten_as(action_batches, [ + flat_action_batches[i][batch_pos] + for i in range(len(flat_action_batches)) + ])) + return out def _to_column_format(rnn_state_rows): diff --git a/rllib/examples/autoregressive_action_dist.py b/rllib/examples/autoregressive_action_dist.py index 3a720e6fc..8a529771a 100644 --- a/rllib/examples/autoregressive_action_dist.py +++ b/rllib/examples/autoregressive_action_dist.py @@ -21,8 +21,7 @@ from ray.rllib.models import ModelCatalog from ray.rllib.models.tf.tf_action_dist import Categorical, ActionDistribution from ray.rllib.models.tf.misc import normc_initializer from ray.rllib.models.tf.tf_modelv2 import TFModelV2 -from ray.rllib.utils.tuple_actions import TupleActions -from ray.rllib.utils import try_import_tf +from ray.rllib.utils.framework import try_import_tf tf = try_import_tf() @@ -78,7 +77,7 @@ class BinaryAutoregressiveOutput(ActionDistribution): self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2) # return the action tuple - return TupleActions([a1, a2]) + return (a1, a2) def sample(self): # first, sample a1 @@ -91,7 +90,7 @@ class BinaryAutoregressiveOutput(ActionDistribution): self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2) # return the action tuple - return TupleActions([a1, a2]) + return (a1, a2) def logp(self, actions): a1, a2 = actions[:, 0], actions[:, 1] diff --git a/rllib/examples/nested_action_spaces.py b/rllib/examples/nested_action_spaces.py new file mode 100644 index 000000000..acaa66203 --- /dev/null +++ b/rllib/examples/nested_action_spaces.py @@ -0,0 +1,97 @@ +import argparse +import gym +from gym.spaces import Dict, Tuple, Box, Discrete +import numpy as np +import sys + +import ray +from ray.tune.registry import register_env +from ray.rllib.utils import try_import_tree +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.space_utils import flatten_space + +tf = try_import_tf() +tree = try_import_tree() + +parser = argparse.ArgumentParser() +parser.add_argument("--run", type=str, default="PPO") +parser.add_argument("--stop", type=int, default=90) +parser.add_argument("--max-trainstop", type=int, default=90) +parser.add_argument("--num-cpus", type=int, default=0) + + +class NestedSpaceRepeatAfterMeEnv(gym.Env): + """Env for which policy has to repeat the (possibly complex) observation. + """ + + def __init__(self, config): + self.observation_space = config.get( + "space", Tuple([Discrete(2), + Dict({ + "a": Box(-1.0, 1.0, (2, )) + })])) + self.action_space = self.observation_space + self.flattened_action_space = flatten_space(self.action_space) + self.episode_len = config.get("episode_len", 100) + + def reset(self): + self.steps = 0 + return self._next_obs() + + def step(self, action): + self.steps += 1 + action = tree.flatten(action) + reward = 0.0 + for a, o, space in zip(action, self.current_obs_flattened, + self.flattened_action_space): + # Box: -abs(diff). + if isinstance(space, gym.spaces.Box): + reward -= np.abs(np.sum(a - o)) + # Discrete: +1.0 if exact match. + if isinstance(space, gym.spaces.Discrete): + reward += 1.0 if a == o else 0.0 + done = self.steps >= self.episode_len + return self._next_obs(), reward, done, {} + + def _next_obs(self): + self.current_obs = self.observation_space.sample() + self.current_obs_flattened = tree.flatten(self.current_obs) + return self.current_obs + + +if __name__ == "__main__": + args = parser.parse_args() + ray.init(num_cpus=args.num_cpus or None) + register_env("NestedSpaceRepeatAfterMeEnv", + lambda c: NestedSpaceRepeatAfterMeEnv(c)) + + config = { + "env": "NestedSpaceRepeatAfterMeEnv", + "env_config": { + "space": Dict({ + "a": Tuple( + [Dict({ + "d": Box(-10.0, 10.0, ()), + "e": Discrete(2) + })]), + "b": Box(-10.0, 10.0, (2, )), + "c": Discrete(4) + }), + }, + "gamma": 0.0, # No history in Env (bandit problem). + "num_workers": 0, + "num_envs_per_worker": 20, + "entropy_coeff": 0.00005, # We don't want high entropy in this Env. + "num_sgd_iter": 20, + "vf_loss_coeff": 0.01, + "lr": 0.0003 + } + + import ray.rllib.agents.ppo as ppo + trainer = ppo.PPOTrainer(config=config) + for _ in range(100): + results = trainer.train() + print(results) + if results["episode_reward_mean"] > args.stop: + sys.exit(0) # Learnt, exit gracefully. + sys.exit(1) # Done, but did not learn, exit with error. diff --git a/rllib/examples/rock_paper_scissors_multiagent.py b/rllib/examples/rock_paper_scissors_multiagent.py index c942bbf82..6c6804541 100644 --- a/rllib/examples/rock_paper_scissors_multiagent.py +++ b/rllib/examples/rock_paper_scissors_multiagent.py @@ -97,7 +97,7 @@ class AlwaysSameHeuristic(Policy): info_batch=None, episodes=None, **kwargs): - return list(state_batches[0]), state_batches, {} + return state_batches[0], state_batches, {} def learn_on_batch(self, samples): pass @@ -168,32 +168,29 @@ def run_heuristic_vs_learned(args, use_lstm=False, trainer="PG"): else: return random.choice(["always_same", "beat_last"]) - tune.run( - trainer, - stop={"timesteps_total": args.stop}, - config={ - "env": RockPaperScissorsEnv, - "gamma": 0.9, - "num_workers": 0, - "num_envs_per_worker": 4, - "rollout_fragment_length": 10, - "train_batch_size": 200, - "multiagent": { - "policies_to_train": ["learned"], - "policies": { - "always_same": (AlwaysSameHeuristic, Discrete(3), - Discrete(3), {}), - "beat_last": (BeatLastHeuristic, Discrete(3), Discrete(3), - {}), - "learned": (None, Discrete(3), Discrete(3), { - "model": { - "use_lstm": use_lstm - } - }), - }, - "policy_mapping_fn": select_policy, + config = { + "env": RockPaperScissorsEnv, + "gamma": 0.9, + "num_workers": 0, + "num_envs_per_worker": 4, + "rollout_fragment_length": 10, + "train_batch_size": 200, + "multiagent": { + "policies_to_train": ["learned"], + "policies": { + "always_same": (AlwaysSameHeuristic, Discrete(3), Discrete(3), + {}), + "beat_last": (BeatLastHeuristic, Discrete(3), Discrete(3), {}), + "learned": (None, Discrete(3), Discrete(3), { + "model": { + "use_lstm": use_lstm + } + }), }, - }) + "policy_mapping_fn": select_policy, + }, + } + tune.run(trainer, stop={"timesteps_total": args.stop}, config=config) def run_with_custom_entropy_loss(args): @@ -217,11 +214,11 @@ def run_with_custom_entropy_loss(args): if __name__ == "__main__": args = parser.parse_args() + run_heuristic_vs_learned(args, use_lstm=False) + print("run_heuristic_vs_learned(w/o lstm): ok.") run_same_policy(args) print("run_same_policy: ok.") run_heuristic_vs_learned(args, use_lstm=True) - print("run_heuristic_vs_learned(w/ lstm): ok.") - run_heuristic_vs_learned(args, use_lstm=False) - print("run_heuristic_vs_learned (w/o lstm): ok.") + print("run_heuristic_vs_learned (w/ lstm): ok.") run_with_custom_entropy_loss(args) print("run_with_custom_entropy_loss: ok.") diff --git a/rllib/examples/rollout_worker_custom_workflow.py b/rllib/examples/rollout_worker_custom_workflow.py index 167722cef..f6a323317 100644 --- a/rllib/examples/rollout_worker_custom_workflow.py +++ b/rllib/examples/rollout_worker_custom_workflow.py @@ -7,6 +7,7 @@ collection and policy optimization. import argparse import gym +import numpy as np import ray from ray import tune @@ -43,7 +44,8 @@ class CustomPolicy(TestPolicy): episodes=None, **kwargs): # return random actions - return [self.action_space.sample() for _ in obs_batch], [], {} + return np.array([self.action_space.sample() + for _ in obs_batch]), [], {} def learn_on_batch(self, samples): # implement your learning code here diff --git a/rllib/models/catalog.py b/rllib/models/catalog.py index 09110a845..289542f9c 100644 --- a/rllib/models/catalog.py +++ b/rllib/models/catalog.py @@ -12,18 +12,22 @@ from ray.rllib.models.preprocessors import get_preprocessor from ray.rllib.models.tf.fcnet_v1 import FullyConnectedNetwork from ray.rllib.models.tf.lstm_v1 import LSTM from ray.rllib.models.tf.modelv1_compat import make_v1_wrapper -from ray.rllib.models.tf.tf_action_dist import Categorical, MultiCategorical, \ - Deterministic, DiagGaussian, MultiActionDistribution, Dirichlet +from ray.rllib.models.tf.tf_action_dist import Categorical, \ + Deterministic, DiagGaussian, Dirichlet, \ + MultiActionDistribution, MultiCategorical from ray.rllib.models.tf.tf_modelv2 import TFModelV2 from ray.rllib.models.tf.visionnet_v1 import VisionNetwork from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 from ray.rllib.models.torch.torch_action_dist import TorchCategorical, \ - TorchMultiCategorical, TorchDeterministic, TorchDiagGaussian -from ray.rllib.utils import try_import_tf + TorchDeterministic, TorchDiagGaussian, \ + TorchMultiActionDistribution, TorchMultiCategorical +from ray.rllib.utils import try_import_tf, try_import_tree from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI from ray.rllib.utils.error import UnsupportedSpaceException +from ray.rllib.utils.space_utils import flatten_space tf = try_import_tf() +tree = try_import_tree() logger = logging.getLogger(__name__) @@ -134,7 +138,8 @@ class ModelCatalog: # Dist_type is given directly as a class. elif type(dist_type) is type and \ issubclass(dist_type, ActionDistribution) and \ - dist_type is not MultiActionDistribution: + dist_type not in ( + MultiActionDistribution, TorchMultiActionDistribution): dist = dist_type # Box space -> DiagGaussian OR Deterministic. elif isinstance(action_space, gym.spaces.Box): @@ -154,24 +159,21 @@ class ModelCatalog: # Discrete Space -> Categorical. elif isinstance(action_space, gym.spaces.Discrete): dist = Categorical if framework == "tf" else TorchCategorical - # Tuple Space -> MultiAction. - elif dist_type is MultiActionDistribution or \ - isinstance(action_space, gym.spaces.Tuple): - if framework == "torch": - # TODO(sven): implement - raise NotImplementedError( - "Tuple action spaces not supported for Pytorch.") - child_dist = [] - input_lens = [] - for action in action_space.spaces: - dist, action_size = ModelCatalog.get_action_dist( - action, config) - child_dist.append(dist) - input_lens.append(action_size) + # Tuple/Dict Spaces -> MultiAction. + elif dist_type in (MultiActionDistribution, + TorchMultiActionDistribution) or \ + isinstance(action_space, (gym.spaces.Tuple, gym.spaces.Dict)): + flat_action_space = flatten_space(action_space) + child_dists_and_in_lens = tree.map_structure( + lambda s: ModelCatalog.get_action_dist( + s, config, framework=framework), flat_action_space) + child_dists = [e[0] for e in child_dists_and_in_lens] + input_lens = [e[1] for e in child_dists_and_in_lens] return partial( - MultiActionDistribution, + (TorchMultiActionDistribution + if framework == "torch" else MultiActionDistribution), action_space=action_space, - child_distributions=child_dist, + child_distributions=child_dists, input_lens=input_lens), sum(input_lens) # Simplex -> Dirichlet. elif isinstance(action_space, Simplex): @@ -186,12 +188,6 @@ class ModelCatalog: TorchMultiCategorical return partial(dist, input_lens=action_space.nvec), \ int(sum(action_space.nvec)) - # Dict -> TODO(sven) - elif isinstance(action_space, gym.spaces.Dict): - # TODO(sven): implement - raise NotImplementedError( - "Dict action spaces are not supported, consider using " - "gym.spaces.Tuple instead") # Unknown type -> Error. else: raise NotImplementedError("Unsupported args: {} {}".format( @@ -217,39 +213,38 @@ class ModelCatalog: elif isinstance(action_space, gym.spaces.MultiDiscrete): return (tf.as_dtype(action_space.dtype), (None, ) + action_space.shape) - elif isinstance(action_space, gym.spaces.Tuple): + elif isinstance(action_space, (gym.spaces.Tuple, gym.spaces.Dict)): + flat_action_space = flatten_space(action_space) size = 0 all_discrete = True - for i in range(len(action_space.spaces)): - if isinstance(action_space.spaces[i], gym.spaces.Discrete): + for i in range(len(flat_action_space)): + if isinstance(flat_action_space[i], gym.spaces.Discrete): size += 1 else: all_discrete = False - size += np.product(action_space.spaces[i].shape) + size += np.product(flat_action_space[i].shape) size = int(size) return (tf.int64 if all_discrete else tf.float32, (None, size)) - elif isinstance(action_space, gym.spaces.Dict): - raise NotImplementedError( - "Dict action spaces are not supported, consider using " - "gym.spaces.Tuple instead") else: - raise NotImplementedError("action space {}" - " not supported".format(action_space)) + raise NotImplementedError( + "Action space {} not supported".format(action_space)) @staticmethod @DeveloperAPI - def get_action_placeholder(action_space, name=None): + def get_action_placeholder(action_space, name="action"): """Returns an action placeholder consistent with the action space Args: action_space (Space): Action space of the target gym env. + name (str): An optional string to name the placeholder by. + Default: "action". Returns: action_placeholder (Tensor): A placeholder for the actions """ dtype, shape = ModelCatalog.get_action_shape(action_space) - return tf.placeholder(dtype, shape=shape, name=(name or "action")) + return tf.placeholder(dtype, shape=shape, name=name) @staticmethod @DeveloperAPI diff --git a/rllib/models/tf/tf_action_dist.py b/rllib/models/tf/tf_action_dist.py index 8d5fd1d63..8197b7f6e 100644 --- a/rllib/models/tf/tf_action_dist.py +++ b/rllib/models/tf/tf_action_dist.py @@ -2,13 +2,15 @@ import numpy as np import functools from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.utils import MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT, \ + SMALL_NUMBER, try_import_tree from ray.rllib.utils.annotations import override, DeveloperAPI -from ray.rllib.utils import try_import_tf, try_import_tfp, SMALL_NUMBER, \ - MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT -from ray.rllib.utils.tuple_actions import TupleActions +from ray.rllib.utils.framework import try_import_tf, try_import_tfp +from ray.rllib.utils.space_utils import get_base_struct_from_space tf = try_import_tf() tfp = try_import_tfp() +tree = try_import_tree() @DeveloperAPI @@ -350,70 +352,82 @@ class Deterministic(TFActionDistribution): class MultiActionDistribution(TFActionDistribution): - """Action distribution that operates for list of actions. + """Action distribution that operates on a set of actions. Args: inputs (Tensor list): A list of tensors from which to compute samples. """ - def __init__(self, inputs, model, action_space, child_distributions, - input_lens): - # skip TFActionDistribution init + def __init__(self, inputs, model, *, child_distributions, input_lens, + action_space): ActionDistribution.__init__(self, inputs, model) - self.input_lens = input_lens - split_inputs = tf.split(inputs, self.input_lens, axis=1) - child_list = [] - for i, distribution in enumerate(child_distributions): - child_list.append(distribution(split_inputs[i], model)) - self.child_distributions = child_list + + self.action_space_struct = get_base_struct_from_space(action_space) + + input_lens = np.array(input_lens, dtype=np.int32) + split_inputs = tf.split(inputs, input_lens, axis=1) + self.flat_child_distributions = tree.map_structure( + lambda dist, input_: dist(input_, model), child_distributions, + split_inputs) @override(ActionDistribution) def logp(self, x): - split_indices = [] - for dist in self.child_distributions: + # Single tensor input (all merged). + if isinstance(x, (tf.Tensor, np.ndarray)): + split_indices = [] + for dist in self.flat_child_distributions: + if isinstance(dist, Categorical): + split_indices.append(1) + else: + split_indices.append(tf.shape(dist.sample())[1]) + split_x = tf.split(x, split_indices, axis=1) + # Structured or flattened (by single action component) input. + else: + split_x = tree.flatten(x) + + def map_(val, dist): + # Remove extra categorical dimension. if isinstance(dist, Categorical): - split_indices.append(1) - else: - split_indices.append(tf.shape(dist.sample())[1]) - split_list = tf.split(x, split_indices, axis=1) - for i, distribution in enumerate(self.child_distributions): - # Remove extra categorical dimension - if isinstance(distribution, Categorical): - split_list[i] = tf.cast( - tf.squeeze(split_list[i], axis=-1), tf.int32) - log_list = [ - distribution.logp(split_x) for distribution, split_x in zip( - self.child_distributions, split_list) - ] - return functools.reduce(lambda a, b: a + b, log_list) + val = tf.cast(tf.squeeze(val, axis=-1), tf.int32) + return dist.logp(val) + + # Remove extra categorical dimension and take the logp of each + # component. + flat_logps = tree.map_structure(map_, split_x, + self.flat_child_distributions) + + return functools.reduce(lambda a, b: a + b, flat_logps) @override(ActionDistribution) def kl(self, other): kl_list = [ - distribution.kl(other_distribution) - for distribution, other_distribution in zip( - self.child_distributions, other.child_distributions) + d.kl(o) for d, o in zip(self.flat_child_distributions, + other.flat_child_distributions) ] return functools.reduce(lambda a, b: a + b, kl_list) @override(ActionDistribution) def entropy(self): - entropy_list = [s.entropy() for s in self.child_distributions] + entropy_list = [d.entropy() for d in self.flat_child_distributions] return functools.reduce(lambda a, b: a + b, entropy_list) @override(ActionDistribution) def sample(self): - return TupleActions([s.sample() for s in self.child_distributions]) + child_distributions = tree.unflatten_as(self.action_space_struct, + self.flat_child_distributions) + return tree.map_structure(lambda s: s.sample(), child_distributions) @override(ActionDistribution) def deterministic_sample(self): - return TupleActions( - [s.deterministic_sample() for s in self.child_distributions]) + child_distributions = tree.unflatten_as(self.action_space_struct, + self.flat_child_distributions) + return tree.map_structure(lambda s: s.deterministic_sample(), + child_distributions) @override(TFActionDistribution) def sampled_action_logp(self): - p = self.child_distributions[0].sampled_action_logp() - for c in self.child_distributions[1:]: + p = self.flat_child_distributions[0].sampled_action_logp() + for c in self.flat_child_distributions[1:]: p += c.sampled_action_logp() return p diff --git a/rllib/policy/eager_tf_policy.py b/rllib/policy/eager_tf_policy.py index 479befa9d..62a5ecf30 100644 --- a/rllib/policy/eager_tf_policy.py +++ b/rllib/policy/eager_tf_policy.py @@ -580,8 +580,6 @@ def build_eager_tf_policy(name, def _initialize_loss_with_dummy_batch(self): # Dummy forward pass to initialize any policy attributes, etc. - action_dtype, action_shape = ModelCatalog.get_action_shape( - self.action_space) dummy_batch = { SampleBatch.CUR_OBS: np.array( [self.observation_space.sample()]), diff --git a/rllib/policy/policy.py b/rllib/policy/policy.py index 09fc96b5b..3b1df5260 100644 --- a/rllib/policy/policy.py +++ b/rllib/policy/policy.py @@ -2,9 +2,13 @@ from abc import ABCMeta, abstractmethod import gym import numpy as np +from ray.rllib.utils import try_import_tree from ray.rllib.utils.annotations import DeveloperAPI from ray.rllib.utils.exploration.exploration import Exploration from ray.rllib.utils.from_config import from_config +from ray.rllib.utils.space_utils import get_base_struct_from_space + +tree = try_import_tree() # By convention, metrics from optimizing the loss can be reported in the # `grad_info` dict returned by learn_on_batch() / compute_grads() via this key. @@ -47,6 +51,7 @@ class Policy(metaclass=ABCMeta): """ self.observation_space = observation_space self.action_space = action_space + self.action_space_struct = get_base_struct_from_space(action_space) self.config = config # The global timestep, broadcast down from time to time from the # driver. @@ -157,7 +162,7 @@ class Policy(metaclass=ABCMeta): timestep=timestep) if clip_actions: - action = clip_action(action, self.action_space) + action = clip_action(action, self.action_space_struct) # Return action, internal state(s), infos. return action, [s[0] for s in state_out], \ @@ -385,27 +390,22 @@ class Policy(metaclass=ABCMeta): return exploration -def clip_action(action, space): - """ - Called to clip actions to the specified range of this policy. +def clip_action(action, action_space): + """Clips all actions in `flat_actions` according to the given Spaces. - Arguments: - action: Single action. - space: Action space the actions should be present in. + Args: + flat_actions (List[np.ndarray]): The (flattened) list of single action + components. List will have len=1 for "primitive" action Spaces. + flat_space (List[Space]): The (flattened) list of single action Space + objects. Has to be of same length as `flat_actions`. Returns: - Clipped batch of actions. + List[np.ndarray]: Flattened list of single clipped "primitive" actions. """ - if isinstance(space, gym.spaces.Box): - return np.clip(action, space.low, space.high) - elif isinstance(space, gym.spaces.Tuple): - if type(action) not in (tuple, list): - raise ValueError("Expected tuple space for actions {}: {}".format( - action, space)) - out = [] - for a, s in zip(action, space.spaces): - out.append(clip_action(a, s)) - return out - else: - return action + def map_(a, s): + if isinstance(s, gym.spaces.Box): + a = np.clip(a, s.low, s.high) + return a + + return tree.map_structure(map_, action, action_space) diff --git a/rllib/policy/tests/test_policy.py b/rllib/policy/tests/test_policy.py index 58c0fc0a0..89f68d5dc 100644 --- a/rllib/policy/tests/test_policy.py +++ b/rllib/policy/tests/test_policy.py @@ -1,3 +1,4 @@ +import numpy as np import random from ray.rllib.policy.policy import Policy @@ -22,7 +23,7 @@ class TestPolicy(Policy): explore=None, timestep=None, **kwargs): - return [random.choice([0, 1])] * len(obs_batch), [], {} + return np.array([random.choice([0, 1])] * len(obs_batch)), [], {} @override(Policy) def compute_log_likelihoods(self, @@ -31,4 +32,4 @@ class TestPolicy(Policy): state_batches=None, prev_action_batch=None, prev_reward_batch=None): - return [random.random()] * len(obs_batch) + return np.array([random.random()] * len(obs_batch)) diff --git a/rllib/tests/test_multi_agent_pendulum.py b/rllib/tests/test_multi_agent_pendulum.py index 9e27af4b0..5ba73ee06 100644 --- a/rllib/tests/test_multi_agent_pendulum.py +++ b/rllib/tests/test_multi_agent_pendulum.py @@ -36,7 +36,7 @@ class TestMultiAgentPendulum(unittest.TestCase): "sgd_minibatch_size": 64, "num_sgd_iter": 10, "model": { - "fcnet_hiddens": [64, 64], + "fcnet_hiddens": [128, 128], }, "batch_mode": "complete_episodes", }, diff --git a/rllib/tests/test_nested_spaces.py b/rllib/tests/test_nested_observation_spaces.py similarity index 100% rename from rllib/tests/test_nested_spaces.py rename to rllib/tests/test_nested_observation_spaces.py diff --git a/rllib/tests/test_rollout_worker.py b/rllib/tests/test_rollout_worker.py index 9b73c43ee..78d38da82 100644 --- a/rllib/tests/test_rollout_worker.py +++ b/rllib/tests/test_rollout_worker.py @@ -29,7 +29,7 @@ class MockPolicy(TestPolicy): explore=None, timestep=None, **kwargs): - return [random.choice([0, 1])] * len(obs_batch), [], {} + return np.array([random.choice([0, 1])] * len(obs_batch)), [], {} def postprocess_trajectory(self, batch, diff --git a/rllib/tests/test_supported_spaces.py b/rllib/tests/test_supported_spaces.py index 6d8871a1b..ccafffce7 100644 --- a/rllib/tests/test_supported_spaces.py +++ b/rllib/tests/test_supported_spaces.py @@ -1,5 +1,5 @@ import gym -from gym.spaces import Box, Discrete, Tuple, Dict, MultiDiscrete +from gym.spaces import Box, Dict, Discrete, Tuple, MultiDiscrete from gym.envs.registration import EnvSpec import numpy as np import unittest @@ -31,6 +31,13 @@ ACTION_SPACES_TO_TEST = { [Discrete(2), Discrete(3), Box(-1.0, 1.0, (5, ), dtype=np.float32)]), + "dict": Dict({ + "action_choice": Discrete(3), + "parameters": Box(-1.0, 1.0, (1, ), dtype=np.float32), + "yet_another_nested_dict": Dict({ + "a": Tuple([Discrete(2), Discrete(3)]) + }) + }), } OBSERVATION_SPACES_TO_TEST = { @@ -89,9 +96,6 @@ def check_support(alg, config, stats, check_bounds=False, name=None): try: if a_name in covered_a and o_name in covered_o: stat = "skip" # speed up tests by avoiding full grid - # TODO(sven): Add necessary torch distributions. - elif torch and a_name in ["tuple", "multidiscrete"]: - stat = "unsupported" else: a = get_agent_class(alg)(config=config, env="stub_env") if alg not in ["DDPG", "ES", "ARS", "SAC"]: diff --git a/rllib/utils/exploration/random.py b/rllib/utils/exploration/random.py index 30a7d3f44..1d73d7ed5 100644 --- a/rllib/utils/exploration/random.py +++ b/rllib/utils/exploration/random.py @@ -1,14 +1,14 @@ from gym.spaces import Discrete, MultiDiscrete, Tuple import numpy as np +import tree from typing import Union from ray.rllib.models.action_dist import ActionDistribution from ray.rllib.utils.annotations import override from ray.rllib.utils.exploration.exploration import Exploration +from ray.rllib.utils import force_tuple from ray.rllib.utils.framework import try_import_tf, try_import_torch, \ TensorType -from ray.rllib.utils.tuple_actions import TupleActions -from ray.rllib.utils import force_tuple tf = try_import_tf() torch, _ = try_import_torch() @@ -75,10 +75,7 @@ class Random(Exploration): false_fn=false_fn) # TODO(sven): Move into (deterministic_)sample(logp=True|False) - if isinstance(action, TupleActions): - batch_size = tf.shape(action[0][0])[0] - else: - batch_size = tf.shape(action)[0] + batch_size = tf.shape(tree.flatten(action)[0])[0] logp = tf.zeros(shape=(batch_size, ), dtype=tf.float32) return action, logp diff --git a/rllib/utils/exploration/stochastic_sampling.py b/rllib/utils/exploration/stochastic_sampling.py index 95b91b05e..a6431745b 100644 --- a/rllib/utils/exploration/stochastic_sampling.py +++ b/rllib/utils/exploration/stochastic_sampling.py @@ -2,14 +2,15 @@ from typing import Union from ray.rllib.models.action_dist import ActionDistribution from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.utils import try_import_tree from ray.rllib.utils.annotations import override from ray.rllib.utils.exploration.exploration import Exploration from ray.rllib.utils.framework import try_import_tf, try_import_torch, \ TensorType -from ray.rllib.utils.tuple_actions import TupleActions tf = try_import_tf() torch, _ = try_import_torch() +tree = try_import_tree() class StochasticSampling(Exploration): @@ -55,11 +56,7 @@ class StochasticSampling(Exploration): false_fn=lambda: deterministic_sample) def logp_false_fn(): - # TODO(sven): Move into (deterministic_)sample(logp=True|False) - if isinstance(sample, TupleActions): - batch_size = tf.shape(action[0])[0] - else: - batch_size = tf.shape(action)[0] + batch_size = tf.shape(tree.flatten(action)[0])[0] return tf.zeros(shape=(batch_size, ), dtype=tf.float32) logp = tf.cond( @@ -67,8 +64,7 @@ class StochasticSampling(Exploration): true_fn=lambda: action_dist.sampled_action_logp(), false_fn=logp_false_fn) - return TupleActions(action) if isinstance(sample, TupleActions) \ - else action, logp + return action, logp @staticmethod def _get_torch_exploration_action(action_dist, explore): diff --git a/rllib/utils/space_utils.py b/rllib/utils/space_utils.py index 6e1f066fc..9081d7398 100644 --- a/rllib/utils/space_utils.py +++ b/rllib/utils/space_utils.py @@ -85,10 +85,10 @@ def flatten_to_single_ndarray(input_): >>> # 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0 >>> # ]) """ - # Concatenate tuple actions - if isinstance(input_, (list, tuple)): + # Concatenate complex inputs. + if isinstance(input_, (list, tuple, dict)): expanded = [] - for in_ in input_: + for in_ in tree.flatten(input_): expanded.append(np.reshape(in_, [-1])) input_ = np.concatenate(expanded, axis=0).flatten() return input_ diff --git a/rllib/utils/torch_ops.py b/rllib/utils/torch_ops.py index 08452f426..52bdc7cff 100644 --- a/rllib/utils/torch_ops.py +++ b/rllib/utils/torch_ops.py @@ -1,17 +1,9 @@ import numpy as np -import logging -from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils import try_import_torch, try_import_tree torch, _ = try_import_torch() - -logger = logging.getLogger(__name__) - -try: - import tree -except (ImportError, ModuleNotFoundError) as e: - logger.warning("`dm-tree` is not installed! Run `pip install dm-tree`.") - raise e +tree = try_import_tree() def global_norm(tensors): diff --git a/rllib/utils/tuple_actions.py b/rllib/utils/tuple_actions.py index c2d9f299c..71d2d65d0 100644 --- a/rllib/utils/tuple_actions.py +++ b/rllib/utils/tuple_actions.py @@ -1,10 +1,18 @@ from collections import namedtuple +from ray.rllib.utils.deprecation import deprecation_warning +# NOTE: This is a deprecated class. Use native python tuples +# or dicts (both arbitrarily nested) for multi-actions from here on. class TupleActions(namedtuple("TupleActions", ["batches"])): """Used to return tuple actions as a list of batches per tuple element.""" def __new__(cls, batches): + # Throw an informative error if used. + deprecation_warning( + old="TupleActions", + new="`native python tuples (arbitrarily nested)`", + error=True) return super(TupleActions, cls).__new__(cls, batches) def numpy(self):