diff --git a/train.py b/train.py
index 02769f7..047f1ee 100644
--- a/train.py
+++ b/train.py
@@ -146,7 +146,10 @@ def main():
         from_pixels=(args.encoder_type == 'pixel'),
         height=args.image_size,
         width=args.image_size,
-        frame_skip=args.action_repeat
+        frame_skip=args.action_repeat,
+        img_source=args.img_source,
+        resource_files=args.resource_files,
+        total_frames=args.total_frames
     )
     env.seed(args.seed)