Fix batch add for async

test/base/test_buffer.py

@@ -1368,26 +1368,31 @@ def test_custom_key():
             'done':
             np.array([False]),
             'returns':
-            np.array([74.70343082])
+            np.array([74.70343082]),
+            'info': 
+            Batch(),
+            'policy':
+            Batch(),
         }
     )
     buffer_size = len(batch.rew)
     buffer = ReplayBuffer(buffer_size)
     buffer.add(batch)
     sampled_batch, _ = buffer.sample(1)
     # Check if they have the same keys
-    assert set(batch.__dict__.keys()) == set(sampled_batch.__dict__.keys()), \
+    assert set(batch.keys()) == set(sampled_batch.keys()), \
         "Batches have different keys: {} and {}".format(
-            set(batch.__dict__.keys()), set(sampled_batch.__dict__.keys()))
+            set(batch.keys()), set(sampled_batch.keys()))
     # Compare the values for each key
-    for key in batch.__dict__.keys():
+    for key in batch.keys():
         if isinstance(batch.__dict__[key], np.ndarray
                       ) and isinstance(sampled_batch.__dict__[key], np.ndarray):
             assert np.allclose(batch.__dict__[key], sampled_batch.__dict__[key]), \
                 "Value mismatch for key: {}".format(key)
-        else:
-            assert batch.__dict__[key] == sampled_batch.__dict__[key], \
-                "Value mismatch for key: {}".format(key)
+        if isinstance(batch.__dict__[key], Batch
+                      ) and isinstance(sampled_batch.__dict__[key], Batch):
+            assert batch.__dict__[key].is_empty()
+            assert sampled_batch.__dict__[key].is_empty()
 
 
 if __name__ == '__main__':

tianshou/data/buffer/base.py

@@ -231,6 +231,10 @@ def add(
         episode_reward is 0.
         """
         # preprocess batch
+        new_batch = Batch()
+        for key in batch.keys():
+            new_batch.__dict__[key] = batch[key]
+        batch = new_batch
         batch.__dict__["done"] = np.logical_or(batch.terminated, batch.truncated)
         assert set(["obs", "act", "rew", "terminated", "truncated",
                     "done"]).issubset(batch.keys())