From 519950a576e9ce0e02092f79ed1d941ead9090a2 Mon Sep 17 00:00:00 2001 From: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> Date: Tue, 29 Sep 2026 00:59:37 +0100 Subject: [PATCH] Replace existing TensorSpec bounds before packing new values --- .../v1/tensor_spec_bounds_replacement_test.py | 101 ++++++++++++++++++ dm_env_rpc/v1/tensor_spec_utils.py | 6 +- 2 files changed, 103 insertions(+), 4 deletions(-) create mode 100644 dm_env_rpc/v1/tensor_spec_bounds_replacement_test.py diff --git a/dm_env_rpc/v1/tensor_spec_bounds_replacement_test.py b/dm_env_rpc/v1/tensor_spec_bounds_replacement_test.py new file mode 100644 index 0000000..4c8bb3e --- /dev/null +++ b/dm_env_rpc/v1/tensor_spec_bounds_replacement_test.py @@ -0,0 +1,101 @@ +# Copyright 2019 DeepMind Technologies Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""Regression tests for replacing existing TensorSpec bounds.""" + +from absl.testing import absltest +from absl.testing import parameterized +import numpy as np + +from dm_env_rpc.v1 import dm_env_rpc_pb2 +from dm_env_rpc.v1 import tensor_spec_utils +from dm_env_rpc.v1 import tensor_utils + +_DTYPES = (np.float32, np.float64, np.int8, np.int32, np.int64, np.uint8, + np.uint32, np.uint64) + + +class BoundsReplacementTest(parameterized.TestCase): + + def make_spec(self, dtype, shape): + return dm_env_rpc_pb2.TensorSpec( + name='observation', + shape=shape, + dtype=tensor_utils.np_type_to_data_type(dtype)) + + def assert_replacement(self, spec, minimum, maximum): + expected = self.make_spec(tensor_utils.data_type_to_np_type(spec.dtype), + tuple(spec.shape)) + tensor_spec_utils.set_bounds(expected, minimum, maximum) + self.assertIsNone(tensor_spec_utils.set_bounds(spec, minimum, maximum)) + self.assertEqual(spec, expected) + self.assertEqual(spec.SerializeToString(), expected.SerializeToString()) + actual_bounds = tensor_spec_utils.bounds(spec) + expected_bounds = tensor_spec_utils.bounds(expected) + np.testing.assert_array_equal(actual_bounds.min, expected_bounds.min) + np.testing.assert_array_equal(actual_bounds.max, expected_bounds.max) + + @parameterized.product(dtype=_DTYPES, shape=((), (2,))) + def test_repeated_updates_replace_values(self, dtype, shape): + spec = self.make_spec(dtype, shape) + for lower, upper in ((1, 9), (2, 8), (2, 8), (0, 7)): + minimum = np.full(shape, lower, dtype=dtype) + maximum = np.full(shape, upper, dtype=dtype) + self.assert_replacement(spec, minimum, maximum) + np.testing.assert_array_equal(minimum, np.full(shape, lower, dtype=dtype)) + np.testing.assert_array_equal(maximum, np.full(shape, upper, dtype=dtype)) + + @parameterized.product(dtype=(np.int32, np.float64), + vector_first=(False, True)) + def test_scalar_and_vector_bounds_can_replace_each_other( + self, dtype, vector_first): + spec = self.make_spec(dtype, (2,)) + shapes = [(2,), ()] if vector_first else [(), (2,)] + for shape in shapes: + self.assert_replacement(spec, np.full(shape, 1, dtype=dtype), + np.full(shape, 6, dtype=dtype)) + + @parameterized.product(dtype=(np.int8, np.int32), + present=((False, False), (False, True), (True, False))) + def test_clearing_either_bound_preserves_the_other_replacement( + self, dtype, present): + spec = self.make_spec(dtype, ()) + tensor_spec_utils.set_bounds(spec, dtype(1), dtype(5)) + minimum = dtype(2) if present[0] else None + maximum = dtype(4) if present[1] else None + self.assert_replacement(spec, minimum, maximum) + self.assertEqual(spec.HasField('min'), present[0]) + self.assertEqual(spec.HasField('max'), present[1]) + + @parameterized.parameters(([1, 2, 3], 5), (1, [4, 5, 6]), ([5, 1], [2, 3]), + (-(2**40), 5), (1, 2**40)) + def test_failed_validation_does_not_clear_existing_bounds( + self, minimum, maximum): + spec = self.make_spec(np.int32, (2,)) + tensor_spec_utils.set_bounds(spec, [1, 2], [5, 6]) + original = spec.SerializeToString() + with self.assertRaises(ValueError): + tensor_spec_utils.set_bounds(spec, minimum, maximum) + self.assertEqual(spec.SerializeToString(), original) + + def test_variable_length_spec_retains_scalar_bounds(self): + spec = self.make_spec(np.int32, (-1, 2)) + for minimum, maximum in ((1, 8), (2, 5)): + self.assert_replacement(spec, minimum, maximum) + self.assertEqual(len(spec.min.int32s.array), 1) + self.assertEqual(len(spec.max.int32s.array), 1) + + +if __name__ == '__main__': + absltest.main() diff --git a/dm_env_rpc/v1/tensor_spec_utils.py b/dm_env_rpc/v1/tensor_spec_utils.py index cbe1f2b..733c9fe 100644 --- a/dm_env_rpc/v1/tensor_spec_utils.py +++ b/dm_env_rpc/v1/tensor_spec_utils.py @@ -199,12 +199,10 @@ def set_bounds(tensor_spec: dm_env_rpc_pb2.TensorSpec, minimum, maximum): tensor_spec.name, minimum, maximum)) packer = tensor_utils.get_packer(np_type) + tensor_spec.ClearField('min') if has_min: packer.pack(tensor_spec.min, minimum) - else: - tensor_spec.ClearField('min') + tensor_spec.ClearField('max') if has_max: packer.pack(tensor_spec.max, maximum) - else: - tensor_spec.ClearField('max')