def test_broadcasts_to_multidimensional_arrays(self):
tensor = dm_env_rpc_pb2.Tensor()
tensor.int32s.array[:] = [4]
tensor.shape[:] = [2, 2]
unpacked = tensor_utils.unpack_tensor(tensor)
expected = np.array([[4, 4], [4, 4]], dtype=np.int32)
np.testing.assert_array_equal(expected, unpacked)
def test_integer_broadcasts_1_element_to_all_elements(self):
tensor = dm_env_rpc_pb2.Tensor()
tensor.int32s.array[:] = [1]
tensor.shape[:] = [4]
unpacked = tensor_utils.unpack_tensor(tensor)
expected = np.array([1, 1, 1, 1], dtype=np.int32)
np.testing.assert_array_equal(expected, unpacked)
def test_unpack_multidimensional_arrays(self):
tensor = dm_env_rpc_pb2.Tensor()
tensor.floats.array[:] = [1, 2, 3, 4, 5, 6, 7, 8]
tensor.shape[:] = [2, 4]
round_trip = tensor_utils.unpack_tensor(tensor)
expected = np.array([[1, 2, 3, 4], [5, 6, 7, 8]])
np.testing.assert_array_equal(expected, round_trip)
def test_negative_dimension_in_matrix(self):
tensor = dm_env_rpc_pb2.Tensor()
tensor.int32s.array[:] = [1, 2, 3, 4, 5, 6]
tensor.shape[:] = [2, -1]
unpacked = tensor_utils.unpack_tensor(tensor)
expected = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32)
np.testing.assert_array_equal(expected, unpacked)
def test_negative_dimension_single_element(self):
tensor = dm_env_rpc_pb2.Tensor()
tensor.int32s.array[:] = [1]
tensor.shape[:] = [-1]
unpacked = tensor_utils.unpack_tensor(tensor)
expected = np.array([1], dtype=np.int32)
np.testing.assert_array_equal(expected, unpacked)
def test_unsigned_integer_broadcasts_1_element_to_all_elements(self):
tensor = dm_env_rpc_pb2.Tensor()
tensor.uint8s.array = b'\x01'
tensor.shape[:] = [4]
unpacked = tensor_utils.unpack_tensor(tensor)
expected = np.array([1, 1, 1, 1], dtype=np.uint8)
np.testing.assert_array_equal(expected, unpacked)
def test_string_broadcasts_1_element_to_all_elements(self):
tensor = dm_env_rpc_pb2.Tensor()
tensor.strings.array[:] = ['foo']
tensor.shape[:] = [4]
unpacked = tensor_utils.unpack_tensor(tensor)
expected = np.array(['foo', 'foo', 'foo', 'foo'], dtype=np.str_)
np.testing.assert_array_equal(expected, unpacked)
def pack_tensor(value,
dtype=None,
try_compress=False) -> dm_env_rpc_pb2.Tensor:
"""Encodes the given value as a tensor.
Args:
value: A scalar (float, int, string, etc.), protobuf message, NumPy array,
or nested lists.
dtype: The type to pack the data to. If set to None, will attempt to detect
the correct type automatically. Either a dm_env_rpc DataType enum or
NumPy type is acceptable.
try_compress: A bool, whether to try and encode the tensor in less space or
not. This will increase the computational cost of the packing, but may
reduce the on-the-wire size of the tensor. There are no guarantees that
any compression will actually happen.
Raises:
ValueError: If `value` is a jagged array, not a primitive type, nested
iterable of primitive types or protobuf messages, or all elements can't be
cast to the same type or the requested type.
Returns:
A dm_env_rpc Tensor proto containing the data.
"""
packed = dm_env_rpc_pb2.Tensor()
value = np.asarray(value)
if value.dtype == np.object:
# Because Numpy doesn't truly support variable-length string arrays, users
# tend to use arrays of Numpy objects instead. Iff a user provides an array
# of objects and a string dtype argument, automatically convert the value to
# an array of strings.
if np.issubdtype(_DM_ENV_RPC_DTYPE_TO_NUMPY_DTYPE.get(dtype, dtype),
np.str_):
for item in value.flat:
if not isinstance(item, str):
raise TypeError(
f'Requested string dtype but not all elements are '
'Python string types. At least one element was '
f'{type(item)}.')
value = np.array(value, dtype=np.str_)
elif dtype is not None:
value = value.astype(dtype=_DM_ENV_RPC_DTYPE_TO_NUMPY_DTYPE.get(
dtype, dtype),
copy=False,
casting='same_kind')
packed.shape[:] = value.shape
packer = _TYPE_TO_PACKER[value.dtype.type]
if (try_compress and np.all(value == next(value.flat))):
# All elements are the same. Pack in to a single value.
packer.pack(packed, value.ravel()[0:1])
else:
packer.pack(packed, value)
return packed
def pack_tensor(value, dtype=None, try_compress=False):
"""Encodes the given value as a tensor.
Args:
value: A scalar (float, int, string, etc.), protobuf message, NumPy array,
or nested lists.
dtype: The type to pack the data to. If set to None, will attempt to detect
the correct type automatically. Either a dm_env_rpc DataType enum or
NumPy type is acceptable.
try_compress: A bool, whether to try and encode the tensor in less space or
not. This will increase the computational cost of the packing, but may
reduce the on-the-wire size of the tensor. There are no guarantees that
any compression will actually happen.
Raises:
ValueError: If `value` is a jagged array, not a primitive type, nested
iterable of primitive types or protobuf messages, or all elements can't be
cast to the same type or the requested type.
Returns:
A dm_env_rpc Tensor proto containing the data.
"""
packed = dm_env_rpc_pb2.Tensor()
value = np.asarray(value)
# For efficiency, only check that the first element is a protobuf message.
if value.dtype == np.object and value.size > 0 and not isinstance(
value.item(0), message.Message):
raise ValueError(
'Could not convert value to a tensor of primitive types: '
f'{value}. Are the iterables jagged? Are the data types '
'not primitive scalar types like strings, floats, or '
'integers? Or are the elements not protobuf messages?')
if dtype is not None:
value = value.astype(dtype=_DM_ENV_RPC_DTYPE_TO_NUMPY_DTYPE.get(
dtype, dtype),
copy=False,
casting='same_kind')
packed.shape[:] = value.shape
packer = _TYPE_TO_PACKER[value.dtype.type]
if (try_compress and np.all(value == next(value.flat))):
# All elements are the same. Pack in to a single value.
packer.pack(packed, value.ravel()[0:1])
else:
packer.pack(packed, value)
return packed
def test_can_unpack(self):
packer = tensor_utils.get_packer(np.int32)
tensor = dm_env_rpc_pb2.Tensor()
tensor.int32s.array[:] = [1, 2, 3]
np.testing.assert_array_equal([1, 2, 3], packer.unpack(tensor))
def test_scalar_with_too_many_elements_raises_error(self):
tensor = dm_env_rpc_pb2.Tensor()
tensor.int32s.array[:] = [1, 2, 3]
with self.assertRaisesRegex(ValueError, '3 element'):
tensor_utils.unpack_tensor(tensor)
def test_can_pack(self):
packer = tensor_utils.get_packer(np.int32)
tensor = dm_env_rpc_pb2.Tensor()
packer.pack(tensor, np.asarray([1, 2, 3]))
self.assertEqual([1, 2, 3], tensor.int32s.array)
def test_setting_tensor_data_with_wrong_type(self):
tensor = dm_env_rpc_pb2.Tensor()
with self.assertRaises(TypeError):
tensor.floats.array[:] = ['hello!']
def test_two_negative_dimensions_in_matrix(self):
tensor = dm_env_rpc_pb2.Tensor()
tensor.int32s.array[:] = [1, 2, 3, 4, 5, 6]
tensor.shape[:] = [-1, -2]
with self.assertRaisesRegex(ValueError, 'one unknown dimension'):
tensor_utils.unpack_tensor(tensor)
def test_which_is_set(self):
tensor = dm_env_rpc_pb2.Tensor()
tensor.floats.array[:] = [1, 2]
self.assertEqual('floats', tensor.WhichOneof('payload'))
def test_too_many_elements(self):
tensor = dm_env_rpc_pb2.Tensor()
tensor.floats.array[:] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
tensor.shape[:] = [2, 4]
with self.assertRaisesRegex(ValueError, 'cannot reshape array'):
tensor_utils.unpack_tensor(tensor)
def test_setting_tensor_data(self): tensor = dm_env_rpc_pb2.Tensor() tensor.floats.array[:] = [1, 2]
