def _create_test_value(spec):
"""Creates a scalar test value consistent with the TensorSpec `spec`."""
if _is_numeric_type(spec.dtype):
return tensor_spec_utils.bounds(spec).min,
else:
np_type = tensor_utils.data_type_to_np_type(spec.dtype)
return np_type()
def tensor_spec_to_dm_env_spec(
tensor_spec: dm_env_rpc_pb2.TensorSpec) -> specs.Array:
"""Returns a dm_env spec given a dm_env_rpc TensorSpec.
Args:
tensor_spec: A dm_env_rpc TensorSpec protobuf.
Returns:
Either a DiscreteArray, BoundedArray, StringArray or Array, depending on the
content of the TensorSpec.
"""
np_type = tensor_utils.data_type_to_np_type(tensor_spec.dtype)
if tensor_spec.HasField('min') or tensor_spec.HasField('max'):
bounds = tensor_spec_utils.bounds(tensor_spec)
if (not tensor_spec.shape and np.issubdtype(np_type, np.integer)
and bounds.min == 0 and tensor_spec.HasField('max')):
return specs.DiscreteArray(num_values=bounds.max + 1,
dtype=np_type,
name=tensor_spec.name)
else:
return specs.BoundedArray(shape=tensor_spec.shape,
dtype=np_type,
name=tensor_spec.name,
minimum=bounds.min,
maximum=bounds.max)
else:
if tensor_spec.dtype == dm_env_rpc_pb2.DataType.STRING:
return specs.StringArray(shape=tensor_spec.shape,
name=tensor_spec.name)
else:
return specs.Array(shape=tensor_spec.shape,
dtype=np_type,
name=tensor_spec.name)
def unpack(
self, dm_env_rpc_tensors: Mapping[int, dm_env_rpc_pb2.Tensor]
) -> MutableMapping[str, Any]:
"""Unpacks a dm_env_rpc uid-to-tensor map to a name-keyed Python dict.
Args:
dm_env_rpc_tensors: A dict mapping UIDs to dm_env_rpc tensor protos.
Returns:
A dict mapping names to scalars and arrays.
"""
unpacked = {}
for uid, tensor in dm_env_rpc_tensors.items():
name = self._uid_to_name[uid]
dm_env_rpc_spec = self.name_to_spec(name)
_assert_shapes_match(tensor, dm_env_rpc_spec)
tensor_dtype = tensor_utils.get_tensor_type(tensor)
spec_dtype = tensor_utils.data_type_to_np_type(
dm_env_rpc_spec.dtype)
if tensor_dtype != spec_dtype:
raise ValueError(
'Received dm_env_rpc tensor {} with dtype {} but spec has dtype {}.'
.format(name, tensor_dtype, spec_dtype))
tensor_unpacked = tensor_utils.unpack_tensor(tensor)
unpacked[name] = tensor_unpacked
return unpacked
def set_bounds(tensor_spec: dm_env_rpc_pb2.TensorSpec, minimum, maximum):
"""Modifies `tensor_spec` to have its inclusive bounds set.
Packs `minimum` in to `tensor_spec.min` and `maximum` in to `tensor_spec.max`.
Args:
tensor_spec: An instance of a dm_env_rpc TensorSpec proto. It should
already have its `name`, `dtype` and `shape` attributes set.
minimum: The minimum value that elements in the described tensor can obtain.
A scalar, iterable of scalars, or None. If None, `min` will be cleared on
`tensor_spec`.
maximum: The maximum value that elements in the described tensor can obtain.
A scalar, iterable of scalars, or None. If None, `max` will be cleared on
`tensor_spec`.
"""
np_type = tensor_utils.data_type_to_np_type(tensor_spec.dtype)
if not issubclass(np_type, np.number):
raise ValueError(f'TensorSpec has non-numeric type "{np_type}".')
np_type_bounds = _np_range_info(np_type)
has_min = minimum is not None
has_max = maximum is not None
if has_min:
minimum = np.asarray(minimum)
if minimum.size != 1 and minimum.shape != tuple(tensor_spec.shape):
raise ValueError(
f'minimum has shape {minimum.shape}, which is incompatible with '
f"tensor_spec {tensor_spec.name}'s shape {tensor_spec.shape}.")
if has_max:
maximum = np.asarray(maximum)
if maximum.size != 1 and maximum.shape != tuple(tensor_spec.shape):
raise ValueError(
f'maximum has shape {maximum.shape}, which is incompatible with '
f"tensor_spec {tensor_spec.name}'s shape {tensor_spec.shape}.")
if ((has_min and not _can_cast(minimum, np_type)) or
(has_max and not _can_cast(maximum, np_type))):
raise ValueError(
_BOUNDS_CANNOT_BE_SAFELY_CAST_TO_DTYPE.format(
name=tensor_spec.name,
minimum=minimum,
maximum=maximum,
dtype=dm_env_rpc_pb2.DataType.Name(tensor_spec.dtype)))
if (has_min and has_max and np.any(maximum < minimum)):
raise ValueError('TensorSpec "{}" has min {} larger than max {}.'.format(
tensor_spec.name, minimum, maximum))
packer = tensor_utils.get_packer(np_type)
if has_min:
packer.pack(tensor_spec.min, minimum)
else:
tensor_spec.ClearField('min')
if has_max:
packer.pack(tensor_spec.max, maximum)
else:
tensor_spec.ClearField('max')
def test_all_observation_dtypes_match_spec_dtypes(self):
response = self.step(requested_observations=self.observation_uids)
for uid, observation in response.observations.items():
spec = self.specs.observations[uid]
with self.subTest(uid=uid, name=spec.name):
spec_type = tensor_utils.data_type_to_np_type(spec.dtype)
tensor_type = tensor_utils.get_tensor_type(observation)
self.assertEqual(spec_type, tensor_type)
def test_cannot_send_wrong_numeric_type_action(self):
for uid, spec in self.numeric_actions.items():
with self.subTest(uid=uid, name=spec.name):
np_type = tensor_utils.data_type_to_np_type(spec.dtype)
wrong_dtype = (np.int32 if np_type == np.issubdtype(
np_type, np.inexact) else np.float32)
tensor = _create_test_tensor(spec, dtype=wrong_dtype)
with self.assertRaises(error.DmEnvRpcError):
self.step(actions={uid: tensor})
def bounds(tensor_spec: dm_env_rpc_pb2.TensorSpec) -> Bounds:
"""Gets the inclusive bounds of `tensor_spec`.
Args:
tensor_spec: An instance of a dm_env_rpc TensorSpec proto.
Returns:
A named tuple (`min`, `max`) of inclusive bounds.
Raises:
ValueError: `tensor_spec` does not have a numeric dtype, or the type of its
`min` or `max` does not match its dtype, or the the bounds are invalid in
some way.
"""
np_type = tensor_utils.data_type_to_np_type(tensor_spec.dtype)
tensor_spec_type = dm_env_rpc_pb2.DataType.Name(tensor_spec.dtype).lower()
if not issubclass(np_type, np.number):
raise ValueError('TensorSpec "{}" has non-numeric type {}.'.format(
tensor_spec.name, tensor_spec_type))
# Check min payload type matches the tensor type.
min_which = tensor_spec.min.WhichOneof('payload')
if min_which and not min_which.startswith(tensor_spec_type):
raise ValueError(
'TensorSpec "{}" has dtype {} but min type {}.'.format(
tensor_spec.name, tensor_spec_type, min_which))
# Check max payload type matches the tensor type.
max_which = tensor_spec.max.WhichOneof('payload')
if max_which and not max_which.startswith(tensor_spec_type):
raise ValueError(
'TensorSpec "{}" has dtype {} but max type {}.'.format(
tensor_spec.name, tensor_spec_type, max_which))
dtype_bounds = _np_range_info(np_type)
min_bound = _get_value(tensor_spec.min, tensor_spec.shape,
dtype_bounds.min)
max_bound = _get_value(tensor_spec.max, tensor_spec.shape,
dtype_bounds.max)
if not _can_cast(min_bound, np_type) or not _can_cast(max_bound, np_type):
raise ValueError(
_BOUNDS_CANNOT_BE_SAFELY_CAST_TO_DTYPE.format(
name=tensor_spec.name,
minimum=min_bound,
maximum=max_bound,
dtype=tensor_spec_type))
if np.any(max_bound < min_bound):
raise ValueError(
'TensorSpec "{}" has min {} larger than max {}.'.format(
tensor_spec.name, min_bound, max_bound))
return Bounds(np_type(min_bound), np_type(max_bound))
def _below_min(spec):
"""Returns a value below spec's min or None if none."""
if not spec.HasField('min'):
return None
np_type = tensor_utils.data_type_to_np_type(spec.dtype)
min_type_value = _np_range_info(np_type).min
if min_type_value < tensor_spec_utils.bounds(spec).min:
return min_type_value
else:
return None
def _above_max(spec):
"""Returns a value above spec's max or None if none."""
if not spec.HasField('max'):
return None
np_type = tensor_utils.data_type_to_np_type(spec.dtype)
max_type_value = _np_range_info(np_type).max
if max_type_value > tensor_spec_utils.bounds(spec).max:
return max_type_value
else:
return None
def tensor_spec_to_dm_env_spec(tensor_spec):
"""Returns the dm_env Array or BoundedArray given a dm_env_rpc TensorSpec."""
tensor_type = tensor_utils.data_type_to_np_type(tensor_spec.dtype)
if tensor_spec.HasField('min') or tensor_spec.HasField('max'):
return specs.BoundedArray(
shape=tensor_spec.shape,
dtype=tensor_type,
name=tensor_spec.name,
minimum=_find_extreme(tensor_spec, tensor_type, 'min'),
maximum=_find_extreme(tensor_spec, tensor_type, 'max'))
else:
return specs.Array(shape=tensor_spec.shape,
dtype=tensor_type,
name=tensor_spec.name)
def _is_numeric_type(dtype):
return (dtype != dm_env_rpc_pb2.DataType.PROTO and np.issubdtype(
tensor_utils.data_type_to_np_type(dtype), np.number))
def test_unknown_type(self):
with self.assertRaises(TypeError):
tensor_utils.data_type_to_np_type(30)
def test_proto_type(self):
with self.assertRaises(TypeError):
tensor_utils.data_type_to_np_type(dm_env_rpc_pb2.DataType.PROTO)
def test_float(self):
self.assertEqual(
np.float32,
tensor_utils.data_type_to_np_type(dm_env_rpc_pb2.DataType.FLOAT))
