Repository: tum-pbs/apebench

Issue:
Add option to cut bptt in unrolled training methodologies
Requires forwarding an option to `Trainax`.

Discussion:

Patch:
diff --git a/apebench/_base_scenario.py b/apebench/_base_scenario.py
index 63c0e70..5b30068 100644
--- a/apebench/_base_scenario.py
+++ b/apebench/_base_scenario.py
@@ -349,18 +349,30 @@ class BaseScenario(eqx.Module, ABC):
             )
         elif train_args[0].lower() == "sup":
             num_rollout_steps = int(train_args[1])
+            if len(train_args) > 2:
+                cut_bptt = train_args[2].lower() == "true"
+            else:
+                cut_bptt = False
             trainer = tx.trainer.SupervisedTrainer(
                 train_trjs,
                 optimizer=optimizer,
                 num_training_steps=self.num_training_steps,
                 batch_size=self.batch_size,
                 num_rollout_steps=num_rollout_steps,
-                cut_bptt=False,
+                cut_bptt=cut_bptt,
                 time_level_weights=None,
                 callback_fn=callback_fn,
             )
         elif train_args[0].lower() == "div":
             num_rollout_steps = int(train_args[1])
+            if len(train_args) > 2:
+                cut_bptt = train_args[2].lower() == "true"
+            else:
+                cut_bptt = False
+            if len(train_args) > 3:
+                cut_div_chain = train_args[3].lower() == "true"
+            else:
+                cut_div_chain = False
             trainer = tx.trainer.DivertedChainBranchOneTrainer(
                 train_trjs,
                 ref_stepper=ref_stepper,
@@ -368,7 +380,8 @@ class BaseScenario(eqx.Module, ABC):
                 num_training_steps=self.num_training_steps,
                 batch_size=self.batch_size,
                 num_rollout_steps=num_rollout_steps,
-                cut_bptt=False,
+                cut_bptt=cut_bptt,
+                cut_div_chain=cut_div_chain,
                 time_level_weights=None,
                 callback_fn=callback_fn,
             )


Test patch:
diff --git a/tests/test_trainer_parsing.py b/tests/test_trainer_parsing.py
new file mode 100644
index 0000000..3d74dc7
--- /dev/null
+++ b/tests/test_trainer_parsing.py
@@ -0,0 +1,65 @@
+import apebench
+
+
+def test_trainer_parsing():
+    advection_scenario = apebench.scenarios.difficulty.Advection()
+
+    trainer = advection_scenario.get_trainer(
+        train_config="one",
+    )
+
+    assert trainer.loss_configuration.num_rollout_steps == 1
+    del trainer
+
+    trainer = advection_scenario.get_trainer(train_config="sup;3")
+
+    assert trainer.loss_configuration.num_rollout_steps == 3
+    assert trainer.loss_configuration.cut_bptt is False
+    del trainer
+
+    trainer = advection_scenario.get_trainer(train_config="sup;3;True")
+
+    assert trainer.loss_configuration.num_rollout_steps == 3
+    assert trainer.loss_configuration.cut_bptt is True
+    del trainer
+
+    trainer = advection_scenario.get_trainer(train_config="sup;3;False")
+
+    assert trainer.loss_configuration.num_rollout_steps == 3
+    assert trainer.loss_configuration.cut_bptt is False
+    del trainer
+
+    trainer = advection_scenario.get_trainer(train_config="div;3")
+
+    assert trainer.loss_configuration.num_rollout_steps == 3
+    assert trainer.loss_configuration.cut_bptt is False
+    assert trainer.loss_configuration.cut_div_chain is False
+    del trainer
+
+    trainer = advection_scenario.get_trainer(train_config="div;3;True")
+
+    assert trainer.loss_configuration.num_rollout_steps == 3
+    assert trainer.loss_configuration.cut_bptt is True
+    assert trainer.loss_configuration.cut_div_chain is False
+    del trainer
+
+    trainer = advection_scenario.get_trainer(train_config="div;3;False")
+
+    assert trainer.loss_configuration.num_rollout_steps == 3
+    assert trainer.loss_configuration.cut_bptt is False
+    assert trainer.loss_configuration.cut_div_chain is False
+    del trainer
+
+    trainer = advection_scenario.get_trainer(train_config="div;3;True;True")
+
+    assert trainer.loss_configuration.num_rollout_steps == 3
+    assert trainer.loss_configuration.cut_bptt is True
+    assert trainer.loss_configuration.cut_div_chain is True
+    del trainer
+
+    trainer = advection_scenario.get_trainer(train_config="div;3;False;True")
+
+    assert trainer.loss_configuration.num_rollout_steps == 3
+    assert trainer.loss_configuration.cut_bptt is False
+    assert trainer.loss_configuration.cut_div_chain is True
+    del trainer