From fc925a867d51fa539f0418f0ff48f996180bb4e3 Mon Sep 17 00:00:00 2001
From: DiadkoShmek <207617272+DiadkoShmek@users.noreply.github.com>
Date: Mon, 2 Nov 2026 19:35:54 +0200
Subject: [PATCH] Reject NaN values in bounded array validation

---
 dm_env/specs.py           |  2 +-
 dm_env/specs_test.py      | 14 ++++++++++++++
 dm_env/test_utils_test.py | 15 +++++++++++++++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/dm_env/specs.py b/dm_env/specs.py
index 0dc989a..fe7fa43 100644
--- a/dm_env/specs.py
+++ b/dm_env/specs.py
@@ -251,7 +251,7 @@ def __eq__(self, other):
   def validate(self, value):
     value = np.asarray(value)
     super(BoundedArray, self).validate(value)
-    if (value < self.minimum).any() or (value > self.maximum).any():
+    if not (value >= self.minimum).all() or not (value <= self.maximum).all():
       self._fail_validation(_OUT_OF_BOUNDS, self.minimum, value, self.maximum)
     return value
 
diff --git a/dm_env/specs_test.py b/dm_env/specs_test.py
index 63a243e..f8b8750 100644
--- a/dm_env/specs_test.py
+++ b/dm_env/specs_test.py
@@ -289,6 +289,20 @@ def testValidateBoundsFloat(self, minimum, maximum, value, is_valid):
           specs._OUT_OF_BOUNDS % (spec.minimum, value, spec.maximum)):
         spec.validate(value)
 
+  @parameterized.parameters(
+      dict(shape=(), minimum=0., maximum=1., value=np.nan),
+      dict(shape=(2,), minimum=-1., maximum=1., value=[0., np.nan]),
+      dict(shape=(2,), minimum=[-1., 0.], maximum=[0., 1.],
+           value=[np.nan, 0.5]),
+      dict(shape=(), minimum=-np.inf, maximum=np.inf, value=np.nan),
+  )
+  def testValidateRejectsNaN(self, shape, minimum, maximum, value):
+    spec = specs.BoundedArray(shape, np.float64, minimum, maximum)
+    value = np.asarray(value, dtype=spec.dtype)
+    with self.assertRaisesWithLiteralMatch(
+        ValueError, specs._OUT_OF_BOUNDS % (spec.minimum, value, spec.maximum)):
+      spec.validate(value)
+
   def testValidateReturnsValue(self):
     spec = specs.BoundedArray([1], np.int32, minimum=0, maximum=1)
     validated_value = spec.validate(np.array([0], dtype=np.int32))
diff --git a/dm_env/test_utils_test.py b/dm_env/test_utils_test.py
index 628c73c..e703784 100644
--- a/dm_env/test_utils_test.py
+++ b/dm_env/test_utils_test.py
@@ -200,6 +200,21 @@ def wrapped_method(
     ]
 )
 
+TestNaNDiscount = _make_test_case_with_expected_failures(
+    name='TestNaNDiscount',
+    timestep_sequence=[
+        FIRST,
+        MID._replace(discount=np.nan),  # NaN is not in the interval [0, 1].
+        MID,
+        LAST,
+    ],
+    expected_failures=[
+        ('test_step_after_reset', ValueError),
+        ('test_step_on_fresh_environment', ValueError),
+        ('test_longer_action_sequence', ValueError),
+    ]
+)
+
 TestInvalidObservation = _make_test_case_with_expected_failures(
     name='TestInvalidObservation',
     timestep_sequence=[