- height=args.picoclvr_height,
- width=args.picoclvr_width,
- logger=log_string,
- device=device,
- )
-
-elif args.task == "world":
- task = tasks.World(
- nb_train_samples=args.nb_train_samples,
- nb_test_samples=args.nb_test_samples,
- batch_size=args.batch_size,
- vqae_nb_epochs=args.world_vqae_nb_epochs,