def __init__(self, env, idx):
"""Creates an observation filter wrapper.
Args:
env: Environment to wrap.
idx: Array of indexes pointing to elements to include in output.
Raises:
ValueError: If observation spec is nested.
ValueError: If indexes are not single-dimensional.
ValueError: If no index is provided.
ValueError: If one of the indexes is out of bounds.
"""
super(ObservationFilterWrapper, self).__init__(env)
idx = np.array(idx)
if tf.nest.is_nested(env.observation_spec()):
raise ValueError(
'ObservationFilterWrapper only works with single-array '
'observations (not nested).')
if len(idx.shape) != 1:
raise ValueError('ObservationFilterWrapper only works with '
'single-dimensional indexes for filtering.')
if idx.shape[0] < 1:
raise ValueError(
'At least one index needs to be provided for filtering.')
if not np.all(idx < env.observation_spec().shape[0]):
raise ValueError('One of the indexes is out of bounds.')
self._idx = idx
self._observation_spec = array_spec.update_spec_shape(
env.observation_spec(), idx.shape)
def _update_shape(spec):
return array_spec.update_spec_shape(spec,
(self._history_length,) + spec.shape)