Fake environment with nested observations.
@@ -20,7 +20,7 @@ order to test or interact with other components.
"""
import threading
-from typing import List, Sequence, Optional
+from typing import List, Mapping, Sequence, Optional
from acme import core
from acme import specs
@@ -166,16 +166,14 @@ class Environment(dm_env.Environment):
return self._spec.discounts
-class DiscreteEnvironment(Environment):
- """Discrete state and action fake environment."""
+class _BaseDiscreteEnvironment(Environment):
+ """Discrete action fake environment."""
def __init__(self,
*,
num_actions: int = 1,
- num_observations: int = 1,
action_dtype=np.int32,
- obs_dtype=np.int32,
- obs_shape: Sequence[int] = (),
+ observation_spec: types.NestedSpec,
discount_spec: Optional[types.NestedSpec] = None,
reward_spec: Optional[types.NestedSpec] = None,
**kwargs):
@@ -187,21 +185,77 @@ class DiscreteEnvironment(Environment):
discount_spec = specs.BoundedArray((), np.float32, 0.0, 1.0)
actions = specs.DiscreteArray(num_actions, dtype=action_dtype)
- observations = specs.BoundedArray(
- shape=obs_shape,
- dtype=obs_dtype,
- minimum=obs_dtype(0),
- maximum=obs_dtype(num_observations - 1))
super().__init__(
spec=specs.EnvironmentSpec(
- observations=observations,
+ observations=observation_spec,
actions=actions,
rewards=reward_spec,
discounts=discount_spec),
**kwargs)
+class DiscreteEnvironment(_BaseDiscreteEnvironment):
+ """Discrete state and action fake environment."""
+
+ def __init__(self,
+ *,
+ num_actions: int = 1,
+ num_observations: int = 1,
+ action_dtype=np.int32,
+ obs_dtype=np.int32,
+ obs_shape: Sequence[int] = (),
+ discount_spec: Optional[types.NestedSpec] = None,
+ reward_spec: Optional[types.NestedSpec] = None,
+ **kwargs):
+ """Initialize the environment."""
+ observations_spec = specs.BoundedArray(
+ shape=obs_shape,
+ dtype=obs_dtype,
+ minimum=obs_dtype(0),
+ maximum=obs_dtype(num_observations - 1))
+
+ super().__init__(
+ num_actions=num_actions,
+ action_dtype=action_dtype,
+ observation_spec=observations_spec,
+ discount_spec=discount_spec,
+ reward_spec=reward_spec,
+ **kwargs)
+
+
+class NestedDiscreteEnvironment(_BaseDiscreteEnvironment):
+ """Discrete action fake environment with nested discrete state."""
+
+ def __init__(self,
+ *,
+ num_observations: Mapping[str, int],
+ num_actions: int = 1,
+ action_dtype=np.int32,
+ obs_dtype=np.int32,
+ obs_shape: Sequence[int] = (),
+ discount_spec: Optional[types.NestedSpec] = None,
+ reward_spec: Optional[types.NestedSpec] = None,
+ **kwargs):
+ """Initialize the environment."""
+
+ observations_spec = {}
+ for key in num_observations:
+ observations_spec[key] = specs.BoundedArray(
+ shape=obs_shape,
+ dtype=obs_dtype,
+ minimum=obs_dtype(0),
+ maximum=obs_dtype(num_observations[key] - 1))
+
+ super().__init__(
+ num_actions=num_actions,
+ action_dtype=action_dtype,
+ observation_spec=observations_spec,
+ discount_spec=discount_spec,
+ reward_spec=reward_spec,
+ **kwargs)
+
+
class ContinuousEnvironment(Environment):
"""Continuous state and action fake environment."""
