diff --git a/CHANGELOG.md b/CHANGELOG.md index b8996d9fd..7d249511f 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/test/base/test_batch.py b/test/base/test_batch.py index 3e9a9fb63..36a73f9c5 100644 --- a/test/base/test_batch.py +++ b/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 diff --git a/tianshou/data/batch.py b/tianshou/data/batch.py index e688ecd63..6dee849ce 100644 --- a/tianshou/data/batch.py +++ b/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: