fix: warn on implicit zero-fill and preserve empty dicts in Batch

CHANGELOG.md

@@ -1,3 +1,9 @@
+# Unreleased
+
+**Fixes:**
+* `data`:
+  * `Batch`: Warn on implicit zero-fill when stacking batches with mismatched keys, and preserve empty dicts instead of silently dropping them #1296
+
 # Release 2.0.1 (2026-04-02)
 
 This is a maintenance release.

test/base/test_batch.py

@@ -1,5 +1,6 @@
 import copy
 import pickle
+import warnings
 from itertools import starmap
 from typing import Any, cast
 
@@ -956,3 +957,130 @@ def test_len_batch_with_dist() -> None:
         with pytest.raises(TypeError):
             # scalar batches have no len
             len(batch_with_dist[0])
+
+
+class TestBatchNoneAndEmptyHandling:
+    """Tests for issues #1088 (None replaced with 0) and #1089 (empty dict dropped)."""
+
+    @staticmethod
+    def test_empty_dict_preserves_length() -> None:
+        """Issue #1089: mixing empty and non-empty dicts should preserve length."""
+        b = Batch(info=[{"a": 1}, {}])
+        assert len(b.info) == 2
+        assert np.array_equal(b.info.a, np.array([1, 0]))
+
+    @staticmethod
+    def test_empty_dict_at_beginning() -> None:
+        """Issue #1089: empty dict at index 0 should not be dropped."""
+        b = Batch(info=[{}, {"a": 1}])
+        assert len(b.info) == 2
+        assert np.array_equal(b.info.a, np.array([0, 1]))
+
+    @staticmethod
+    def test_multiple_empty_dicts() -> None:
+        """Issue #1089: multiple empty dicts interspersed should all be preserved."""
+        b = Batch(info=[{}, {"a": 1}, {}, {"a": 2}, {}])
+        assert len(b.info) == 5
+        assert np.array_equal(b.info.a, np.array([0, 1, 0, 2, 0]))
+
+    @staticmethod
+    def test_all_empty_dicts_stack() -> None:
+        """Stacking all-empty dicts/Batches should still return an empty Batch."""
+        b = Batch.stack([Batch(), Batch(), Batch()])
+        assert len(b.get_keys()) == 0
+
+    @staticmethod
+    def test_all_empty_dicts_in_list() -> None:
+        """A list of all empty dicts should produce an empty Batch."""
+        b = Batch(info=[{}, {}, {}])
+        assert len(b.info.get_keys()) == 0
+
+    @staticmethod
+    def test_empty_dict_with_nested_batch() -> None:
+        """Issue #1089: empty dicts mixed with nested structures."""
+        b = Batch(info=[{"inner": {"x": 1}}, {}])
+        assert len(b.info) == 2
+        assert len(b.info.inner) == 2
+        assert np.array_equal(b.info.inner.x, np.array([1, 0]))
+
+    @staticmethod
+    def test_missing_key_warns_for_numeric_in_setitem() -> None:
+        """Issue #1088: __setitem__ should warn when filling 0 for missing numeric key."""
+        b = Batch(a=[1, 2], env_num=[3, 4])
+        with pytest.warns(
+            UserWarning,
+            match=r"Key 'env_num' is not found in the value Batch",
+        ):
+            b[1] = Batch(a=99)
+        assert b.env_num[1] == 0
+        assert b.a[1] == 99
+
+    @staticmethod
+    def test_missing_key_warns_for_numeric_in_stack() -> None:
+        """Issue #1088: stack_ should warn when filling 0 for missing numeric key."""
+        with pytest.warns(
+            UserWarning,
+            match=r"Key 'env_num' is not present in all batches during stacking",
+        ):
+            b = Batch(info=[{"a": 1, "env_num": 3}, {"a": 2}])
+        assert np.array_equal(b.info.a, np.array([1, 2]))
+        assert np.array_equal(b.info.env_num, np.array([3, 0]))
+
+    @staticmethod
+    def test_missing_key_no_warn_for_object_type() -> None:
+        """Non-numeric types (object arrays) use None and should not warn."""
+        b = Batch(a=["hello", "world"], b=["x", "y"])
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            b[0] = Batch(a="replaced")
+        assert b.a[0] == "replaced"
+        assert b.b[0] is None
+
+    @staticmethod
+    def test_missing_key_batch_type_fills_empty() -> None:
+        """Batch-type values use empty Batch() for missing key at the outer level."""
+        b = Batch(a=[1, 2], sub=Batch())
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            # 'sub' is an empty Batch, so assigning with missing 'sub' triggers
+            # the Batch branch (not numeric), which fills with Batch() without warning
+            b[0] = Batch(a=99)
+        assert b.a[0] == 99
+
+    @staticmethod
+    def test_missing_key_torch_tensor_warns() -> None:
+        """Issue #1088: torch tensors should also warn when filling 0."""
+        b = Batch(a=torch.tensor([1, 2]), extra=torch.tensor([3, 4]))
+        with pytest.warns(
+            UserWarning,
+            match=r"Key 'extra' is not found in the value Batch",
+        ):
+            b[0] = Batch(a=torch.tensor(99))
+        assert b.extra[0] == 0
+
+    @staticmethod
+    def test_stack_partial_keys_preserves_values() -> None:
+        """Partial keys during stack should correctly set present values."""
+        with pytest.warns(UserWarning):
+            b = Batch.stack(
+                [
+                    Batch(a=1, b=2),
+                    Batch(a=3),
+                ]
+            )
+        assert np.array_equal(b.a, np.array([1, 3]))
+        assert np.array_equal(b.b, np.array([2, 0]))
+
+    @staticmethod
+    def test_hasnull_detects_object_none_but_not_numeric_zero() -> None:
+        """Verify that hasnull works correctly: detects None in object arrays,
+        but cannot detect 0-filled numeric missing values (documenting current behavior).
+        """
+        # Object array with None is detectable
+        b_obj = Batch(a=[1, 2], b=["x", None])
+        assert b_obj.hasnull() is True
+
+        # Numeric array with 0-filled missing is NOT detectable by hasnull
+        with pytest.warns(UserWarning):
+            b_num = Batch.stack([Batch(a=1, env_num=3), Batch(a=2)])
+        assert b_num.hasnull() is False  # 0 is not null

tianshou/data/batch.py

@@ -227,6 +227,51 @@ def _parse_value(obj: Any) -> Union["Batch", np.ndarray, torch.Tensor] | None:
     return obj
 
 
+def _validate_and_convert_batches(
+    batches: Sequence,
+) -> tuple[list["Batch"], bool]:
+    """Convert input batches to Batch objects, preserving empty entries (fixes #1089)."""
+    batch_list: list[Batch] = []
+    has_any_nonempty = False
+    for batch in batches:
+        if isinstance(batch, dict):
+            batch_list.append(Batch(batch) if len(batch) > 0 else Batch())
+            if len(batch) > 0:
+                has_any_nonempty = True
+        elif isinstance(batch, Batch):
+            batch_list.append(batch)
+            if len(batch.get_keys()) != 0:
+                has_any_nonempty = True
+        else:
+            raise ValueError(f"Cannot concatenate {type(batch)} in Batch.stack_")
+    return batch_list, has_any_nonempty
+
+
+def _warn_numeric_zero_fill(
+    data: dict[str, Any],
+    indices_missing_keys: dict[str, list[int]],
+) -> None:
+    """Emit a warning for keys where missing entries were filled with 0."""
+    for key, missing_indices in indices_missing_keys.items():
+        if not missing_indices:
+            continue
+        val = data.get(key)
+        if val is None:
+            continue
+        is_numeric = isinstance(val, torch.Tensor) or (
+            isinstance(val, np.ndarray) and issubclass(val.dtype.type, np.bool_ | np.number)
+        )
+        if is_numeric:
+            warnings.warn(
+                f"Key '{key}' is not present in all batches during "
+                f"stacking (missing at indices {missing_indices}). "
+                f"Filling missing entries with 0 for numeric type "
+                f"({type(val).__name__}), which may mask truly missing "
+                f"values. Consider using None or np.nan to represent "
+                f"missing data explicitly.",
+            )
+
+
 def alloc_by_keys_diff(
     meta: "BatchProtocol",
     batch: "BatchProtocol",
@@ -629,11 +674,9 @@ class Batch(BatchProtocol):
 
     def __init__(
         self,
-        batch_dict: dict
-        | BatchProtocol
-        | Sequence[dict | BatchProtocol]
-        | np.ndarray
-        | None = None,
+        batch_dict: (
+            dict | BatchProtocol | Sequence[dict | BatchProtocol] | np.ndarray | None
+        ) = None,
         copy: bool = False,
         **kwargs: Any,
     ) -> None:
@@ -788,6 +831,12 @@ def __setitem__(self, index: str | IndexType, value: Any) -> None:
                 elif isinstance(val, torch.Tensor) or (
                     isinstance(val, np.ndarray) and issubclass(val.dtype.type, np.bool_ | np.number)
                 ):
+                    warnings.warn(
+                        f"Key '{key}' is not found in the value Batch during "
+                        f"item assignment. Filling with 0 for numeric type "
+                        f"({type(val).__name__}), which may mask missing values. "
+                        f"Consider using None or np.nan to represent missing data.",
+                    )
                     self.__dict__[key][index] = 0
                 else:
                     self.__dict__[key][index] = None
@@ -1039,18 +1088,8 @@ def cat(batches: Sequence[dict | TBatch]) -> TBatch:
         return batch  # type: ignore
 
     def stack_(self, batches: Sequence[dict | BatchProtocol], axis: int = 0) -> None:
-        # check input format
-        batch_list = []
-        for batch in batches:
-            if isinstance(batch, dict):
-                if len(batch) > 0:
-                    batch_list.append(Batch(batch))
-            elif isinstance(batch, Batch):
-                if len(batch.get_keys()) != 0:
-                    batch_list.append(batch)
-            else:
-                raise ValueError(f"Cannot concatenate {type(batch)} in Batch.stack_")
-        if len(batch_list) == 0:
+        batch_list, has_any_nonempty = _validate_and_convert_batches(batches)
+        if not has_any_nonempty:
             return
         batches = batch_list
         if len(self.get_keys()) != 0:
@@ -1098,22 +1137,42 @@ def stack_(self, batches: Sequence[dict | BatchProtocol], axis: int = 0) -> None
         for key in keys_reserve:
             # reserved keys
             self.__dict__[key] = Batch()
+        if keys_partial:
+            indices_missing_keys: dict[str, list[int]] = {key: [] for key in keys_partial}
         for key in keys_partial:
+            # Collect all values for this partial key; missing entries use Batch()
+            key_values: list[Batch | Any] = []
+            has_nested_batch = False
             for i, batch in enumerate(batches):
                 if key not in batch.__dict__:
+                    indices_missing_keys[key].append(i)
+                    key_values.append(Batch())
                     continue
-                value = batch.get(key)
-                # TODO: fix code/annotations s.t. the ignores can be removed
-                if (
-                    isinstance(value, Batch)  # type: ignore
-                    and len(value.get_keys()) == 0  # type: ignore
-                ):
-                    continue  # type: ignore
-                try:
-                    self.__dict__[key][i] = value
-                except KeyError:
-                    self.__dict__[key] = create_value(value, len(batches))
-                    self.__dict__[key][i] = value
+                val = batch.get(key)
+                if isinstance(val, Batch) and len(val.get_keys()) == 0:
+                    indices_missing_keys[key].append(i)
+                    key_values.append(Batch())
+                    continue
+                if isinstance(val, Batch | dict):
+                    has_nested_batch = True
+                key_values.append(val)
+            # For nested Batch/dict values, use recursive Batch.stack to
+            # handle differing nested keys correctly (fixes regression
+            # where Batch.stack([{"info": {"a": 1}}, {}, {"info": {"b": 2}}])
+            # would raise ValueError).
+            if has_nested_batch:
+                self.__dict__[key] = Batch.stack(key_values, axis)
+            else:
+                for i, val in enumerate(key_values):
+                    if isinstance(val, Batch) and len(val.get_keys()) == 0:
+                        continue
+                    try:
+                        self.__dict__[key][i] = val
+                    except KeyError:
+                        self.__dict__[key] = create_value(val, len(batches))
+                        self.__dict__[key][i] = val
+        if keys_partial:
+            _warn_numeric_zero_fill(self.__dict__, indices_missing_keys)
 
     @staticmethod
     def stack(batches: Sequence[dict | TBatch], axis: int = 0) -> TBatch: