Merge pull request #3 from rowanmcallister/patch-1

Update train.py
2020-10-20 13:12:01 -07:00 · 2020-10-20 13:12:01 -07:00 · 15008bbd3b
commit 15008bbd3b
parent 0fc1b8bd37 3f5320c78f
1 changed files with 1 additions and 1 deletions
--- a/train.py
+++ b/train.py
@ -318,7 +318,7 @@ def main():
        )

    # stack several consecutive frames together
-    if args.encoder_type == 'pixel':
+    if args.encoder_type.startswith('pixel'):
        env = utils.FrameStack(env, k=args.frame_stack)
        eval_env = utils.FrameStack(eval_env, k=args.frame_stack)