-def create_data_and_processors(nb_train_samples, nb_test_samples, nb_epochs=10):
- steps = [True] + [False] * 30 + [True]
- train_input = generate_episodes(nb_train_samples, steps)
- test_input = generate_episodes(nb_test_samples, steps)
+def create_data_and_processors(
+ nb_train_samples,
+ nb_test_samples,
+ mode,
+ nb_steps,
+ nb_epochs=10,
+ device=torch.device("cpu"),
+):
+ assert mode in ["first_last"]
+
+ if mode == "first_last":
+ steps = [True] + [False] * (nb_steps + 1) + [True]
+
+ train_input, train_actions = generate_episodes(nb_train_samples, steps)
+ train_input, train_actions = train_input.to(device), train_actions.to(device)
+ test_input, test_actions = generate_episodes(nb_test_samples, steps)
+ test_input, test_actions = test_input.to(device), test_actions.to(device)