rxnn 0.2.29__py3-none-any.whl → 0.2.31__py3-none-any.whl

rxnn/training/base.py

@@ -8,6 +8,7 @@ import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel
 from typing import Callable
 from .callbacks import TrainerCallback
+from .ddp import get_os_ddp_config, distributed_value_mean
 
 
 class BaseTrainer(ABC):
@@ -91,8 +92,7 @@ class BaseTrainer(ABC):
             optimizer = self.optimizer
 
         if self.use_ddp:
-            rank = int(os.environ['RANK'])
-            world_size = int(os.environ['WORLD_SIZE'])
+            rank, world_size = get_os_ddp_config()
             dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
             self.model = DistributedDataParallel(self.model, device_ids=[self.device.index])
             train_sampler = torch.utils.data.DistributedSampler(dataset, shuffle=True)
@@ -218,10 +218,9 @@ class BaseTrainer(ABC):
         if self.validation_dataset:
             self.validation_steps = 0
             val_loss, val_metrics = self.validate(batch_size)
-            val_loss_tensor = torch.tensor(val_loss).to(self.device)
             if self.use_ddp:
-                dist.all_reduce(val_loss_tensor, op=dist.ReduceOp.SUM)
-                val_loss = val_loss_tensor.item() / dist.get_world_size()
+                val_loss = distributed_value_mean(val_loss, device=self.device)
+
             self.validation_metrics[epoch] = val_metrics
 
             if self.writer:

rxnn/training/bml.py

@@ -7,6 +7,7 @@ import torch.distributed as dist
 from ..transformers.models import ReactiveTransformerDecoder
 from ..training.base import BaseTrainer
 from .models import MLMTrainingModel, JointTrainingModel
+from .ddp import distributed_mean
 
 class MLMTrainer(BaseTrainer):
     def __init__(
@@ -96,8 +97,7 @@ class MLMTrainer(BaseTrainer):
         acc = (correct / total * 100) if total > 0 else torch.tensor(0.0).to(self.device)
         node_acc = acc.item()
         if self.use_ddp:
-            dist.all_reduce(acc, op=dist.ReduceOp.SUM)
-            acc = acc / dist.get_world_size()
+            acc = distributed_mean(acc)
 
         metrics = {
             'accuracy': acc.item(),
@@ -198,8 +198,7 @@ class AutoregressiveTrainer(BaseTrainer):
         acc = (correct / total * 100) if total > 0 else torch.tensor(0.0).to(self.device)
         node_acc = acc.item()
         if self.use_ddp:
-            dist.all_reduce(acc, op=dist.ReduceOp.SUM)
-            acc = acc / dist.get_world_size()
+            acc = distributed_mean(acc)
 
         metrics = {
             'accuracy': acc.item(),
@@ -347,14 +346,10 @@ class JointLMTrainer(BaseTrainer):
         node_mlm_acc = mlm_acc.item()
         node_alm_acc = alm_acc.item()
         if self.use_ddp:
-            dist.all_reduce(avg_dec_loss, op=dist.ReduceOp.SUM)
-            dist.all_reduce(avg_enc_loss, op=dist.ReduceOp.SUM)
-            dist.all_reduce(mlm_acc, op=dist.ReduceOp.SUM)
-            dist.all_reduce(alm_acc, op=dist.ReduceOp.SUM)
-            avg_dec_loss = avg_dec_loss / dist.get_world_size()
-            avg_enc_loss = avg_enc_loss / dist.get_world_size()
-            mlm_acc = mlm_acc / dist.get_world_size()
-            alm_acc = alm_acc / dist.get_world_size()
+            avg_dec_loss = distributed_mean(avg_dec_loss)
+            avg_enc_loss = distributed_mean(avg_enc_loss)
+            mlm_acc = distributed_mean(mlm_acc)
+            alm_acc = distributed_mean(alm_acc)
 
         metrics = {
             'accuracy': {

rxnn/training/callbacks.py

@@ -536,6 +536,9 @@ class MrlTrainerCallback:
     def on_reward(self, actor: nn.Module, reward: float, generated: str, reference: str, saved_data: str, eval_mode: bool) -> None:
         pass
 
+    def on_update_epoch_start(self, actor: nn.Module, critic: nn.Module, global_epoch: int, update_epoch: int) -> None:
+        pass
+
     def on_batch_updated(self, actor: nn.Module, epoch: int, step: int, policy_loss: float) -> None:
         pass
 
@@ -543,6 +546,9 @@ class MrlTrainerCallback:
                           critic_loss: float) -> None:
         pass
 
+    def on_update_epoch_end(self, actor: nn.Module, critic: nn.Module, global_epoch: int, update_epoch: int, policy_loss: float, critic_loss: float) -> None:
+        pass
+
     def on_training_end(self, actor: nn.Module, critic: nn.Module, curriculum_config: dict) -> None:
         pass
 
@@ -572,6 +578,9 @@ class MrlPrintCallback(MrlTrainerCallback):
                   reference: dict[str, torch.Tensor], saved_data: dict[str, torch.Tensor], eval_mode: bool) -> None:
         print(f"{'Eval' if eval_mode else 'Train'} | Collected reward {reward}")
 
+    def on_update_epoch_start(self, actor: nn.Module, critic: nn.Module, global_epoch: int, update_epoch: int) -> None:
+        print(f'Epoch {global_epoch} | Starting update epoch {update_epoch}')
+
     def on_batch_updated(self, actor: nn.Module, epoch: int, step: int, policy_loss: float) -> None:
         print(f'Epoch {epoch} | Step {step} - updated policy loss {policy_loss}')
 
@@ -579,6 +588,9 @@ class MrlPrintCallback(MrlTrainerCallback):
                           critic_loss: float) -> None:
         print(f'Epoch {epoch} | Step {step} - updated critic loss {critic_loss}')
 
+    def on_update_epoch_end(self, actor: nn.Module, critic: nn.Module, global_epoch: int, update_epoch: int, policy_loss: float, critic_loss: float) -> None:
+        print(f'Epoch {global_epoch} | Update epoch {update_epoch} - mean policy loss {policy_loss} | mean critic loss {critic_loss}')
+
     def on_training_end(self, actor: nn.Module, critic: nn.Module, curriculum_config: dict) -> None:
         print(f'Finished training for {curriculum_config["steps"]} steps in {curriculum_config["strategy"]} strategy.')
 

rxnn/training/ddp.py

@@ -0,0 +1,26 @@
+import torch
+import torch.distributed as dist
+import os
+from ..utils import set_random_seed
+
+def get_os_ddp_config():
+    rank = int(os.environ['RANK'])
+    world_size = int(os.environ['WORLD_SIZE'])
+    return rank, world_size
+
+def distributed_mean(x: torch.Tensor) -> torch.Tensor:
+    """Average tensor across all devices"""
+    x = x.clone()
+    dist.all_reduce(x, op=dist.ReduceOp.SUM)
+    x /= dist.get_world_size()
+    return x
+
+def distributed_value_mean(value: float, device: torch.device = None) -> float:
+    """Average float value across all devices"""
+    tensor = torch.tensor(value, device=device)
+    reduced = distributed_mean(tensor)
+    return reduced.item()
+
+def set_distributed_random_seed(seed: int):
+    rank = dist.get_rank() if dist.is_initialized() else get_os_ddp_config()[0]
+    set_random_seed(seed + rank)

rxnn/training/mrl.py

@@ -3,7 +3,7 @@ from torch.utils.data import DataLoader, DistributedSampler
 from torch.utils.tensorboard import SummaryWriter
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel
-from typing import Optional, TypedDict, Union, TypeAlias, Literal
+from typing import Optional, TypedDict, Union, TypeAlias, Literal, Callable
 from enum import Enum
 import random, os
 from ..transformers.sampler import BatchSampler
@@ -13,7 +13,7 @@ from .utils import smart_concat, smart_concat_critic_states, TokenizedDict
 from .rl import RlAlgorithm
 from .reward import MrlRewardMode, MrlRewardModel
 from .models import MrlActorAction, MrlActorModel, MrlCriticModel
-
+from .ddp import get_os_ddp_config, distributed_mean
 
 class MrlConfig(TypedDict):
     lr: float
@@ -24,6 +24,11 @@ class MrlConfig(TypedDict):
     critic_max_len: int
     weight_decay: float
     critic_weight_decay: float
+    update_epochs: int
+    pad_token_id: int
+    end_token_id: int
+    callbacks: Optional[list[MrlTrainerCallback]]
+    memory_aware_critic: bool
 
 
 class MrlStrategy(Enum):
@@ -31,8 +36,11 @@ class MrlStrategy(Enum):
     MULTI_STEP_STRATEGY = 2
     LONG_RANGE_STRATEGY = 3
 
+
+UnfreezeStrategyFn = Callable[[int], None]
 UnfreezeItem = Union[int, tuple[int, float]]
-UnfreezeEpochsStrategy: TypeAlias = Union[int, tuple[UnfreezeItem, UnfreezeItem, UnfreezeItem, int]]
+UnfreezeEpochsStrategy: TypeAlias = Union[int, tuple[UnfreezeItem, UnfreezeItem, UnfreezeItem, int], UnfreezeStrategyFn]
+
 
 class CurriculumConfig(TypedDict):
     steps: int
@@ -52,6 +60,7 @@ class CurriculumConfig(TypedDict):
     critic_lr: Optional[float]
     weight_decay: Optional[float]
     critic_weight_decay: Optional[float]
+    update_epochs: Optional[int]
 
 
 class SamplerConfig(TypedDict):
@@ -66,6 +75,7 @@ class MrlTrajectoryStep(TypedDict):
     log_probs: torch.Tensor
     reward: list[float]
     reference: TokenizedDict
+    done: bool
 
 
 class MrlTrajectoryEpisode(TypedDict):
@@ -84,21 +94,25 @@ class MRLTrainer:
             rl_algorithm: RlAlgorithm,
             sampler_config: Optional[SamplerConfig] = None,
             log_dir: str = None,
-            pad_token_id: int = 0,
-            end_token_id: int = 3,
             use_ddp: bool = False,
             use_amp: bool = False,
             dtype: torch.dtype = torch.float32,
-            callbacks: list[MrlTrainerCallback] = None,
-
     ):
         """
-        Trainer for Memory Reinforcement Learning (MRL) in Reactive Transformer.
+        Trainer for Memory Reinforcement Learning (MRL) algorithm for reactive models and Attention-Based Memory System.
 
         Args:
-            actor: MRL Actor model with encoder, decoder and memory attention.
-            critic: Critic network for advantage estimation.
-            config: Configuration dictionary with hyperparameters.
+            actor (MrlActorModel): MRL Actor model with encoder, decoder and memory attention.
+            critic (MrlCriticModel): MRL Critic network for advantage estimation.
+            reward (MrlRewardModel): MRL Reward model or extension.
+            device (torch.device): Device used for training.
+            config (MrlConfig): Configuration dictionary with hyperparameters.
+            rl_algorithm (RlAlgorithm): Reinforcement Learning algorithm (currently only PPO available).
+            sampler_config (SamplerConfig): Sampler configuration.
+            log_dir (str): Log directory for TensorBoard logs.
+            use_ddp (bool): Use Distributed Data Parallel mode.
+            use_amp (bool): Use AMP Autocast for training.
+            dtype (torch.dtype): Data type used in training - in AMP mode it's auto cast, otherwise data and model are transformed to this type
         """
         self.actor = actor
         self.critic = critic
@@ -107,6 +121,10 @@ class MRLTrainer:
         self.device = device
         self.max_seq_len = config.get('max_seq_len', 256)
         self.critic_max_len = config.get('critic_max_len', 512)
+        self.memory_aware_critic = config.get('memory_aware_critic', False)
+        # Internal update epochs config
+        self.shared_update_epochs = config.get('update_epochs', 10)
+        self.update_epochs = self.shared_update_epochs
 
         # Move models to device
         if use_amp:
@@ -117,14 +135,15 @@ class MRLTrainer:
             self.critic.to(self.device, dtype=dtype)
 
         # Batch Sampler for answer generation
-        self.generator = BatchSampler(self.actor, self.device, end_token_id=end_token_id)
+        self.generator = None
         self.sampler_config = SamplerConfig(
             temperature=1.0,
             top_k=None,
             top_p=None,
         ) if sampler_config is None else sampler_config
 
-        self.pad_token_id = pad_token_id
+        self.pad_token_id = config.get('pad_token_id', 0)
+        self.end_token_id = config.get('end_token_id', 3)
 
         self.use_ddp = use_ddp
         self.use_amp = use_amp
@@ -172,7 +191,7 @@ class MRLTrainer:
         self.eval_dataset = None
         self.random_resets_ratio = 0.0
         self.strategy = None
-        self.shared_callbacks = callbacks if callbacks else []
+        self.shared_callbacks = config.get('callbacks', [])
         self.callbacks = []
         self.global_epoch = 0
         self.global_epochs_count = 0
@@ -187,8 +206,8 @@ class MRLTrainer:
     ) -> tuple[torch.optim.Optimizer, torch.optim.Optimizer]:
         if memory_lr is not None:
             optimizer = torch.optim.AdamW([
-                { 'params': self.actor.not_memory_parameters(), 'lr': lr },
-                { 'params': self.actor.memory_parameters(), 'lr': memory_lr },
+                {'params': self.actor.not_memory_parameters(), 'lr': lr},
+                {'params': self.actor.memory_parameters(), 'lr': memory_lr},
             ],
                 weight_decay=weight_decay,
             )
@@ -207,12 +226,10 @@ class MRLTrainer:
 
         return optimizer, critic_optimizer
 
-
     def _init_steps(self):
         return {
             'collect': 0,
-            'critic': 0,
-            'rl': 0,
+            'train': 0,
             'eval': 0,
         }
 
@@ -221,9 +238,12 @@ class MRLTrainer:
         self.epoch_step[step_type] += 1
         self.stage_step[step_type] += 1
 
-    def reset_stm(self) -> bool:
+    def reset_stm(self, force: bool = False) -> bool:
         """Reset Short-Term Memory state with random reset ratio."""
-        if self.random_resets_ratio == 1.0:
+        if force:
+            self.actor.reset_memory()
+            return True
+        elif self.random_resets_ratio == 1.0:
             self.actor.reset_memory()
             return True
         else:
@@ -351,7 +371,7 @@ class MRLTrainer:
                     # state from existing one, instead of new random one)
                     reset_done = self.reset_stm()
 
-                    # 4. Reset reward prev data running mean - it's calculated for multi-step retention, we have to reset it before episode
+                    # 4. Reset reward prev data running mean - it's calculated for multistep retention, we have to reset it before episode
                     self.reward.reset_running_mean()
 
                     # 5. Get first batch of interactions (data to save) and follow-up interactions for current episode, based on curriculum step
@@ -406,6 +426,7 @@ class MRLTrainer:
                             'log_probs': log_probs.detach().cpu(),
                             'reward': reward,
                             'reference': interaction['answer'],
+                            'done': is_last_interaction,
                         }
                         episode_steps.append(trajectory)
                         episode_rewards.append(reward)
@@ -432,201 +453,268 @@ class MRLTrainer:
 
         return trajectories
 
-    def _critic_loss(self, inputs: TokenizedDict, rewards: torch.Tensor) -> torch.Tensor:
-        # 1. Calculate values with critic encoder
-        values = self.critic(
-            inputs['input_ids'],
-            attention_mask=inputs['attention_mask'],
-        ).squeeze()
-        # 2. Calculate critic loss
-        loss = self.rl_algorithm.critic_loss(values, rewards)
-        return loss
-
     def _critic_writer(self, critic_loss: float, epoch: int):
         if self.writer is not None:
-            self.writer.add_scalar('Loss/critic (global)', critic_loss, self.global_step['critic'])
+            self.writer.add_scalar('Loss/critic (global)', critic_loss, self.global_step['train'])
             self.writer.add_scalar(f'Loss/critic (steps: {self.curriculum_steps}, epoch: {epoch})', critic_loss,
-                                   self.epoch_step['critic'])
+                                   self.epoch_step['train'])
             self.writer.add_scalar(f'Loss/critic (steps: {self.curriculum_steps})', critic_loss,
-                                   self.stage_step['critic'])
-
-    def update_critic(self, states: list[tuple[TokenizedDict, TokenizedDict, TokenizedDict]],
-                      rewards: list[torch.Tensor], epoch: int):
-        """Update critic network using MSE loss."""
-        # 1. Run critic updates for all collected batches
-        critic_losses = []
-        for step_idx, (state, reward) in enumerate(zip(states, rewards)):
-            self._increment_steps('critic')
-            # 2. Move state batches to training device (GPU)
-            prev_query, prev_answer, next_query = self._move_multiple_batches(*state)
-
-            # 3. Reset critic gradients
-            self.critic_optimizer.zero_grad()
-
-            # 4. Run critic and calculate loss - in autocast on/off mode
-            if self.use_amp:
-                # Move tensors to training device and calculate loss in autocast mode
-                batch_rewards = reward.to(self.device)
-                with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype):
-                    # Concatenate state into single critic input sequence
-                    inputs = smart_concat_critic_states(
-                        prev_query, prev_answer, next_query,
-                        max_length=self.critic_max_len,
-                        pad_token_id=self.pad_token_id,
-                    )
-                    loss = self._critic_loss(inputs, batch_rewards)
-                # Run backpropagation with scaler
-                self.critic_scaler.scale(loss).backward()
-                # Unscale and clip gradients
-                self.critic_scaler.unscale_(self.critic_optimizer)
-                torch.nn.utils.clip_grad_norm_(self.critic.parameters(), max_norm=1.0, error_if_nonfinite=False)
-                # Run scaled optimization step
-                self.critic_scaler.step(self.critic_optimizer)
-                self.critic_scaler.update()
-            else:
-                # Concatenate state into single critic input sequence
-                inputs = smart_concat_critic_states(
-                    prev_query, prev_answer, next_query,
-                    max_length=self.critic_max_len,
-                    pad_token_id=self.pad_token_id,
-                )
-                # Calculate loss
-                loss = self._critic_loss(inputs, reward.to(self.device, dtype=self.dtype))
-                # Run backpropagation
-                loss.backward()
-                # Clip gradients
-                torch.nn.utils.clip_grad_norm_(self.critic.parameters(), max_norm=1.0, error_if_nonfinite=False)
-                # Run optimizer step
-                self.critic_optimizer.step()
-            critic_loss = loss.item()
-            self._critic_writer(critic_loss, epoch)
-
-            # 5. Run "on critic updated" callbacks
-            for cb in self.callbacks:
-                cb.on_critic_updated(self.actor, self.critic, epoch, step_idx, critic_loss)
-
-            # 6. Accumulate loss for epoch callbacks
-            critic_losses.append(critic_loss)
-
-        # 7. Calculate mean loss for epoch callbacks
-        critic_mean_loss = torch.tensor(critic_losses).mean().item()
-
-        return critic_mean_loss
-
-    def _critic_advantages(self, critic_state: TokenizedDict, rewards: torch.Tensor) -> torch.Tensor:
-        with torch.no_grad():
-            values = self.critic(critic_state['input_ids'],
-                                 attention_mask=critic_state['attention_mask']).squeeze()
-        return self.rl_algorithm.calculate_advantages(rewards, values)
+                                   self.stage_step['train'])
 
     def _rl_writer(self, policy_loss: float, epoch: int):
         if self.writer is not None:
-            self.writer.add_scalar('Loss/policy (global)', policy_loss, self.global_step['rl'])
+            self.writer.add_scalar('Loss/policy (global)', policy_loss, self.global_step['train'])
             self.writer.add_scalar(f'Loss/policy (steps: {self.curriculum_steps}, epoch: {epoch})', policy_loss,
-                                   self.epoch_step['rl'])
-            self.writer.add_scalar(f'Loss/policy (steps: {self.curriculum_steps})', policy_loss, self.stage_step['rl'])
+                                   self.epoch_step['train'])
+            self.writer.add_scalar(f'Loss/policy (steps: {self.curriculum_steps})', policy_loss,
+                                   self.stage_step['train'])
 
-    def rl_step(self, trajectories: list[MrlTrajectoryEpisode], epoch: int):
+    def update_critic(self, state: tuple[TokenizedDict, TokenizedDict, TokenizedDict], ref_values: torch.Tensor,
+                      epoch: int) -> float:
+        # 1. Reset critic gradients
+        self.critic_optimizer.zero_grad()
+
+        # 2. Update critic - with autocast on/off
+        if self.use_amp:
+            with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype):
+                # 2.1 Concat states and calculate critic loss
+                critic_state = smart_concat_critic_states(*state, max_length=self.critic_max_len,
+                                                          pad_token_id=self.pad_token_id)
+                values = self.critic(critic_state['input_ids'], attention_mask=critic_state['attention_mask']).squeeze()
+                critic_loss = self.rl_algorithm.critic_loss(values, ref_values)
+            # 2.2 Run backpropagation with scaler
+            self.critic_scaler.scale(critic_loss).backward()
+            # 2.3 Unscale and clip gradients
+            self.critic_scaler.unscale_(self.critic_optimizer)
+            torch.nn.utils.clip_grad_norm_(self.critic.parameters(), max_norm=1.0, error_if_nonfinite=False)
+            # 2.4 Run scaled optimization step
+            self.critic_scaler.step(self.critic_optimizer)
+            self.critic_scaler.update()
+        else:
+            # 2.1 Concat states and calculate critic loss
+            critic_state = smart_concat_critic_states(*state, max_length=self.critic_max_len,
+                                                      pad_token_id=self.pad_token_id)
+            values = self.critic(critic_state['input_ids'], attention_mask=critic_state['attention_mask']).squeeze()
+            critic_loss = self.rl_algorithm.critic_loss(values, ref_values)
+            # 2.2 Run backpropagation
+            critic_loss.backward()
+            # 2.3 Clip gradients
+            torch.nn.utils.clip_grad_norm_(self.critic.parameters(), max_norm=1.0, error_if_nonfinite=False)
+            # 2.4 Run optimizer step
+            self.critic_optimizer.step()
+        # 3. Get float loss value for callbacks/writer
+        critic_loss_item = critic_loss.item()
+
+        # 4. Write to TensorBoard
+        self._critic_writer(critic_loss_item, epoch)
+
+        # 5. Run "on critic updated" callbacks
+        for cb in self.callbacks:
+            cb.on_critic_updated(self.actor, self.critic, epoch, self.epoch_step['train'], critic_loss_item)
+        # 6. Return loss item
+        return critic_loss_item
+
+    def update_actor(self, state: tuple[TokenizedDict, TokenizedDict, TokenizedDict], action: TokenizedDict,
+                     advantages: torch.Tensor, old_log_probs: torch.Tensor, epoch: int) -> float:
+        # 1. Reset actor gradients
+        self.optimizer.zero_grad()
+        # 2. Unpack state dicts
+        query, answer, next_query = state
+
+        # 3. Encode and update STM on each step, to include encoder and memory attention gradients in loss (skip if it was updated before with memory aware critic)
+        if not self.memory_aware_critic:
+            self.encode_and_update_stm(query, answer)
+
+        # 4. Update actor - with autocast on/off
+        if self.use_amp:
+            with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype):
+                # 4.1 Concatenate next query and action and get action logits from decoder
+                inputs = smart_concat(next_query, action, max_length=self.max_seq_len,
+                                      pad_token_id=self.pad_token_id)
+                logits = self.actor(inputs['input_ids'], attention_mask=inputs['attention_mask'],
+                                    action=MrlActorAction.DECODE)
+                # 4.2 Calculate policy loss with selected algorithm
+                policy_loss = self.rl_algorithm.policy_loss(next_query, action, logits, old_log_probs,
+                                                            advantages)
+            # 4.3 Run backpropagation with scaler
+            self.scaler.scale(policy_loss).backward(retain_graph=True)
+            # 4.4 Unscale and clip gradient norms
+            self.scaler.unscale_(self.optimizer)
+            torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
+                                           error_if_nonfinite=False)
+            # 4.5 Run scaled optimization step
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
+        else:
+            # 4.1 Concatenate next query and action and get action logits from decoder
+            inputs = smart_concat(next_query, action, max_length=self.max_seq_len,
+                                  pad_token_id=self.pad_token_id)
+            logits = self.actor(inputs['input_ids'], attention_mask=inputs['attention_mask'],
+                                action=MrlActorAction.DECODE)
+            # 4.2 Calculate policy loss with selected algorithm
+            policy_loss = self.rl_algorithm.policy_loss(next_query, action, logits, old_log_probs, advantages)
+            # 4.3 Run backpropagation
+            policy_loss.backward(retain_graph=True)
+            # 4.4 Clip gradient norms
+            torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
+                                           error_if_nonfinite=False)
+            # 4.5 Run scaled optimization step
+            self.optimizer.step()
+        # 5. Get float loss value for callbacks/writer
+        policy_loss_item = policy_loss.item()
+
+        # 6. Write to TensorBoard
+        self._rl_writer(policy_loss_item, epoch)
+
+        # 7. Run "on batch updated" callback
+        for cb in self.callbacks:
+            cb.on_batch_updated(self.actor, epoch, self.epoch_step['train'], policy_loss_item)
+
+        # 8. Return loss item
+        return policy_loss_item
+
+    def rl_step(self, trajectories: list[MrlTrajectoryEpisode], advantages: torch.Tensor, ref_values: torch.Tensor,
+                epoch: int, batch_size: int) -> tuple[torch.Tensor, torch.Tensor]:
         """Perform PPO update step using trajectories."""
         # 1. Run update separately for episodes in trajectory - we have to reset memory before each episode, and update
         # memory, based on collected episode data
         all_losses = []
-        trajectories_len = len(trajectories)
+        critic_losses = []
         for episode_idx, episode in enumerate(trajectories):
             episode_steps = episode['steps']
             should_reset_stm = episode['reset_stm']
 
-            # 2. Reset memory for current batch episode
+            # 2. Get advantages and reference values for current full episode (batch_size * episode_steps)
+            num_steps = len(episode_steps)
+            start = episode_idx * num_steps
+            end = start + num_steps
+            episode_critic_values = ref_values[start:end]
+            episode_advantages = advantages[start:end]
+
+            # 3. Reset memory for current batch episode
             if should_reset_stm:
-                self.reset_stm()
+                self.reset_stm(force=True)
 
-            # 3. Run episode steps - each episode has number of steps depending on curriculum stage. Each step is run for all batch
-            for step in episode_steps:
-                self._increment_steps('rl')
-                state, action, reward, log_probs = step['state'], step['action'], step['reward'], step['log_probs']
+            # 4. Run episode steps - each episode has number of steps depending on curriculum stage. Each step is run for all batch
+            for step_idx, step in enumerate(episode_steps):
+                self._increment_steps('train')
+                # 5. Get and move to device collected states, action and log probs
+                state, action, _, log_probs = step['state'], step['action'], step['reward'], step['log_probs']
                 query, answer, next_query = self._move_multiple_batches(*state)
                 action = self._move_batch(action)
                 log_probs = log_probs.to(self.device)
-                rewards = torch.tensor(reward).to(self.device)
-
-                # 4. Compute advantages using critic
-                if self.use_amp:
-                    with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype):
-                        critic_state = smart_concat_critic_states(query, answer, next_query,
-                                                                  max_length=self.critic_max_len,
-                                                                  pad_token_id=self.pad_token_id)
-                        advantages = self._critic_advantages(critic_state, rewards)
-                else:
-                    critic_state = smart_concat_critic_states(query, answer, next_query, max_length=self.critic_max_len,
-                                                              pad_token_id=self.pad_token_id)
-                    advantages = self._critic_advantages(critic_state, rewards)
-
-                # 5. Encode and update STM on each step, to include encoder and memory attention gradients in loss
-                self.encode_and_update_stm(query, answer)
-                # 6. Concatenate next query and action and get action logits from decoder
-                if self.use_amp:
-                    with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype):
-                        inputs = smart_concat(next_query, action, max_length=self.max_seq_len,
-                                              pad_token_id=self.pad_token_id)
-                        logits = self.actor(inputs['input_ids'], attention_mask=inputs['attention_mask'],
-                                            action=MrlActorAction.DECODE)
-                else:
-                    inputs = smart_concat(next_query, action, max_length=self.max_seq_len,
-                                          pad_token_id=self.pad_token_id)
-                    logits = self.actor(inputs['input_ids'], attention_mask=inputs['attention_mask'],
-                                        action=MrlActorAction.DECODE)
-
-                # 7. Calculate RL Algorithm (PPO etc.) loss
-                policy_loss = self.rl_algorithm.policy_loss(next_query, action, logits, log_probs, advantages)
-
-                # 8. Reset gradients
-                self.optimizer.zero_grad()
-
-                # 9. Update the model in AMP or regular mode
-                if self.use_amp:
-                    self.scaler.scale(policy_loss).backward(retain_graph=True)
-                    self.scaler.unscale_(self.optimizer)
-                    torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
-                                                   error_if_nonfinite=False)
-                    self.scaler.step(self.optimizer)
-                    self.scaler.update()
-                else:
-                    policy_loss.backward(retain_graph=True)
-                    torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
-                                                   error_if_nonfinite=False)
-                    self.optimizer.step()
 
-                policy_loss_item = policy_loss.item()
-                self._rl_writer(policy_loss_item, epoch)
-                all_losses.append(policy_loss_item)
+                # 6. Select advantages and reference values for current step (batch_size)
+                step_critic_values = episode_critic_values[step_idx]
+                step_advantages = episode_advantages[step_idx]
 
-                # 10. Run "on batch updated" callback
-                for cb in self.callbacks:
-                    cb.on_batch_updated(self.actor, epoch, self.epoch_step['rl'], policy_loss_item)
+                # 7. In memory aware critic version, encode and update STM before critic update, to include its gradients in critic loss too
+                if self.memory_aware_critic:
+                    self.encode_and_update_stm(query, answer)
+
+                # 7. Update critic
+                critic_loss_item = self.update_critic((query, answer, next_query), step_critic_values, epoch)
 
-        return torch.mean(torch.tensor(all_losses)).item()
+                # 8. Accumulate critic loss for epoch callbacks
+                critic_losses.append(critic_loss_item)
 
-    def _critic_states_and_rewards(self, trajectories: list[MrlTrajectoryEpisode]):
-        flat_trajectories = [t for episode in trajectories for t in episode['steps']]
-        states = [t['state'] for t in flat_trajectories]
-        rewards = [torch.tensor(t['reward']) for t in flat_trajectories]
-        return states, rewards
+                # 9. Update actor
+                policy_loss_item = self.update_actor((query, answer, next_query), action, step_advantages, log_probs,
+                                                     epoch)
+                all_losses.append(policy_loss_item)
+        # 10. Return mean losses for epoch callbacks
+        return torch.mean(torch.tensor(all_losses)), torch.mean(torch.tensor(critic_losses))
+
+    def _critic_values_rewards_and_dones(self, trajectories: list[MrlTrajectoryEpisode]):
+        if self.memory_aware_critic:
+            flat_trajectories = [
+                (t, i == 0 and episode['reset_stm'])
+                for episode in trajectories
+                for i, t in enumerate(episode['steps'])
+            ]
+            values = torch.stack([
+                self._critic_values_with_memory(r, *self._move_multiple_batches(*t['state'])) for t, r in flat_trajectories
+            ]).to(self.device)
+            rewards = torch.stack([torch.tensor(t['reward']) for t, _ in flat_trajectories]).to(self.device)
+            dones = torch.stack([torch.tensor(t['done']) for t, _ in flat_trajectories]).to(self.device)
+        else:
+            flat_trajectories = [t for episode in trajectories for t in episode['steps']]
+            values = torch.stack([
+                self._critic_values(*self._move_multiple_batches(*t['state'])) for t in flat_trajectories
+            ]).to(self.device)
+            rewards = torch.stack([torch.tensor(t['reward']) for t in flat_trajectories]).to(self.device)
+            dones = torch.stack([torch.tensor(t['done']) for t in flat_trajectories]).to(self.device)
+        return values, rewards, dones
+
+    def _critic_values_with_memory(self, reset_stm: bool, *moved_state: tuple[TokenizedDict, TokenizedDict, TokenizedDict]) -> torch.Tensor:
+        # 1. Calculate critic values in memory aware version - reset/update STM before calculating values
+        with torch.no_grad():
+            # 2. Reset STM if it was reset in trajectory collection
+            if reset_stm:
+                self.reset_stm(force=True)
+            # 3. Encode and update STM for critic
+            self.encode_and_update_stm(*moved_state)
+            # 4. Get concatenated critic states
+            inputs = smart_concat_critic_states(
+                *moved_state,
+                max_length=self.critic_max_len,
+                pad_token_id=self.pad_token_id,
+            )
+            # 5. Calculate values for current batch
+            return self.critic(inputs['input_ids'],
+                               attention_mask=inputs['attention_mask']).squeeze()
+
+    def _critic_values(self, *moved_state: tuple[TokenizedDict, TokenizedDict, TokenizedDict]) -> torch.Tensor:
+        # 1. Calculate critic values
+        with torch.no_grad():
+            # 2. Get concatenated critic states
+            inputs = smart_concat_critic_states(
+                *moved_state,
+                max_length=self.critic_max_len,
+                pad_token_id=self.pad_token_id,
+            )
+            # 3. Calculate values for current batch
+            return self.critic(inputs['input_ids'],
+                               attention_mask=inputs['attention_mask']).squeeze()
 
     def train_epoch(self, dataloader: DataLoader, epoch: int, batch_size: int):
         """Train for one epoch."""
         # 1. Collect trajectories for current epoch
         trajectories = self.collect_trajectories(dataloader, epoch, batch_size)
 
-        # 2. Flatten trajectories and collect state and rewards for critic update
-        states, rewards = self._critic_states_and_rewards(trajectories)
-        # 3. Update critic model, based on states and rewards
-        critic_loss = self.update_critic(states, rewards, epoch)
+        # 2. Flatten trajectories, call critic and collect values, dones and rewards, and calculate advantages
+        if self.use_amp:
+            with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype):
+                values, rewards, dones = self._critic_values_rewards_and_dones(trajectories)
+                advantages, ref_values = self.rl_algorithm.calculate_advantages(rewards, values, dones)
+        else:
+            values, rewards, dones = self._critic_values_rewards_and_dones(trajectories)
+            advantages, ref_values = self.rl_algorithm.calculate_advantages(rewards, values, dones)
+
+        # 3. Run internal update epochs
+        critic_loss_sum, policy_loss_sum = 0.0, 0.0
+        for update_epoch in range(self.update_epochs):
+            # 4. Run 'on update epoch start' callbacks
+            for cb in self.callbacks:
+                cb.on_update_epoch_start(self.actor, self.critic, epoch, update_epoch)
+
+            # 5. Run RL algorithm step
+            policy_loss, critic_loss = self.rl_step(trajectories[:-1], advantages, ref_values, epoch, batch_size)
+
+            if self.use_ddp:
+                policy_loss = distributed_mean(policy_loss)
+                critic_loss = distributed_mean(critic_loss)
+
+            # 6. Run 'on update epoch end' callbacks
+            for cb in self.callbacks:
+                cb.on_update_epoch_end(self.actor, self.critic, epoch, update_epoch, policy_loss, critic_loss)
 
-        # 4. Run PPO algorithm step
-        policy_loss = self.rl_step(trajectories, epoch)
+            # 7. Accumulate losses for epoch callbacks
+            critic_loss_sum += critic_loss
+            policy_loss_sum += policy_loss
 
-        # 5. Return policy and critic mean losses for epoch callbacks
-        return policy_loss, critic_loss
+        # 8. Return policy and critic mean losses for epoch callbacks
+        return policy_loss_sum / self.update_epochs, critic_loss_sum / self.update_epochs
 
     def _eval_loader(self, batch_size: int):
         if self.use_ddp:
@@ -731,12 +819,11 @@ class MRLTrainer:
                         cb.on_eval_episode_end(self.actor, epoch, self.epoch_step['eval'], avg_episode_reward)
 
         # 15. Calculate average reward
+        avg_reward = (total_reward / count) if count > 0 else torch.tensor(0.0).to(self.device)
         if self.use_ddp:
-            total_sum = dist.all_reduce(total_reward, dist.ReduceOp.SUM)
-            count_sum = dist.all_reduce(count, dist.ReduceOp.SUM)
-            avg_reward = (total_sum / count_sum).item() if count_sum > 0 else 0
-        else:
-            avg_reward = (total_reward / count).item() if count > 0 else 0
+            avg_reward = distributed_mean(avg_reward)
+
+        avg_reward = avg_reward.item()
 
         should_stop_stage = False
         # 16. Run "on eval end" callbacks
@@ -747,55 +834,6 @@ class MRLTrainer:
 
         return should_stop_stage
 
-    def _setup_curriculum_step(self, config: CurriculumConfig) -> tuple[tuple[int, UnfreezeEpochsStrategy], tuple[bool, int, float]]:
-        # 1. Set common fields based on config
-        self.curriculum_steps = config.get('steps', 1)  # number of steps to run in episode
-        self.train_dataset = config.get('dataset', None)  # training dataset for current curriculum stage
-        self.eval_dataset = config.get('eval_dataset', None)  # evaluation dataset for current curriculum stage
-        self.callbacks = config.get('callbacks',
-                                    self.shared_callbacks)  # trainer callbacks for current curriculum stage
-        self.strategy = config.get('strategy',
-                                   MrlStrategy.MULTI_STEP_STRATEGY)  # MRL strategy for given curriculum stage
-        self.reward = config.get('reward_model', self.shared_reward_model)  # MRL Reward Model for curriculum stage
-        if config['lr'] is not None or config['critic_lr'] is not None or config['weight_decay'] is not None or config['critic_weight_decay'] is not None or (config['separate_memory_lr'] and config['memory_lr'] is not None):
-            if config.get('separate_memory_lr', False):
-                self.optim_config = {
-                    'lr': config.get('lr', self.base_optim_config['lr']),
-                    'critic_lr': config.get('critic_lr', self.base_optim_config['critic_lr']),
-                    'weight_decay': config.get('weight_decay', self.base_optim_config['weight_decay']),
-                    'critic_weight_decay': config.get('critic_weight_decay', self.base_optim_config['critic_weight_decay']),
-                    'memory_lr': config.get('memory_lr', self.base_optim_config['memory_lr']),
-                }
-            else:
-                self.optim_config = {
-                    'lr': config.get('lr', self.base_optim_config['lr']),
-                    'critic_lr': config.get('critic_lr', self.base_optim_config['critic_lr']),
-                    'weight_decay': config.get('weight_decay', self.base_optim_config['weight_decay']),
-                    'critic_weight_decay': config.get('critic_weight_decay', self.base_optim_config['critic_weight_decay']),
-                }
-            self.optimizer, self.critic_optimizer = self._init_optimizers(**self.optim_config)
-        elif self.optim_config != self.base_optim_config:
-            self.optim_config = self.base_optim_config
-            self.optimizer, self.critic_optimizer = self._init_optimizers(**self.optim_config)
-
-
-
-
-        # 2. Get epochs and random resets configs
-        epochs = config.get('epochs', 5)  # number of epochs for current stage
-        unfreeze_epoch = config.get('unfreeze_epoch',
-                                    0)  # epoch when components (other than memory) are unfrozen (before epoch starts)
-        random_resets = config.get('random_resets',
-                                   False)  # flag for using random STM resets (recommended, as model should learn transitions between different states)
-        random_resets_from = config.get('random_resets_from', None)  # epoch from which random STM resets are started
-        random_resets_ratio = config.get('random_resets_ratio',
-                                         None)  # ratio of random STM resets - 1.0 is "always reset", 0.0 is "no resets"
-
-        # 3. Reset stage step counter
-        self.stage_step = self._init_steps()
-
-        return (epochs, unfreeze_epoch), (random_resets, random_resets_from, random_resets_ratio)
-
     def _apply_unfreeze_strategy(self, epoch: int, unfreeze_epoch: UnfreezeEpochsStrategy):
         is_staged_unfreeze = isinstance(unfreeze_epoch, tuple)
         if is_staged_unfreeze:
@@ -808,28 +846,31 @@ class MRLTrainer:
                     self.optimizer = self._init_unfreeze_optimizer('update', cross_att_lr)
                     print(f"Activating 'update' unfreeze strategy with custom cross_att_lr: {cross_att_lr}")
             elif epoch == update_epoch:
-                 self.actor.freeze_components('update')
-                 print(f"Activating 'update' unfreeze strategy - mem-att trainable / cross-att frozen / rest model frozen")
+                self.actor.freeze_components('update')
+                print(
+                    f"Activating 'update' unfreeze strategy - mem-att trainable / cross-att frozen / rest model frozen")
 
             if isinstance(fetch_epoch, tuple):
                 switch_epoch, mem_att_lr = fetch_epoch
-                if epoch == fetch_epoch:
+                if epoch == switch_epoch:
                     self.actor.freeze_components('joint')
                     self.optimizer = self._init_unfreeze_optimizer('fetch', mem_att_lr)
                     print(f"Activating 'fetch' unfreeze strategy with custom mem_att_lr: {mem_att_lr}")
             elif epoch == fetch_epoch:
                 self.actor.freeze_components('fetch')
-                print(f"Activating 'fetch' unfreeze strategy - mem-att frozen / cross-att trainable / rest model frozen")
+                print(
+                    f"Activating 'fetch' unfreeze strategy - mem-att frozen / cross-att trainable / rest model frozen")
 
             if isinstance(joint_epoch, tuple):
                 switch_epoch, model_lr = joint_epoch
-                if epoch == joint_epoch:
+                if epoch == switch_epoch:
                     self.actor.unfreeze_components()
                     self.optimizer = self._init_unfreeze_optimizer('joint', model_lr)
                     print(f"Activating 'joint' unfreeze strategy with custom model_lr: {model_lr}")
             elif epoch == joint_epoch:
-                    self.actor.freeze_components('joint')
-                    print(f"Activating 'joint' unfreeze strategy - mem-att/cross-att trainable / rest model frozen")
+                self.actor.freeze_components('joint')
+                print(f"Activating 'joint' unfreeze strategy - mem-att/cross-att trainable / rest model frozen")
+
             if epoch == all_epoch:
                 self.actor.unfreeze_components()
                 self.optimizer = self._init_unfreeze_optimizer('all', 0.)
@@ -871,6 +912,56 @@ class MRLTrainer:
 
         return torch.optim.AdamW(params, weight_decay=self.optim_config['weight_decay'])
 
+    def _setup_curriculum_step(self, config: CurriculumConfig) -> tuple[
+        tuple[int, UnfreezeEpochsStrategy], tuple[bool, int, float]]:
+        # 1. Set common fields based on config
+        self.curriculum_steps = config.get('steps', 1)  # number of steps to run in episode
+        self.train_dataset = config.get('dataset', None)  # training dataset for current curriculum stage
+        self.eval_dataset = config.get('eval_dataset', None)  # evaluation dataset for current curriculum stage
+        self.callbacks = config.get('callbacks',
+                                    self.shared_callbacks)  # trainer callbacks for current curriculum stage
+        self.strategy = config.get('strategy',
+                                   MrlStrategy.MULTI_STEP_STRATEGY)  # MRL strategy for given curriculum stage
+        self.reward = config.get('reward_model', self.shared_reward_model)  # MRL Reward Model for curriculum stage
+        self.update_epochs = config.get('update_epochs', self.shared_update_epochs)  # Internal update epochs
+        if config['lr'] is not None or config['critic_lr'] is not None or config['weight_decay'] is not None or config[
+            'critic_weight_decay'] is not None or (config['separate_memory_lr'] and config['memory_lr'] is not None):
+            if config.get('separate_memory_lr', False):
+                self.optim_config = {
+                    'lr': config.get('lr', self.base_optim_config['lr']),
+                    'critic_lr': config.get('critic_lr', self.base_optim_config['critic_lr']),
+                    'weight_decay': config.get('weight_decay', self.base_optim_config['weight_decay']),
+                    'critic_weight_decay': config.get('critic_weight_decay',
+                                                      self.base_optim_config['critic_weight_decay']),
+                    'memory_lr': config.get('memory_lr', self.base_optim_config['memory_lr']),
+                }
+            else:
+                self.optim_config = {
+                    'lr': config.get('lr', self.base_optim_config['lr']),
+                    'critic_lr': config.get('critic_lr', self.base_optim_config['critic_lr']),
+                    'weight_decay': config.get('weight_decay', self.base_optim_config['weight_decay']),
+                    'critic_weight_decay': config.get('critic_weight_decay',
+                                                      self.base_optim_config['critic_weight_decay']),
+                }
+            self.optimizer, self.critic_optimizer = self._init_optimizers(**self.optim_config)
+        elif self.optim_config != self.base_optim_config:
+            self.optim_config = self.base_optim_config
+            self.optimizer, self.critic_optimizer = self._init_optimizers(**self.optim_config)
+
+        # 2. Get epochs and random resets configs
+        epochs = config.get('epochs', 5)  # number of epochs for current stage
+        unfreeze_epoch = config.get('unfreeze_epoch',
+                                    0)  # epoch when components (other than memory) are unfrozen (before epoch starts)
+        random_resets = config.get('random_resets',
+                                   False)  # flag for using random STM resets (recommended, as model should learn transitions between different states)
+        random_resets_from = config.get('random_resets_from', None)  # epoch from which random STM resets are started
+        random_resets_ratio = config.get('random_resets_ratio',
+                                         None)  # ratio of random STM resets - 1.0 is "always reset", 0.0 is "no resets"
+
+        # 3. Reset stage step counter
+        self.stage_step = self._init_steps()
+
+        return (epochs, unfreeze_epoch), (random_resets, random_resets_from, random_resets_ratio)
 
     def __call__(self, curriculum_config: list[CurriculumConfig], batch_size: int):
         """Start Memory Reinforcement Learning Curriculum."""
@@ -881,29 +972,36 @@ class MRLTrainer:
 
         # 1. Init DDP for distributed training mode
         if self.use_ddp:
-            rank = int(os.environ['RANK'])
-            world_size = int(os.environ['WORLD_SIZE'])
+            rank, world_size = get_os_ddp_config()
             dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
             self.actor = DistributedDataParallel(self.actor, device_ids=[self.device.index])
             self.critic = DistributedDataParallel(self.critic, device_ids=[self.device.index])
 
-        # 2. Run each curriculum step based on config
+        # 2. Init BatchSampler with actor model (we have to run it after DDP init)
+        self.generator = BatchSampler(self.actor, self.device, end_token_id=self.end_token_id)
+
+        # 3. Run each curriculum step based on config
         for current_curriculum_step in curriculum_config:
-            # 3. Setup training config for curriculum step
+            # 4. Setup training config for curriculum step
             epochs_config, random_resets_config = self._setup_curriculum_step(current_curriculum_step)
             epochs, unfreeze_epoch = epochs_config
             random_resets, random_resets_from, random_resets_ratio = random_resets_config
             assert self.train_dataset is not None
 
-            # 4. Freeze all components except memory attention and memory cross-attention layers in decoder/encoder
+            # 5. Freeze all components except memory attention and memory cross-attention layers in decoder/encoder
             if unfreeze_epoch != 0:
-                self.actor.freeze_components('joint')
-                if isinstance(unfreeze_epoch, tuple):
-                    print(f"Starting training with unfreeze strategies - 'warmup' - mem-att/cross-att trainable / rest model frozen")
+                if callable(unfreeze_epoch):
+                    unfreeze_epoch(-1)
                 else:
-                    print(f"Starting training with simple unfreeze - 'joint' - mem-att/cross-att trainable / rest model frozen")
-
-            # 5. Setup train DataLoader
+                    self.actor.freeze_components('joint')
+                    if isinstance(unfreeze_epoch, tuple):
+                        print(
+                            f"Starting training with unfreeze strategies - 'warmup' - mem-att/cross-att trainable / rest model frozen")
+                    else:
+                        print(
+                            f"Starting training with simple unfreeze - 'joint' - mem-att/cross-att trainable / rest model frozen")
+
+            # 6. Setup train DataLoader
             if self.use_ddp:
                 train_sampler = DistributedSampler(self.train_dataset, shuffle=True)
                 dataloader = DataLoader(
@@ -912,6 +1010,7 @@ class MRLTrainer:
                     sampler=train_sampler,
                     pin_memory=True,
                     collate_fn=MrlCurriculumDataset.collate_mrl_batch,
+                    drop_last=True,
                 )
             else:
                 train_sampler = None
@@ -923,65 +1022,68 @@ class MRLTrainer:
                     collate_fn=MrlCurriculumDataset.collate_mrl_batch,
                 )
 
-            # 6. Run selected number of epochs for given curriculum stage
+            # 7. Run selected number of epochs for given curriculum stage
             for epoch in range(epochs):
-                # 7. Increment global epoch
+                # 8. Increment global epoch
                 self.global_epoch += 1
-                # 8. Run "on epoch start" callbacks (log info, etc.)
+                # 9. Run "on epoch start" callbacks (log info, etc.)
                 for cb in self.callbacks:
                     cb.on_epoch_start(self.actor, epoch, epochs, current_curriculum_step, self.global_epoch,
                                       self.global_epochs_count)
 
-                # 9. Reset steps counter for epoch
+                # 10. Reset steps counter for epoch
                 self.epoch_step = self._init_steps()
 
-                # 10. Set random STM resets ratio from selected epoch
+                # 11. Set random STM resets ratio from selected epoch
                 if random_resets and random_resets_from <= epoch:
                     self.random_resets_ratio = random_resets_ratio
                 else:
                     self.random_resets_ratio = 1.0
 
-                # 11. Apply the unfreeze strategy
-                self._apply_unfreeze_strategy(epoch, unfreeze_epoch)
+                # 12. Apply the unfreeze strategy
+                if callable(unfreeze_epoch):
+                    unfreeze_epoch(epoch)
+                else:
+                    self._apply_unfreeze_strategy(epoch, unfreeze_epoch)
 
-                # 12. Set epoch for distributed sampler
+                # 13. Set epoch for distributed sampler
                 if train_sampler is not None:
                     train_sampler.set_epoch(epoch)
 
-                # 13. Run reinforcement learning algorithms for current epoch
+                # 14. Run reinforcement learning algorithms for current epoch
                 policy_loss, critic_loss = self.train_epoch(dataloader, epoch, batch_size)
 
-                # 14. If evaluation dataset is provided, run evaluation steps
+                # 15. If evaluation dataset is provided, run evaluation steps
                 if self.eval_dataset:
                     should_stop_stage = self.evaluate(batch_size, epoch)
                 else:
                     should_stop_stage = False
 
-                # 15. Finally, run "on epoch end" callbacks (save models, etc.)
+                # 16. Finally, run "on epoch end" callbacks (save models, etc.)
                 for cb in self.callbacks:
                     cb.on_epoch_end(self.actor, epoch, epochs, policy_loss, critic_loss, self.global_epoch,
                                     self.global_epochs_count)
 
-                # 16. Synchronize TensorBoard writer
+                # 17. Synchronize TensorBoard writer
                 if self.writer:
                     self.writer.flush()
 
-                # 17. Synchronize devices in DDP mode
+                # 18. Synchronize devices in DDP mode
                 if self.use_ddp:
                     dist.barrier()
 
-                # 18. Finish curriculum stage if rewards are not increased or reached threshold point
+                # 19. Finish curriculum stage if rewards are not increased or reached threshold point
                 if should_stop_stage:
                     break
 
-            # 19. Run "on_training_end" callbacks after each curriculum stage (they have own callbacks)
+            # 20. Run "on_training_end" callbacks after each curriculum stage (they have own callbacks)
             for cb in self.callbacks:
                 cb.on_training_end(self.actor, self.critic, current_curriculum_step)
 
-        # 20. Training end - finish processes after all curriculum stages
+        # 21. Training end - finish processes after all curriculum stages
         if self.use_ddp:
             dist.destroy_process_group()
 
-        # 21. Close writer
+        # 22. Close writer
         if self.writer:
             self.writer.close()

rxnn/training/rl.py

@@ -2,8 +2,9 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from abc import ABC, abstractmethod
-from typing import TypedDict
+from typing import TypedDict, Optional
 from .utils import TokenizedDict
+from .ddp import distributed_mean
 
 
 class RlAlgorithm(ABC):
@@ -17,21 +18,34 @@ class RlAlgorithm(ABC):
         pass
 
     @abstractmethod
-    def calculate_advantages(self, rewards: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
+    def calculate_advantages(self, rewards: torch.Tensor, values: torch.Tensor, dones: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         pass
 
     def critic_loss(self, rewards: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
         return self.critic_loss(rewards, values)
 
+
 class PPOConfig(TypedDict):
-    clip_eps: float
+    clip_eps: Optional[float]
+    gae_lambda: Optional[float]
+    gae_gamma: Optional[float]
+    entropy_coef: Optional[float]
+    use_distributed_advantage_norm: Optional[bool]
+
 
 class PPOAlgorithm(RlAlgorithm):
-    def __init__(self, config: PPOConfig):
+    def __init__(self, config: Optional[PPOConfig] = None):
         super(PPOAlgorithm, self).__init__()
 
+        if config is None:
+            config = {}
+
         # PPO Config
         self.clip_eps = config.get('clip_eps', 0.2)
+        self.gae_lambda = config.get('gae_lambda', 0.95)
+        self.gae_gamma = config.get('gae_gamma', 0.99)
+        self.entropy_coef = config.get('entropy_coef', 0.01)
+        self.use_distributed_advantage_norm = config.get('use_distributed_advantage_norm', False)
 
     def policy_loss(self, query: TokenizedDict, answer: TokenizedDict, logits: torch.Tensor,
                     old_log_probs: torch.Tensor, advantages: torch.Tensor) -> torch.Tensor:
@@ -78,11 +92,37 @@ class PPOAlgorithm(RlAlgorithm):
 
         # d) Entropy bonus
         entropy = -torch.sum(new_probs * new_probs.exp(), dim=-1).mean()
-        policy_loss -= 0.01 * entropy
+        policy_loss -= self.entropy_coef * entropy
 
         return policy_loss
 
-    def calculate_advantages(self, rewards: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
-        advantages = rewards - values
-        normalized_advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
-        return normalized_advantages
+    def _compute_gae(self, rewards: torch.Tensor, values: torch.Tensor, next_value: torch.Tensor, dones: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        T, B = rewards.shape
+        advantages = torch.zeros_like(rewards, device=values.device)
+        last_advantage = 0
+        last_value = next_value.detach()
+
+        for t in reversed(range(T)):
+            if t == T - 1:
+                next_values = last_value
+            else:
+                next_values = values[t + 1]
+
+            # Mask next values if episode ended
+            next_values = next_values * ~dones[t]
+            delta = rewards[t] + self.gae_gamma * next_values - values[t]
+            advantages[t] = delta + self.gae_gamma * self.gae_lambda * last_advantage
+            last_advantage = advantages[t]
+
+        returns = advantages + values
+        return advantages, returns
+
+    def calculate_advantages(self, rewards: torch.Tensor, values: torch.Tensor, dones: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        advantages, ref_values = self._compute_gae(rewards[:-1], values[:-1], values[-1], dones[:-1])
+        if self.use_distributed_advantage_norm:
+            mean_advantage = distributed_mean(advantages.mean())
+            std_advantage = distributed_mean(advantages.std())
+            normalized_advantages = (advantages - mean_advantage) / (std_advantage + 1e-8)
+        else:
+            normalized_advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
+        return normalized_advantages, ref_values

{rxnn-0.2.29.dist-info → rxnn-0.2.31.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.2.29
+Version: 0.2.31
 Summary: RxNN: Reactive Neural Networks Platform
 License: Apache-2.0
 Keywords: deep-learning,ai,machine-learning

{rxnn-0.2.29.dist-info → rxnn-0.2.31.dist-info}/RECORD

@@ -11,14 +11,15 @@ rxnn/memory/stm.py,sha256=IH_3INw7FdI013t56ui3Zq9GPUq-k3HeZGjx6BerS4g,3888
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/rxt/models.py,sha256=r8wZeeNTC2VAhiiNe4y7LrbnB4wjFu_cupKiGkpdgjI,13002
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/training/base.py,sha256=_xik1GXE4RJ_nxwqLQ1ccXA5pRtBCi-jL-jeRFBdHBU,11851
-rxnn/training/bml.py,sha256=FJszaQXOLx2ZHBa1CQpyMrG8i4Kj14E-gzDAEK_Ei5k,17272
-rxnn/training/callbacks.py,sha256=-N0MQPpZQaUWCINdTOsjul4bDGbGr2JgQBqOIXBLS6o,35053
+rxnn/training/base.py,sha256=TGz_37RfI1qLI31GNRV5rLowW1kAHnJwqPm7DNfLfe4,11730
+rxnn/training/bml.py,sha256=hw6gLpLkGvqLzxIvBg4MvCc5r8cHpEm2RDyh7nH6CtE,16914
+rxnn/training/callbacks.py,sha256=p72lbzFAmFjpcUvyy4aUB3qd53I8C6Sk5w9nQvsKgTk,35852
 rxnn/training/dataset.py,sha256=7hTilFWPpqUEc6zNcMqBPjxFKxCfvTKKF3E8tVlwccQ,51250
+rxnn/training/ddp.py,sha256=VsNBjn3cY-uUj8hbsW7oKvb0_ZKnXnJ2KgObm-Mr9i4,836
 rxnn/training/models.py,sha256=2KhNT7yx0AgUke4nmsFqzQKx_YYp78QvsLWYZjWeUgQ,6812
-rxnn/training/mrl.py,sha256=MnLaYWxblc5cF261R5PNjIvddVQVNxyjAkEYtchBn9E,49299
+rxnn/training/mrl.py,sha256=Aimiiqf_4p6dp5Ty9pY9VwetySBS_OFpCQlcVHVkO4Q,55124
 rxnn/training/reward.py,sha256=7MTVdNm5HnWmt6zFDi3TAYmnVSL_-24riOoY2F7z4x8,11290
-rxnn/training/rl.py,sha256=j-KNLoZjhaEKasYNOc8DxHtwvknAgAJFwvXKot6otFA,3272
+rxnn/training/rl.py,sha256=eL3C0yryiNBgl_xb-D-5dyYUtK4V4-K4t3a60x5ir28,5142
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
 rxnn/training/utils.py,sha256=Bw8nZLKIt7NQpUVCYkb_79kWKChVFOYgYXwODo4SvNc,5718
@@ -32,7 +33,7 @@ rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
 rxnn/transformers/sampler.py,sha256=t6iiQTdLQ0TakUWnnhKkb5DKF2F_9-thXHBydDF3fxg,17389
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.29.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.29.dist-info/METADATA,sha256=WVEyKmyYbMOb5sm7vjjnCN9j8ABz0QfGJCYkQbWvwT8,25960
-rxnn-0.2.29.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.29.dist-info/RECORD,,
+rxnn-0.2.31.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.31.dist-info/METADATA,sha256=zxD2qPHL_QrFH1bYZrMv4odbXE4B_YIVEpGDzV2MYEI,25960
+rxnn-0.2.31.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.31.dist-info/RECORD,,