hongshaorou/Tianshou
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix #56

Browse Source
This commit is contained in:
Trinkle23897 2020-05-29 08:03:37 +08:00
parent de556fd22d
commit d2b2fa87c0
2 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
test/base/test_batch.py
View File

@ -22,9 +22,10 @@ def test_batch():
def test_batch_over_batch(): def test_batch_over_batch():
batch = Batch(a=[3, 4, 5], b=[4, 5, 6]) batch = Batch(a=[3, 4, 5], b=[4, 5, 6])
batch2 = Batch(b=batch, c=[6, 7, 8]) batch2 = Batch(c=[6, 7, 8], b=batch)
batch2.b.b[-1] = 0 batch2.b.b[-1] = 0
print(batch2) print(batch2)
assert batch2.values()[-1] == batch2.c
assert batch2[-1].b.b == 0 assert batch2[-1].b.b == 0

4
tianshou/data/batch.py
View File

@ -131,6 +131,10 @@ class Batch(object):
return sorted([ return sorted([
i for i in self.__dict__ if i[0] != '_'] + list(self._meta)) i for i in self.__dict__ if i[0] != '_'] + list(self._meta))
def values(self) -> List[Any]:
"""Return self.values()."""
return [self[k] for k in self.keys()]
def get(self, k: str, d: Optional[Any] = None) -> Union['Batch', Any]: def get(self, k: str, d: Optional[Any] = None) -> Union['Batch', Any]:
"""Return self[k] if k in self else d. d defaults to None.""" """Return self[k] if k in self else d. d defaults to None."""
if k in self.__dict__ or k in self._meta: if k in self.__dict__ or k in self._meta: