hongshaorou/dreamerv3-torch
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

make sure "is_first" is set 0 at beginning

Browse Source
This commit is contained in:
NM512 2023-07-22 21:08:53 +09:00
parent f07d843953
commit 03d91cb2c1
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tools.py
View File

@ -364,6 +364,8 @@ def sample_episodes(episodes, length, seed=0):
ret = { ret = {
k: v[index : min(index + length, total)] for k, v in episode.items() k: v[index : min(index + length, total)] for k, v in episode.items()
} }
if "is_first" in ret:
ret["is_first"][0] = True
else: else:
# 'is_first' comes after 'is_last' # 'is_first' comes after 'is_last'
index = 0 index = 0
@ -374,6 +376,8 @@ def sample_episodes(episodes, length, seed=0):
) )
for k, v in episode.items() for k, v in episode.items()
} }
if "is_first" in ret:
ret["is_first"][size] = True
size = len(next(iter(ret.values()))) size = len(next(iter(ret.values())))
yield ret yield ret