Update train.py

Stack frames for CARLA pixel encoders too
2020-10-20 13:10:49 -07:00 · 2020-10-20 13:10:49 -07:00 · 3f5320c78f
commit 3f5320c78f
parent 0fc1b8bd37
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)