fix: save full rl checkpoint & merge comet experiments into one

Author: levje

TrackToLearn/algorithms/sac_auto.py

@@ -9,7 +9,7 @@
 from TrackToLearn.algorithms.sac import SAC
 from TrackToLearn.algorithms.shared.offpolicy import SACActorCritic
 from TrackToLearn.algorithms.shared.replay import OffPolicyReplayBuffer
-from TrackToLearn.utils.torch_utils import get_device
+from TrackToLearn.utils.torch_utils import get_device, gradients_norm
 from TrackToLearn.algorithms.shared.kl import AdaptiveKLController, FixedKLController
 
 LOG_STD_MAX = 2
@@ -184,6 +184,10 @@ def load_checkpoint(self, checkpoint_file: str):
         self.actor_optimizer.load_state_dict(checkpoint['actor_optimizer'])
         self.critic_optimizer.load_state_dict(checkpoint['critic_optimizer'])
         self.alpha_optimizer.load_state_dict(checkpoint['alpha_optimizer'])
+        if checkpoint.get('replay_buffer', None) is not None:
+            self.replay_buffer.load_state_dict(checkpoint['replay_buffer'])
+        if checkpoint.get('log_alpha', None) is not None:
+            self.log_alpha = checkpoint['log_alpha']
 
     def save_checkpoint(self, checkpoint_file: str, **extra_info):
         """
@@ -200,6 +204,8 @@ def save_checkpoint(self, checkpoint_file: str, **extra_info):
             'actor_optimizer': self.actor_optimizer.state_dict(),
             'critic_optimizer': self.critic_optimizer.state_dict(),
             'alpha_optimizer': self.alpha_optimizer.state_dict(),
+            'replay_buffer': self.replay_buffer.state_dict(),
+            'log_alpha': self.log_alpha,
             **extra_info
         }
 
@@ -273,18 +279,6 @@ def update(
         # Total critic loss
         critic_loss = loss_q1 + loss_q2
 
-        losses = {
-            # 'actor_loss': actor_loss.detach(),
-            # 'alpha_loss': alpha_loss.detach(),
-            # 'critic_loss': critic_loss.detach(),
-            # 'loss_q1': loss_q1.detach(),
-            # 'loss_q2': loss_q2.detach(),
-            # 'entropy': alpha.detach(),
-            # 'Q1': current_Q1.mean().detach(),
-            # 'Q2': current_Q2.mean().detach(),
-            # 'backup': backup.mean().detach(),
-        }
-
         # Optimize the temperature
         self.alpha_optimizer.zero_grad()
         alpha_loss.backward()
@@ -313,7 +307,30 @@ def update(
             self.target.actor.parameters()
         ):
             target_param.data.copy_(
-                self.tau * param.data + (1 - self.tau) * target_param.data
-            )
+                self.tau * param.data + (1 - self.tau) * target_param.data)
+
+        # Compute the norm of the gradients to plot.
+        alpha_norm = self.log_alpha.grad.norm(2).cpu().detach().numpy()
+        critic_norm = gradients_norm(self.agent.critic)
+        actor_norm = gradients_norm(self.agent.actor)
+
+        # print("alpha_norm: ", type(alpha_norm))
+        # print("critic_norm: ", type(critic_norm))
+        # print("actor_norm: ", type(actor_norm))
+
+        losses = {
+            # 'actor_loss': actor_loss.detach(),
+            # 'alpha_loss': alpha_loss.detach(),
+            # 'critic_loss': critic_loss.detach(),
+            # 'loss_q1': loss_q1.detach(),
+            # 'loss_q2': loss_q2.detach(),
+            # 'entropy': alpha.detach(),
+            # 'Q1': current_Q1.mean().detach(),
+            # 'Q2': current_Q2.mean().detach(),
+            # 'backup': backup.mean().detach(),
+            "alpha_norm": alpha_norm,
+            "critic_norm": critic_norm,
+            "actor_norm": actor_norm,
+        }
 
         return losses

TrackToLearn/algorithms/shared/replay.py

@@ -10,6 +10,7 @@
 from TrackToLearn.utils.torch_utils import get_device, get_device_str
 
 device = get_device()
+rb_type = torch.float32
 
 class OffPolicyReplayBuffer(object):
     """ Replay buffer to store transitions. Implemented in a "ring-buffer"
@@ -40,16 +41,17 @@ def __init__(
 
 
         self.state = torch.zeros(
-            (self.max_size, state_dim), dtype=torch.float32)
+            (self.max_size, state_dim), dtype=rb_type)
         self.action = torch.zeros(
-            (self.max_size, action_dim), dtype=torch.float32)
+            (self.max_size, action_dim), dtype=rb_type)
         self.next_state = torch.zeros(
-            (self.max_size, state_dim), dtype=torch.float32)
+            (self.max_size, state_dim), dtype=rb_type)
         self.reward = torch.zeros(
-            (self.max_size, 1), dtype=torch.float32)
+            (self.max_size, 1), dtype=rb_type)
         self.not_done = torch.zeros(
-            (self.max_size, 1), dtype=torch.float32)
-
+            (self.max_size, 1), dtype=rb_type)
+    
+    def _pin_to_memory(self):
         if get_device_str() == "cuda":
             self.state = self.state.pin_memory()
             self.action = self.action.pin_memory()
@@ -162,6 +164,28 @@ def load_from_file(self, path):
         """
         pass
 
+    def state_dict(self):
+        size = self.size
+        return {
+            "state": self.state[:size],
+            "action": self.action[:size],
+            "next_state": self.next_state[:size],
+            "reward": self.reward[:size],
+            "not_done": self.not_done[:size],
+            "ptr": self.ptr,
+            "size": self.size
+        }
+    
+    def load_state_dict(self, state_dict):
+        self.size = state_dict["size"]
+        self.ptr = state_dict["ptr"]
+
+        self.state[:self.size] = state_dict["state"]
+        self.action[:self.size] = state_dict["action"]
+        self.next_state[:self.size] = state_dict["next_state"]
+        self.reward[:self.size] = state_dict["reward"]
+        self.not_done[:self.size] = state_dict["not_done"]
+
 class OnPolicyReplayBuffer(object):
     """ Replay buffer to store transitions. Efficiency could probably be
     improved.

TrackToLearn/experiment/experiment.py

@@ -83,7 +83,7 @@ def setup_comet(self, prefix=''):
         """
         # The comet object that will handle monitors
         self.comet_monitor = CometMonitor(
-            self.comet_experiment, self.name, self.experiment_path,
+            self.comet_experiment, self.experiment_path,
             prefix, use_comet=self.use_comet)
         print(self.hyperparameters)
         self.comet_monitor.log_parameters(self.hyperparameters)

TrackToLearn/runners/tractoracle_predict.py

@@ -67,12 +67,12 @@ def test(self):
             parse_args=False,
             auto_metric_logging=False,
             disabled=True)
+        oracle_experiment.set_name(self.id)
 
         print("Done.")
 
         oracle_trainer = OracleTrainer(
             oracle_experiment,
-            self.id,
             root_dir,
             self.oracle_train_steps,
             enable_checkpointing=True,

TrackToLearn/trainers/oracle/oracle_monitor.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 from TrackToLearn.oracles.transformer_oracle import LightningLikeModule
-from TrackToLearn.trainers.oracle.oracle_monitor import OracleMonitor
+from TrackToLearn.utils.comet_monitor import OracleMonitor
 from TrackToLearn.utils.torch_utils import get_device
 from TrackToLearn.algorithms.shared.utils import \
     (add_item_to_means, mean_losses, add_losses, get_mean_item)
@@ -68,7 +68,6 @@ def reset(self):
 class OracleTrainer(object):
     def __init__(self,
                  experiment,
-                 experiment_id,
                  saving_path,
                  max_epochs,
                  use_comet=True,
@@ -81,7 +80,6 @@ def __init__(self,
                  metrics_prefix=None,
                  ):
         self.experiment = experiment
-        self.experiment_id = experiment_id
         self.saving_path = saving_path
 
         self.auto_checkpointing_enabled = enable_auto_checkpointing
@@ -95,7 +93,6 @@ def __init__(self,
         self.hooks_manager = HooksManager(OracleHookEvent)
         self.oracle_monitor = OracleMonitor(
             experiment=self.experiment,
-            experiment_id=self.experiment_id,
             use_comet=use_comet,
             metrics_prefix=metrics_prefix
         )
@@ -109,7 +106,6 @@ def save_hyperparameters(self):
 
         hyperparameters = self.oracle_model.hyperparameters
         hyperparameters.update({
-            'experiment_id': self.experiment_id,
             'saving_path': self.saving_path,
             'max_epochs': self.max_epochs,
             'val_interval': self.val_interval,

TrackToLearn/trainers/oracle/oracle_trainer.py

@@ -39,8 +39,7 @@ def __init__(
         self,
         rlhf_train_dto: dict,
         trainer_cls: TrackToLearnTraining,
-        agent_experiment: CometExperiment = None,
-        oracle_experiment: CometExperiment = None
+        comet_experiment: CometExperiment = None
     ):
         # Only load the parameters from the parent instead of calling
         # the full constructor twice. (As we call it for the agent_trainer
@@ -73,8 +72,14 @@ def __init__(
         self.agent_train_steps = rlhf_train_dto['agent_train_steps']
         self.num_workers = rlhf_train_dto['num_workers']
         self.rlhf_inter_npv = rlhf_train_dto['rlhf_inter_npv']
+        
         self.disable_oracle_training = rlhf_train_dto.get(
             'disable_oracle_training', False)
+        if self.disable_oracle_training:
+            LOGGER.warning("Oracle training is disabled. The dataset will "
+                           "be augmented to evaluate the oracles during the "
+                           "agent's training.")
+
         self.batch_size = rlhf_train_dto['batch_size']
         self.oracle_batch_size = rlhf_train_dto['oracle_batch_size']
         grad_accumulation_steps = rlhf_train_dto.get(
@@ -91,15 +96,15 @@ def __init__(
 
         ################################################
         # Start by initializing the agent trainer.     #
-        if agent_experiment is None:
-            agent_experiment = CometExperiment(project_name=self.experiment,
+        if comet_experiment is None:
+            comet_experiment = CometExperiment(project_name=self.experiment,
                                                workspace=rlhf_train_dto['workspace'], parse_args=False,
                                                auto_metric_logging=False,
                                                disabled=not self.use_comet)
 
-            agent_experiment.set_name(self.name)
+        comet_experiment.set_name(self.name)
 
-        self.agent_trainer: TrackToLearnTraining = trainer_cls(rlhf_train_dto, agent_experiment)
+        self.agent_trainer: TrackToLearnTraining = trainer_cls(rlhf_train_dto, comet_experiment)
         _ = self.agent_trainer.setup_environment_and_info()
         self.get_alg = self.agent_trainer.get_alg
 
@@ -110,17 +115,6 @@ def __init__(
 
         ################################################
         # Continue by initializing the oracle trainer. #
-        # Need this to avoid erasing the RL agent's experiment
-        # when creating a new one.
-        if oracle_experiment is None:
-            comet_ml.config.set_global_experiment(None)
-            oracle_experiment = CometExperiment(project_name="TractOracleRLHF",
-                                                       workspace=rlhf_train_dto['workspace'], parse_args=False,
-                                                       auto_metric_logging=False,
-                                                       disabled=not self.use_comet)
-
-            oracle_experiment_id = '-'.join([self.experiment, self.name])
-
         dataset_to_augment = rlhf_train_dto.get('dataset_to_augment', None)
         self.dataset_manager = StreamlineDatasetManager(saving_path=self.oracle_training_dir,
                                                         dataset_to_augment_path=dataset_to_augment,
@@ -130,8 +124,7 @@ def __init__(
         # because we will want to save the checkpoints only when we improve the 
         # total agent. We manually checkpoint those oracles instead.
         self.oracle_reward_trainer = OracleTrainer(
-            oracle_experiment,
-            oracle_experiment_id,
+            comet_experiment,
             self.oracle_training_dir,
             self.oracle_train_steps,
             enable_auto_checkpointing=False,
@@ -143,8 +136,7 @@ def __init__(
         )
 
         self.oracle_crit_trainer = OracleTrainer(
-            oracle_experiment,
-            oracle_experiment_id,
+            comet_experiment,
             self.oracle_training_dir,
             self.oracle_train_steps,
             enable_auto_checkpointing=False,
@@ -293,10 +285,7 @@ def rl_train(
         while i < max_ep: 
             self.start_finetuning_epoch(i, do_warmup)
 
-            if self.disable_oracle_training:
-                LOGGER.info("Oracle training is disabled. Only the agent will be trained and the dataset will not be augmented.\n",
-                                 "This is equivalent to just training the agent for an additional {} ({} x {}) epochs.".format(self.agent_train_steps*max_ep, max_ep, self.agent_train_steps))
-            elif not do_warmup:
+            if not do_warmup:
                 total_added = 0
 
                 with tqdm(total=self.nb_new_streamlines_per_iter,
@@ -353,9 +342,9 @@ def rl_train(
                     prettier_dict(data_stats, title="Dataset stats (iter {})".format(i)))
 
                 # Train reward model
-                LOGGER.info("Training reward model...")
-                self.train_reward()
-                self.train_stopping_criterion()
+                if not self.disable_oracle_training:
+                    self.train_reward()
+                    self.train_stopping_criterion()
 
             # Train the RL agent
             agent_nb_steps = self.agent_train_steps if not do_warmup else self.warmup_agent_steps
@@ -369,14 +358,16 @@ def rl_train(
                                         max_ep=agent_nb_steps,
                                         starting_ep=current_ep,
                                         save_model_dir=self.model_dir,
-                                        test_before_training=False
+                                        test_before_training=do_warmup or i == 0
                                         )
-            current_ep += self.agent_train_steps
 
             self.end_finetuning_epoch(i, do_warmup)
 
-            if not do_warmup:
+            if do_warmup:
+                current_ep += self.warmup_agent_steps
+            else:
                 self.backuper.backup(step=i)
+                current_ep += self.agent_train_steps
                 i += 1
             do_warmup = False