17 lines
311 B
Python
17 lines
311 B
Python
|
import numpy as np
|
||
|
|
|
||
|
|
from tianshou.data import Batch
|
||
|
|
|
||
|
|
|
||
|
|
def test_batch():
|
||
|
|
batch = Batch(obs=[0], np=np.zeros([3, 4]))
|
||
|
|
batch.update(obs=[1])
|
||
|
|
assert batch.obs == [1]
|
||
|
|
batch.append(batch)
|
||
|
|
assert batch.obs == [1, 1]
|
||
|
|
assert batch.np.shape == (6, 4)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
test_batch()
|