patch multibinary, collect mode with custom dset

Author: jakegrigsby

amago/agent.py

@@ -289,8 +289,6 @@ def __init__(
         self.critics = critic_type(**ac_kwargs, num_critics=num_critics)
         self.target_critics = critic_type(**ac_kwargs, num_critics=num_critics)
         self.maximized_critics = critic_type(**ac_kwargs, num_critics=num_critics)
-        if self.multibinary:
-            ac_kwargs["cont_dist_kind"] = "multibinary"
         self.actor = actor_type(**ac_kwargs)
         self.target_actor = actor_type(**ac_kwargs)
         # full weight copy to targets

amago/cli_utils.py

@@ -464,7 +464,6 @@ def make_experiment_learn_only(experiment: amago.Experiment) -> amago.Experiment
     experiment.parallel_actors = 1
     experiment.always_save_latest = True
     experiment.always_load_latest = False
-    experiment.has_dset_edit_rights = True
     return experiment
 
 
@@ -488,6 +487,7 @@ def make_experiment_collect_only(experiment: amago.Experiment) -> amago.Experime
     experiment.epochs = max(experiment.epochs, 1_000_000)
     # do not delete anything from the collection process
     experiment.has_dset_edit_rights = False
+    experiment.init_dsets()
     return experiment
 
 

amago/nets/traj_encoders.py

@@ -532,6 +532,7 @@ def reset_hidden_state(
         hidden_state.reset(idxs=dones)
         return hidden_state
 
+    @torch.compile
     def forward(
         self,
         seq: torch.Tensor,