+elif args.task == "escape":
+ task = tasks.Escape(
+ nb_train_samples=args.nb_train_samples,
+ nb_test_samples=args.nb_test_samples,
+ batch_size=args.batch_size,
+ height=args.escape_height,
+ width=args.escape_width,
+ T=args.escape_T,
+ nb_walls=args.escape_nb_walls,
+ logger=log_string,
+ device=device,
+ )
+