diff --git a/rllib/agents/ars/ars_tf_policy.py b/rllib/agents/ars/ars_tf_policy.py index b47cef75d..c201000b7 100644 --- a/rllib/agents/ars/ars_tf_policy.py +++ b/rllib/agents/ars/ars_tf_policy.py @@ -11,7 +11,7 @@ from ray.rllib.models import ModelCatalog from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.utils.filter import get_filter from ray.rllib.utils.framework import try_import_tf -from ray.rllib.utils.space_utils import unbatch +from ray.rllib.utils.spaces.space_utils import unbatch tf = try_import_tf() diff --git a/rllib/agents/es/es_tf_policy.py b/rllib/agents/es/es_tf_policy.py index da2c274e5..166129b73 100644 --- a/rllib/agents/es/es_tf_policy.py +++ b/rllib/agents/es/es_tf_policy.py @@ -11,7 +11,8 @@ 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, unbatch +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space, \ + unbatch tf = try_import_tf() tree = try_import_tree() diff --git a/rllib/agents/es/es_torch_policy.py b/rllib/agents/es/es_torch_policy.py index 95d844584..eddefdbb5 100644 --- a/rllib/agents/es/es_torch_policy.py +++ b/rllib/agents/es/es_torch_policy.py @@ -10,7 +10,7 @@ from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.policy.torch_policy_template import build_torch_policy from ray.rllib.utils.filter import get_filter from ray.rllib.utils.framework import try_import_torch -from ray.rllib.utils.space_utils import unbatch +from ray.rllib.utils.spaces.space_utils import unbatch from ray.rllib.utils.torch_ops import convert_to_torch_tensor torch, _ = try_import_torch() diff --git a/rllib/evaluation/episode.py b/rllib/evaluation/episode.py index 2f3d37379..1d903dbb0 100644 --- a/rllib/evaluation/episode.py +++ b/rllib/evaluation/episode.py @@ -5,7 +5,7 @@ import random 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 +from ray.rllib.utils.spaces.space_utils import flatten_to_single_ndarray tree = try_import_tree() diff --git a/rllib/evaluation/sampler.py b/rllib/evaluation/sampler.py index c8e50b1f3..7c4ef67fe 100644 --- a/rllib/evaluation/sampler.py +++ b/rllib/evaluation/sampler.py @@ -18,7 +18,8 @@ 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.space_utils import flatten_to_single_ndarray, unbatch +from ray.rllib.utils.spaces.space_utils import flatten_to_single_ndarray, \ + unbatch from ray.rllib.utils.tf_run_builder import TFRunBuilder tree = try_import_tree() diff --git a/rllib/examples/env/nested_space_repeat_after_me_env.py b/rllib/examples/env/nested_space_repeat_after_me_env.py index c48505fdb..ac3df4b6c 100644 --- a/rllib/examples/env/nested_space_repeat_after_me_env.py +++ b/rllib/examples/env/nested_space_repeat_after_me_env.py @@ -3,7 +3,7 @@ from gym.spaces import Box, Dict, Discrete, Tuple import numpy as np from ray.rllib.utils import try_import_tree -from ray.rllib.utils.space_utils import flatten_space +from ray.rllib.utils.spaces.space_utils import flatten_space tree = try_import_tree() diff --git a/rllib/models/catalog.py b/rllib/models/catalog.py index 2ed0220b6..79e76b3ee 100644 --- a/rllib/models/catalog.py +++ b/rllib/models/catalog.py @@ -5,7 +5,6 @@ import numpy as np from ray.tune.registry import RLLIB_MODEL, RLLIB_PREPROCESSOR, \ RLLIB_ACTION_DIST, _global_registry -from ray.rllib.models.extra_spaces import Simplex from ray.rllib.models.action_dist import ActionDistribution from ray.rllib.models.modelv2 import ModelV2 from ray.rllib.models.preprocessors import get_preprocessor @@ -25,7 +24,8 @@ from ray.rllib.utils import try_import_tf, try_import_tree from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI from ray.rllib.utils.deprecation import deprecation_warning, DEPRECATED_VALUE from ray.rllib.utils.error import UnsupportedSpaceException -from ray.rllib.utils.space_utils import flatten_space +from ray.rllib.utils.spaces.simplex import Simplex +from ray.rllib.utils.spaces.space_utils import flatten_space tf = try_import_tf() tree = try_import_tree() diff --git a/rllib/models/tf/tf_action_dist.py b/rllib/models/tf/tf_action_dist.py index b2f59ec57..d3a56f92d 100644 --- a/rllib/models/tf/tf_action_dist.py +++ b/rllib/models/tf/tf_action_dist.py @@ -7,7 +7,7 @@ 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.framework import try_import_tf, try_import_tfp -from ray.rllib.utils.space_utils import get_base_struct_from_space +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space tf = try_import_tf() tfp = try_import_tfp() diff --git a/rllib/models/torch/torch_action_dist.py b/rllib/models/torch/torch_action_dist.py index ce9c19038..64cc07a9f 100644 --- a/rllib/models/torch/torch_action_dist.py +++ b/rllib/models/torch/torch_action_dist.py @@ -8,7 +8,7 @@ from ray.rllib.utils.annotations import override from ray.rllib.utils.framework import try_import_torch from ray.rllib.utils.numpy import SMALL_NUMBER, MIN_LOG_NN_OUTPUT, \ MAX_LOG_NN_OUTPUT -from ray.rllib.utils.space_utils import get_base_struct_from_space +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space from ray.rllib.utils.torch_ops import atanh torch, nn = try_import_torch() diff --git a/rllib/policy/eager_tf_policy.py b/rllib/policy/eager_tf_policy.py index 62a5ecf30..6c37e1fa5 100644 --- a/rllib/policy/eager_tf_policy.py +++ b/rllib/policy/eager_tf_policy.py @@ -13,7 +13,7 @@ from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.utils import add_mixins from ray.rllib.utils.annotations import override from ray.rllib.utils.framework import try_import_tf -from ray.rllib.utils.space_utils import flatten_to_single_ndarray +from ray.rllib.utils.spaces.space_utils import flatten_to_single_ndarray tf = try_import_tf() logger = logging.getLogger(__name__) diff --git a/rllib/policy/policy.py b/rllib/policy/policy.py index 786fcbbb0..142729f39 100644 --- a/rllib/policy/policy.py +++ b/rllib/policy/policy.py @@ -7,7 +7,7 @@ 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, \ +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space, \ unbatch tree = try_import_tree() diff --git a/rllib/rollout.py b/rllib/rollout.py index 3df0233e1..7f0b8dc48 100755 --- a/rllib/rollout.py +++ b/rllib/rollout.py @@ -16,7 +16,7 @@ from ray.rllib.env.base_env import _DUMMY_AGENT_ID 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.space_utils import flatten_to_single_ndarray +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 diff --git a/rllib/utils/exploration/random.py b/rllib/utils/exploration/random.py index 8da2b7c7a..581883320 100644 --- a/rllib/utils/exploration/random.py +++ b/rllib/utils/exploration/random.py @@ -9,7 +9,7 @@ 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.space_utils import get_base_struct_from_space +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space tf = try_import_tf() torch, _ = try_import_torch() diff --git a/rllib/utils/spaces/__init__.py b/rllib/utils/spaces/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/rllib/models/extra_spaces.py b/rllib/utils/spaces/simplex.py similarity index 100% rename from rllib/models/extra_spaces.py rename to rllib/utils/spaces/simplex.py diff --git a/rllib/utils/space_utils.py b/rllib/utils/spaces/space_utils.py similarity index 100% rename from rllib/utils/space_utils.py rename to rllib/utils/spaces/space_utils.py