Upload lightbulb_wm.py

lightbulb_wm.py

@@ -0,0 +1,1292 @@
+
+import argparse
+import math
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import DataLoader
+import copy
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from torch.amp import autocast, GradScaler
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+
+# Set the device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Train World Model with Transformer outputs.')
+    parser.add_argument('--model_name', type=str, default='gpt2', help='Pretrained model name or path')
+    parser.add_argument('--dataset_name', type=str, default='wikitext', help='Dataset name from HuggingFace Datasets')
+    parser.add_argument('--dataset_config', type=str, default='wikitext-2-raw-v1', help='Dataset configuration name')
+    parser.add_argument('--batch_size', type=int, default=2, help='Batch size')
+    parser.add_argument('--num_epochs', type=int, default=3, help='Number of epochs')
+    parser.add_argument('--max_length', type=int, default=128, help='Maximum sequence length')
+    parser.add_argument('--mcts_iterations', type=int, default=5, help='Number of MCTS Iterations')
+    parser.add_argument('--mcts_exploration_constant', type=float, default=1.414, help='Learning rate')
+    parser.add_argument('--accumulation_steps', type=int, default=4, help='Gradient accumulation steps')
+    parser.add_argument('--learning_rate', type=float, default=1e-4, help='Learning rate')
+    parser.add_argument('--weight_decay', type=float, default=1e-2, help='Weight decay')
+    parser.add_argument('--alpha', type=float, default=0.1, help='Entropy regularization weight')
+    parser.add_argument('--beta', type=float, default=0.1, help='Variance regularization weight')
+    parser.add_argument('--max_grad_norm', type=float, default=1.0, help='Max gradient norm for clipping')
+    parser.add_argument('--save_dir', type=str, default='./models', help='Directory to save the models')
+    parser.add_argument('--temperature', type=float, default=1.0, help='Temperature parameter for entropy and variance')
+    parser.add_argument('--transformer_model_path', type=str, required=True, help='Path to the saved Transformer model')
+    args = parser.parse_args()
+    return args
+
+
+def load_data(args, tokenizer):
+    # Load the dataset
+    dataset = load_dataset(args.dataset_name, args.dataset_config)
+    
+    # Ensure the tokenizer has a padding token
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    def tokenize_function(examples):
+        return tokenizer(examples['text'], truncation=True, max_length=args.max_length)
+
+    tokenized_datasets = dataset.map(
+        tokenize_function,
+        batched=True,
+        num_proc=4,
+        remove_columns=dataset['train'].column_names,
+    )
+
+    # Build inputs and labels for language modeling
+    block_size = args.max_length
+
+    def group_texts(examples):
+        # Concatenate all texts
+        concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
+        total_length = len(concatenated_examples['input_ids'])
+        # We drop the small remainder
+        total_length = (total_length // block_size) * block_size
+        # Split by chunks of block_size
+        result = {
+            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+            for k, t in concatenated_examples.items()
+        }
+        result['labels'] = result['input_ids'].copy()
+        return result
+
+    lm_datasets = tokenized_datasets.map(
+        group_texts,
+        batched=True,
+        num_proc=4,
+    )
+
+    # Create DataLoader
+    train_dataset = lm_datasets['train']
+    eval_dataset = lm_datasets['validation'] if 'validation' in lm_datasets else lm_datasets['test']
+
+    data_collator = lambda data: {
+        'input_ids': torch.tensor([f['input_ids'] for f in data], dtype=torch.long),
+        'labels': torch.tensor([f['labels'] for f in data], dtype=torch.long)
+    }
+
+    train_loader = DataLoader(train_dataset, shuffle=True, batch_size=args.batch_size, collate_fn=data_collator)
+    eval_loader = DataLoader(eval_dataset, shuffle=False, batch_size=args.batch_size, collate_fn=data_collator)
+
+    return train_loader, eval_loader
+
+
+def save_all_models(transformer_model, representation_network, dynamics_network, prediction_network, action_encoder, save_dir, epoch):
+    """
+    Save all models to the specified directory.
+
+    Args:
+        transformer_model (nn.Module): Transformer model.
+        representation_network (nn.Module): Representation network.
+        dynamics_network (nn.Module): Dynamics network.
+        prediction_network (nn.Module): Prediction network.
+        action_encoder (nn.Module): Action encoder.
+        save_dir (str): Directory to save the models.
+        epoch (int): Current epoch number.
+    """
+    os.makedirs(save_dir, exist_ok=True)
+
+    torch.save(transformer_model.state_dict(), os.path.join(save_dir, f'transformer_model_epoch_{epoch}.pt'))
+    torch.save(representation_network.state_dict(), os.path.join(save_dir, f'representation_network_epoch_{epoch}.pt'))
+    torch.save(dynamics_network.state_dict(), os.path.join(save_dir, f'dynamics_network_epoch_{epoch}.pt'))
+    torch.save(prediction_network.state_dict(), os.path.join(save_dir, f'prediction_network_epoch_{epoch}.pt'))
+    torch.save(action_encoder.state_dict(), os.path.join(save_dir, f'action_encoder_epoch_{epoch}.pt'))
+
+    print(f"All models saved for epoch {epoch}.")
+
+
+class RotaryPositionalEncoding(nn.Module):
+    def __init__(self, d_model):
+        super(RotaryPositionalEncoding, self).__init__()
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, d_model, 2).float() / d_model))
+        self.register_buffer('inv_freq', inv_freq)
+
+    def forward(self, x):
+        seq_len, batch_size, _ = x.size()
+        t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
+        sinusoid_inp = torch.einsum("i,j->ij", t, self.inv_freq)
+        sin = sinusoid_inp.sin().unsqueeze(1)  # (seq_len, 1, d_model/2)
+        cos = sinusoid_inp.cos().unsqueeze(1)  # (seq_len, 1, d_model/2)
+
+        x1 = x[..., 0::2]
+        x2 = x[..., 1::2]
+
+        # Apply rotation
+        x_rotated = torch.zeros_like(x)
+        x_rotated[..., 0::2] = x1 * cos - x2 * sin
+        x_rotated[..., 1::2] = x1 * sin + x2 * cos
+
+        return x_rotated
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, d_model, num_heads):
+        super(MultiHeadAttention, self).__init__()
+        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
+        self.d_k = d_model // num_heads
+        self.num_heads = num_heads
+        self.linear_q = nn.Linear(d_model, d_model)
+        self.linear_k = nn.Linear(d_model, d_model)
+        self.linear_v = nn.Linear(d_model, d_model)
+        self.linear_out = nn.Linear(d_model, d_model)
+    
+    def forward(self, query, key, value, mask=None):
+        batch_size = query.size(0)
+        query = self.linear_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
+        key = self.linear_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
+        value = self.linear_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
+        
+        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.d_k)
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e4)
+        attn = F.softmax(scores, dim=-1)
+        output = torch.matmul(attn, value)
+        
+        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k)
+        return self.linear_out(output)
+
+
+class MoE(nn.Module):
+    def __init__(self, d_model, num_experts, d_ff, top_k=2, dropout=0.1):
+        super(MoE, self).__init__()
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.experts = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(d_model, d_ff),
+                nn.GELU() if i % 2 == 0 else nn.SiLU(),
+                nn.Linear(d_ff, d_model)
+            )
+            for i in range(num_experts)
+        ])
+        self.gate = nn.Linear(d_model, num_experts)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        batch_size, seq_len, d_model = x.size()
+        # Compute gating scores
+        gate_scores = self.gate(x)  # (batch_size, seq_len, num_experts)
+        top_k_scores, top_k_indices = torch.topk(gate_scores, self.top_k, dim=-1)  # (batch_size, seq_len, top_k)
+        top_k_scores = F.softmax(top_k_scores, dim=-1)  # (batch_size, seq_len, top_k)
+
+        # Initialize output
+        output = torch.zeros_like(x)
+
+        # Flatten batch and sequence dimensions
+        x_flat = x.view(-1, d_model)  # (batch_size * seq_len, d_model)
+        output_flat = output.view(-1, d_model)
+        top_k_indices_flat = top_k_indices.view(-1, self.top_k)  # (batch_size * seq_len, top_k)
+        top_k_scores_flat = top_k_scores.view(-1, self.top_k)  # (batch_size * seq_len, top_k)
+
+        for k in range(self.top_k):
+            expert_idx_flat = top_k_indices_flat[:, k]  # (batch_size * seq_len)
+            expert_scores_flat = top_k_scores_flat[:, k]  # (batch_size * seq_len)
+            for e in range(self.num_experts):
+                mask = (expert_idx_flat == e)  # Boolean mask
+                if mask.any():
+                    x_masked = x_flat[mask]  # Select tokens for expert e
+                    expert_output = self.experts[e](x_masked)  # Apply expert e
+                    output_flat[mask] += expert_scores_flat[mask].unsqueeze(-1) * expert_output
+
+        output = output_flat.view(batch_size, seq_len, d_model)
+        return self.dropout(output)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model, num_heads, d_ff, num_experts, dropout=0.1, top_k=2):
+        super(TransformerBlock, self).__init__()
+        self.self_attention = MultiHeadAttention(d_model, num_heads)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.cross_attention = MultiHeadAttention(d_model, num_heads)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.moe = MoE(d_model, num_experts, d_ff, top_k, dropout)
+        self.norm3 = nn.LayerNorm(d_model)
+
+    def forward(self, x, mask=None, enc_output=None, enc_mask=None):
+        # Self-attention
+        attn_output = self.self_attention(x, x, x, mask)
+        x = self.norm1(x + attn_output)
+        # Cross-attention (only in decoder)
+        if enc_output is not None:
+            cross_attn_output = self.cross_attention(x, enc_output, enc_output, enc_mask)
+            x = self.norm2(x + cross_attn_output)
+        # Feedforward/MoE
+        moe_output = self.moe(x)
+        return self.norm3(x + moe_output)
+
+
+class Transformer(nn.Module):
+    def __init__(self, input_dim, d_model, num_heads, num_layers, d_ff, num_experts, output_dim, dropout=0.1, top_k=2):
+        super(Transformer, self).__init__()
+        self.embedding = nn.Embedding(input_dim, d_model, padding_idx=input_dim - 1)
+        self.rotary_positional_encoding = RotaryPositionalEncoding(d_model)
+        self.encoder_layers = nn.ModuleList(
+            [TransformerBlock(d_model, num_heads, d_ff, num_experts, dropout, top_k) for _ in range(num_layers)]
+        )
+        self.decoder_layers = nn.ModuleList(
+            [TransformerBlock(d_model, num_heads, d_ff, num_experts, dropout, top_k) for _ in range(num_layers)]
+        )
+        self.output_layer = nn.Linear(d_model, output_dim)
+        self.d_model = d_model
+
+    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
+        # Encoder
+        src = self.embedding(src) * math.sqrt(self.d_model)
+        src = src.transpose(0, 1)  # (batch_size, seq_len, d_model) -> (seq_len, batch_size, d_model)
+        src = self.rotary_positional_encoding(src)
+        src = src.transpose(0, 1)  # (seq_len, batch_size, d_model) -> (batch_size, seq_len, d_model)
+        for layer in self.encoder_layers:
+            src = layer(src, src_mask)
+
+        # Decoder
+        tgt = self.embedding(tgt) * math.sqrt(self.d_model)
+        tgt = tgt.transpose(0, 1)
+        tgt = self.rotary_positional_encoding(tgt)
+        tgt = tgt.transpose(0, 1)
+        for layer in self.decoder_layers:
+            tgt = layer(tgt, tgt_mask, src, src_mask)
+        output = self.output_layer(tgt)
+        return output
+
+    def generate(self, src, tokenizer, max_length=20, temperature=1.0):
+        """
+        Generate sequences using differentiable sampling (Gumbel-Softmax).
+
+        Args:
+            src (torch.Tensor): Source input tensor of shape (batch_size, seq_len)
+            tokenizer (transformers.PreTrainedTokenizer): Tokenizer to access special tokens
+            max_length (int): Maximum length of the generated sequence
+            temperature (float): Temperature parameter for Gumbel-Softmax
+
+        Returns:
+            torch.Tensor: Generated sequences of shape (batch_size, max_length)
+            torch.Tensor: Entropy values for each time step
+            torch.Tensor: Variance values for each time step
+        """
+        batch_size = src.size(0)
+
+        # Encode the source
+        src_enc = self.embedding(src) * math.sqrt(self.d_model)
+        src_enc = src_enc.transpose(0, 1)
+        src_enc = self.rotary_positional_encoding(src_enc)
+        src_enc = src_enc.transpose(0, 1)
+        for layer in self.encoder_layers:
+            src_enc = layer(src_enc)
+
+        # Initialize decoder input with <sos> tokens
+        tgt_seq = torch.full((batch_size, 1), tokenizer.bos_token_id, dtype=torch.long, device=src.device)
+        entropies = []
+        variances = []
+
+        for _ in range(max_length):
+            tgt_emb = self.embedding(tgt_seq) * math.sqrt(self.d_model)
+            tgt_emb = tgt_emb.transpose(0, 1)
+            tgt_emb = self.rotary_positional_encoding(tgt_emb)
+            tgt_emb = tgt_emb.transpose(0, 1)
+            tgt_dec = tgt_emb
+            for layer in self.decoder_layers:
+                tgt_dec = layer(tgt_dec, None, src_enc, None)
+            output = self.output_layer(tgt_dec)  # (batch_size, seq_len, vocab_size)
+            logits = output[:, -1, :]  # Get logits for the last time step
+
+            # Compute token probabilities
+            probs = F.softmax(logits / temperature, dim=-1)  # (batch_size, vocab_size)
+
+            # Compute entropy
+            entropy = -torch.sum(probs * torch.log(probs + 1e-9), dim=-1)  # (batch_size)
+            entropies.append(entropy)
+
+            # Sample token using Gumbel-Softmax
+            gumbel_noise = -torch.log(-torch.log(torch.rand_like(probs) + 1e-9) + 1e-9)
+            y = (logits + gumbel_noise) / temperature
+            y = F.softmax(y, dim=-1)  # (batch_size, vocab_size)
+
+            # Compute variance
+            variance = torch.var(y, dim=-1)  # (batch_size)
+            variances.append(variance)
+
+            # Get token indices (argmax for hard selection)
+            next_tokens = torch.argmax(y, dim=-1, keepdim=True)  # (batch_size, 1)
+            tgt_seq = torch.cat([tgt_seq, next_tokens], dim=1)
+
+        # Stack entropies and variances
+        entropies = torch.stack(entropies, dim=1)  # (batch_size, max_length)
+        variances = torch.stack(variances, dim=1)  # (batch_size, max_length)
+
+        return tgt_seq[:, 1:], entropies, variances  # Exclude the initial <sos> token
+
+
+# Objective Functions
+
+class InfoNCE_Loss(nn.Module):
+    def __init__(self, temperature=0.07):
+        super(InfoNCE_Loss, self).__init__()
+        self.temperature = temperature
+        self.cross_entropy = nn.CrossEntropyLoss()
+    
+    def forward(self, z_i, z_j):
+        """
+        Args:
+            z_i (torch.Tensor): Flattened representations from view i, shape (2n, embed_dim)
+            z_j (torch.Tensor): Flattened representations from view j, shape (2n, embed_dim)
+        
+        Returns:
+            torch.Tensor: InfoNCE loss
+        """
+        n = z_i.size(0)
+        z = torch.cat([z_i, z_j], dim=0)  # Shape: (2n, embed_dim)
+
+        z = F.normalize(z, dim=1)
+        similarity_matrix = torch.matmul(z, z.T)  # Shape: (2n, 2n)
+
+        # Create a mask to exclude self-similarity
+        mask = torch.eye(2 * n, device=z.device, dtype=torch.bool)
+        similarity_matrix = similarity_matrix.masked_fill(mask, -1e4)  # Use a manageable negative value
+
+        # Create labels for contrastive learning
+        labels = torch.arange(n, device=z.device)
+        labels = torch.cat([labels + n, labels], dim=0)  # Shape: (2n,)
+
+        # Apply temperature scaling
+        similarity_matrix /= self.temperature
+
+        # Compute cross-entropy loss
+        loss = self.cross_entropy(similarity_matrix, labels)
+        return loss
+
+
+
+class CovarianceRegularization(nn.Module):
+    def __init__(self, lambda_reg=1e-3):
+        super(CovarianceRegularization, self).__init__()
+        self.lambda_reg = lambda_reg
+    
+    def forward(self, embeddings):
+        """
+        Args:
+            embeddings (torch.Tensor): Embedding tensor, shape (batch_size, embed_dim)
+        
+        Returns:
+            torch.Tensor: Covariance regularization loss
+        """
+        batch_size, embed_dim = embeddings.size()
+        mean = embeddings.mean(dim=0)
+        embeddings_centered = embeddings - mean
+        cov = (embeddings_centered.T @ embeddings_centered) / (batch_size - 1)
+        cov_loss = torch.sum(cov ** 2) - torch.sum(torch.diag(cov) ** 2)
+        return self.lambda_reg * cov_loss
+
+
+class DynamicsPerformanceLoss(nn.Module):
+    def __init__(self, lambda_var=1e-3):
+        super(DynamicsPerformanceLoss, self).__init__()
+        self.lambda_var = lambda_var
+    
+    def forward(self, true_next_state, predicted_next_state):
+        """
+        Args:
+            true_next_state (torch.Tensor): Ground truth next state, shape (batch_size, state_dim)
+            predicted_next_state (torch.Tensor): Predicted next state, shape (batch_size, state_dim)
+        
+        Returns:
+            torch.Tensor: Dynamics performance loss
+        """
+        mse_loss = F.mse_loss(predicted_next_state, true_next_state)
+        variance_loss = torch.var(predicted_next_state, dim=0).mean()
+        return mse_loss + self.lambda_var * variance_loss
+
+
+class ThoughtConsistencyLoss(nn.Module):
+    def __init__(self):
+        super(ThoughtConsistencyLoss, self).__init__()
+    
+    def forward(self, true_next_state, perturbed_next_state):
+        """
+        Args:
+            true_next_state (torch.Tensor): Ground truth next state, shape (batch_size, state_dim)
+            perturbed_next_state (torch.Tensor): Perturbed next state, shape (batch_size, state_dim)
+        
+        Returns:
+            torch.Tensor: Thought-consistency loss
+        """
+        return F.mse_loss(true_next_state, perturbed_next_state)
+
+
+class PolicyValueJointLoss(nn.Module):
+    def __init__(self, lambda_value=0.5):
+        super(PolicyValueJointLoss, self).__init__()
+        self.lambda_value = lambda_value
+        self.cross_entropy = nn.CrossEntropyLoss()
+        self.mse_loss = nn.MSELoss()
+    
+    def forward(self, policy_logits, true_policy, value_pred, true_value):
+        """
+        Args:
+            policy_logits (torch.Tensor): Logits from the policy network, shape (batch_size * seq_len, num_actions)
+            true_policy (torch.Tensor): Ground truth policy, shape (batch_size * seq_len, num_actions)
+            value_pred (torch.Tensor): Predicted values, shape (batch_size * seq_len)
+            true_value (torch.Tensor): Ground truth values, shape (batch_size * seq_len)
+        
+        Returns:
+            torch.Tensor: Combined policy and value loss
+        """
+        policy_logits = policy_logits.view(-1, policy_logits.size(-1))
+        true_policy = true_policy.view(-1, true_policy.size(-1))
+        value_pred = value_pred.view(-1)
+        true_value = true_value.view(-1)
+
+        policy_loss = self.cross_entropy(policy_logits, true_policy.argmax(dim=1))
+        value_loss = self.mse_loss(value_pred, true_value)
+        return policy_loss + self.lambda_value * value_loss
+
+
+
+class ActionDiversityReward(nn.Module):
+    def __init__(self, lambda_div=1e-3):
+        super(ActionDiversityReward, self).__init__()
+        self.lambda_div = lambda_div
+    
+    def forward(self, action_embeddings):
+        """
+        Args:
+            action_embeddings (torch.Tensor): Embeddings of actions, shape (batch_size, embed_dim)
+        
+        Returns:
+            torch.Tensor: Action diversity loss
+        """
+        similarity_matrix = F.cosine_similarity(action_embeddings.unsqueeze(1), action_embeddings.unsqueeze(0), dim=2)
+        # Zero out self-similarity
+        similarity_matrix = similarity_matrix - torch.eye(similarity_matrix.size(0)).to(action_embeddings.device)
+        diversity_loss = torch.sum(similarity_matrix ** 2)
+        return self.lambda_div * diversity_loss
+
+
+class ExpectedThoughtValueLoss(nn.Module):
+    def __init__(self):
+        super(ExpectedThoughtValueLoss, self).__init__()
+    
+    def forward(self, mcts_best_values):
+        """
+        Args:
+            mcts_best_values (torch.Tensor): Best values from MCTS, shape (batch_size)
+        
+        Returns:
+            torch.Tensor: ETV loss
+        """
+        return -mcts_best_values.mean()
+
+
+class ExplorationRegularization(nn.Module):
+    def __init__(self, lambda_expl=1e-3):
+        super(ExplorationRegularization, self).__init__()
+        self.lambda_expl = lambda_expl
+    
+    def forward(self, visit_counts):
+        """
+        Args:
+            visit_counts (torch.Tensor): Visit counts for actions, shape (batch_size, num_actions)
+        
+        Returns:
+            torch.Tensor: Exploration regularization loss
+        """
+        reward = torch.sum(1.0 / (visit_counts + 1), dim=-1)
+        return self.lambda_expl * reward.mean()
+
+
+class KL_DivergenceLoss(nn.Module):
+    def __init__(self):
+        super(KL_DivergenceLoss, self).__init__()
+    
+    def forward(self, old_policy, new_policy):
+        """
+        Args:
+            old_policy (torch.Tensor): Old policy probabilities, shape (batch_size, num_actions)
+            new_policy (torch.Tensor): New policy probabilities, shape (batch_size, num_actions)
+        
+        Returns:
+            torch.Tensor: KL divergence loss
+        """
+        kl_div = F.kl_div(new_policy.log(), old_policy, reduction='batchmean')
+        return kl_div
+
+# MuZero
+
+class ActionEncoder(nn.Module):
+    def __init__(self, vocab_size, embed_dim):
+        super(ActionEncoder, self).__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+    
+    def forward(self, action_sequences):
+        """
+        Args:
+            action_sequences (torch.Tensor): Tensor of shape (batch_size, seq_len)
+        
+        Returns:
+            torch.Tensor: Encoded actions of shape (batch_size, seq_len, embed_dim)
+        """
+        return self.embedding(action_sequences) #.half()  # Convert to half-precision
+
+class RepresentationNetwork(nn.Module):
+    def __init__(self, vocab_dim, d_model, state_dim):
+        super(RepresentationNetwork, self).__init__()
+        self.proj = nn.Linear(vocab_dim, d_model)  # Project from vocab_dim to d_model
+        self.linear = nn.Linear(d_model, state_dim)  # Project from d_model to state_dim
+        self.norm = nn.LayerNorm(state_dim)
+    
+    def forward(self, transformer_output):
+        """
+        Args:
+            transformer_output (torch.Tensor): Shape (batch_size, seq_len, vocab_dim)
+        
+        Returns:
+            torch.Tensor: Encoded state of shape (batch_size, seq_len, state_dim)
+        """
+        # First project down from vocab_dim to d_model
+        projected_output = self.proj(transformer_output)
+        # Then project down from d_model to state_dim
+        state = self.linear(projected_output)
+        state = self.norm(state)
+        return state
+
+
+class DynamicsNetwork(nn.Module):
+    def __init__(self, state_dim, action_dim, hidden_dim):
+        super(DynamicsNetwork, self).__init__()
+        self.rms_norm = nn.LayerNorm(state_dim)
+        self.fc1 = nn.Linear(state_dim + action_dim, hidden_dim)
+        self.activation = nn.GELU()
+        self.fc2 = nn.Linear(hidden_dim, state_dim)
+    
+    def forward(self, state, action):
+        """
+        Args:
+            state (torch.Tensor): Current state, shape (batch_size, seq_len, state_dim)
+            action (torch.Tensor): Action embedding, shape (batch_size, seq_len, action_dim)
+        
+        Returns:
+            torch.Tensor: Predicted next state, shape (batch_size, seq_len, state_dim)
+        """
+        norm_state = self.rms_norm(state)
+        combined = torch.cat([norm_state, action], dim=-1)
+        hidden = self.activation(self.fc1(combined))
+        next_state = self.fc2(hidden)
+        return next_state
+
+class PredictionNetwork(nn.Module):
+    def __init__(self, state_dim, policy_dim, value_dim):
+        super(PredictionNetwork, self).__init__()
+        self.state_dim = state_dim
+        self.rms_norm = nn.LayerNorm(state_dim)
+        self.policy_head = nn.Linear(state_dim, policy_dim)
+        self.value_head = nn.Linear(state_dim, value_dim)
+    
+    def forward(self, state):
+        """
+        Args:
+            state (torch.Tensor): Predicted state, shape (batch_size, seq_len, state_dim)
+        
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: Policy logits and value estimates
+        """
+        norm_state = self.rms_norm(state)
+        policy_logits = self.policy_head(norm_state)
+        value_estimates = self.value_head(norm_state)
+        return policy_logits, value_estimates
+
+
+class MCTSNode:
+    def __init__(self, state, parent=None, action=None):
+        """
+        Initialize an MCTS node.
+
+        Args:
+            state (State): The current state representation.
+            parent (MCTSNode, optional): The parent node. Defaults to None.
+            action (int, optional): The action taken to reach this node. Defaults to None.
+        """
+        self.state = state  # Instance of State class
+        self.parent = parent  # Parent MCTSNode
+        self.action = action  # Action taken to reach this node
+        self.children = {}  # Dict mapping actions to MCTSNode
+        self.visit_count = 0
+        self.value_sum = 0.0
+        self.prior = 0.0  # Prior probability from policy network
+
+    def expand(self, actions, priors):
+        """
+        Expand the node with possible actions and their priors.
+
+        Args:
+            actions (list): List of possible actions (action indices).
+            priors (list): List of prior probabilities corresponding to actions.
+        """
+        for action, prior in zip(actions, priors):
+            if action not in self.children:
+                child_state = self.state.apply_action(action)  # Apply action to get new state
+                child_node = MCTSNode(state=child_state, parent=self, action=action)
+                child_node.prior = float(prior)  # Ensure that prior is a float value
+                self.children[action] = child_node
+
+    def is_leaf(self):
+        """
+        Check if the node is a leaf node (i.e., has no children).
+
+        Returns:
+            bool: True if leaf, False otherwise.
+        """
+        return len(self.children) == 0
+
+    def ucb_score(self, total_visits, exploration_constant=math.sqrt(2)):
+        """
+        Calculate the UCB (Upper Confidence Bound) score for the node.
+
+        Args:
+            total_visits (int): Total number of visits to the parent node.
+            exploration_constant (float, optional): Exploration parameter. Defaults to math.sqrt(2).
+
+        Returns:
+            float: The UCB score.
+        """
+        if self.visit_count == 0:
+            return float('inf')
+        average_value = self.value_sum / self.visit_count
+        exploration_term = exploration_constant * self.prior * math.sqrt(total_visits) / (1 + self.visit_count)
+        return average_value + exploration_term
+
+class MCTS:
+    def __init__(self, prediction_network, dynamics_network, action_encoder, num_iterations=10, exploration_constant=math.sqrt(2)):
+        """
+        Initialize the MCTS.
+
+        Args:
+            prediction_network (nn.Module): The Prediction Network.
+            dynamics_network (nn.Module): The Dynamics Network.
+            num_iterations (int): Number of MCTS iterations per search.
+            exploration_constant (float): Exploration parameter for UCB.
+        """
+        self.action_encoder = action_encoder
+        self.prediction_network = prediction_network
+        self.dynamics_network = dynamics_network
+        self.num_iterations = num_iterations
+        self.exploration_constant = exploration_constant
+
+    def search(self, root_state):
+        """
+        Perform MCTS starting from the root_state.
+
+        Args:
+            root_state: The initial state from which to start MCTS.
+
+        Returns:
+            The best action determined by MCTS.
+        """
+        self.root = MCTSNode(state=root_state)
+
+        for _ in range(self.num_iterations):
+            node = self.select(self.root)
+            value = self.evaluate(node)
+            self.backpropagate(node, value)
+
+        return self.best_action()
+
+    def select(self, node):
+        """
+        Traverse the tree to select a node for evaluation.
+
+        Args:
+            node: The starting node for selection.
+
+        Returns:
+            The node selected for evaluation.
+        """
+        while not node.is_leaf():
+            best_action, best_node = max(node.children.items(),
+                                         key=lambda item: item[1].ucb_score(node.visit_count, self.exploration_constant))
+            node = best_node
+        return node
+
+    def evaluate(self, node):
+        """
+        Evaluate the node by expanding it and predicting its value.
+
+        Args:
+            node: The node to evaluate.
+
+        Returns:
+            The value estimate of the node.
+        """
+        # Use the prediction network to get policy logits and value estimate
+        policy_logits, value_estimate = self.prediction_network(node.state.representation)
+
+        # Convert logits to probabilities
+        policy = F.softmax(policy_logits, dim=-1).detach().cpu().numpy()
+
+        # Expand the node with possible actions and their priors
+        actions = list(range(policy.shape[-1]))  # Assuming actions are indexed from 0 to num_actions-1
+        priors = policy[0].flatten().tolist()  # Convert to a 1D list of floats
+
+        node.expand(actions, priors)
+
+        return value_estimate.mean().item()
+
+
+    def backpropagate(self, node, value):
+        """
+        Backpropagate the value up the tree.
+
+        Args:
+            node: The node to start backpropagation from.
+            value (float): The value to backpropagate.
+        """
+        while node is not None:
+            node.visit_count += 1
+            node.value_sum += value
+            node = node.parent
+
+    def best_action(self):
+        """
+        Choose the action with the highest visit count.
+
+        Returns:
+            The best action.
+        """
+        best_child = max(self.root.children.values(), key=lambda n: n.visit_count)
+        return best_child.action
+
+class State:
+    def __init__(self, representation, dynamics_network, action_encoder):
+        """
+        Initialize the State.
+
+        Args:
+            representation (torch.Tensor): Encoded state representation, shape (batch_size, seq_len, state_dim)
+            dynamics_network (nn.Module): The Dynamics Network to predict next states
+            action_encoder (nn.Module): The Action Encoder to encode actions
+        """
+        self.representation = representation  # Shape: (batch_size, seq_len, state_dim)
+        self.dynamics_network = dynamics_network  # Reference to Dynamics Network
+        self.action_encoder = action_encoder
+
+    def apply_action(self, action):
+        """
+        Apply an action to the current state to get a new state.
+
+        Args:
+            action (int): The action to apply (e.g., token index)
+
+        Returns:
+            State: The new state after applying the action
+        """
+        # Create action sequence filled with action index
+        batch_size, seq_len, _ = self.representation.size()
+        action_sequence = torch.full((batch_size, seq_len), action, dtype=torch.long, device=self.representation.device)
+        # Encode action
+        action_embedding = self.action_encoder(action_sequence)
+        # Predict the next state using the Dynamics Network
+        with torch.no_grad():
+            next_state_representation = self.dynamics_network(self.representation, action_embedding)
+        return State(next_state_representation, self.dynamics_network, self.action_encoder)
+
+
+
+
+class PPOAgent:
+    def __init__(self, policy_network, optimizer, clip_epsilon=0.2, entropy_coef=0.01, value_coef=0.5):
+        self.policy_network = policy_network
+        self.optimizer = optimizer
+        self.clip_epsilon = clip_epsilon
+        self.entropy_coef = entropy_coef
+        self.value_coef = value_coef
+    
+    def compute_loss(self, states, old_log_probs, actions, returns, advantages):
+        # Get policy logits and value estimates
+        policy_logits, value_estimates = self.policy_network(states)
+        batch_size, seq_len, num_actions = policy_logits.size()
+        
+        # Flatten tensors
+        policy_logits = policy_logits.view(-1, num_actions)  # Shape: (batch_size * seq_len, num_actions)
+        value_estimates = value_estimates.view(-1)           # Shape: (batch_size * seq_len)
+        actions = actions.view(-1)                           # Shape: (batch_size * seq_len)
+        old_log_probs = old_log_probs.view(-1)               # Shape: (batch_size * seq_len)
+        returns = returns.view(-1)                           # Shape: (batch_size * seq_len)
+        advantages = advantages.view(-1)                     # Shape: (batch_size * seq_len)
+        
+        # Compute new log probabilities
+        new_log_probs_all = F.log_softmax(policy_logits, dim=-1)  # Shape: (batch_size * seq_len, num_actions)
+        new_log_probs = new_log_probs_all.gather(1, actions.unsqueeze(-1)).squeeze(-1)  # Shape: (batch_size * seq_len)
+        
+        # Compute ratios
+        ratios = torch.exp(new_log_probs - old_log_probs)
+        
+        # PPO surrogate loss
+        surr1 = ratios * advantages
+        surr2 = torch.clamp(ratios, 1 - self.clip_epsilon, 1 + self.clip_epsilon) * advantages
+        policy_loss = -torch.min(surr1, surr2).mean()
+        
+        # Value loss
+        value_loss = F.mse_loss(value_estimates, returns)
+        
+        # Entropy loss
+        entropy = -(new_log_probs * torch.exp(new_log_probs)).mean()
+        
+        # Total loss
+        total_loss = policy_loss + self.value_coef * value_loss - self.entropy_coef * entropy
+        return total_loss
+
+
+
+
+def compute_loss_world_model(predicted_next_state, true_next_state, policy_logits, true_policy, value_estimates, true_value, 
+                             alpha, beta, temperature, lambda_reg, lambda_var, lambda_div, lambda_expl):
+    """
+    Compute the combined loss for the World Model.
+
+    Args:
+        predicted_next_state (torch.Tensor): Predicted next state, shape (batch_size, state_dim)
+        true_next_state (torch.Tensor): Ground truth next state, shape (batch_size, state_dim)
+        policy_logits (torch.Tensor): Policy logits, shape (batch_size, num_actions)
+        true_policy (torch.Tensor): Ground truth policy, shape (batch_size, num_actions)
+        value_estimates (torch.Tensor): Value estimates, shape (batch_size)
+        true_value (torch.Tensor): Ground truth value, shape (batch_size)
+        alpha (float): Entropy regularization weight
+        beta (float): Variance regularization weight
+        temperature (float): Temperature parameter
+        lambda_reg (float): Covariance regularization weight
+        lambda_var (float): Dynamics variance loss weight
+        lambda_div (float): Action diversity reward weight
+        lambda_expl (float): Exploration regularization weight
+
+    Returns:
+        torch.Tensor: Combined loss
+    """
+    # Cross-entropy loss
+    ce_loss = F.cross_entropy(policy_logits, true_policy.argmax(dim=1))
+    
+    # Entropy loss
+    probs = F.softmax(policy_logits / temperature, dim=-1)  # (batch_size, num_actions)
+    entropy = -torch.sum(probs * torch.log(probs + 1e-9), dim=-1)  # (batch_size)
+    entropy_loss = -alpha * torch.mean(entropy)
+    
+    # Variance loss
+    variance = torch.var(probs, dim=-1)  # (batch_size)
+    variance_loss = -beta * torch.mean(variance)
+    
+    # Covariance Regularization
+    cov_reg = CovarianceRegularization(lambda_reg)(predicted_next_state)
+    
+    # Dynamics Performance Loss
+    dynamics_loss = DynamicsPerformanceLoss(lambda_var)(true_next_state, predicted_next_state)
+    
+    # Thought-Consistency Loss
+    perturbed_next_state = predicted_next_state + torch.randn_like(predicted_next_state) * 0.01
+    thought_loss = ThoughtConsistencyLoss()(true_next_state, perturbed_next_state)
+    
+    # Policy-Value Joint Loss
+    pv_loss = PolicyValueJointLoss()(policy_logits, true_policy, value_estimates, true_value)
+    
+    # Action Diversity Reward
+    action_embeddings = predicted_next_state  # Assuming actions are derived from state
+    action_diversity = ActionDiversityReward(lambda_div)(action_embeddings)
+    
+    # Expected Thought Value (ETV) Loss
+    # Placeholder: Replace with actual MCTS best values
+    mcts_best_values = torch.zeros(value_estimates.size(0)).to(device)
+    etv = ExpectedThoughtValueLoss()(mcts_best_values)
+    
+    # Exploration Regularization
+    # Placeholder: Replace with actual visit counts
+    visit_counts = torch.ones(predicted_next_state.size(0), input_dim).to(device)
+    exploration = ExplorationRegularization(lambda_expl)(visit_counts)
+    
+    # KL Divergence Regularization
+    # Placeholder: Replace with actual old and new policies
+    old_policy = F.softmax(policy_logits.detach(), dim=-1)
+    new_policy = F.softmax(policy_logits, dim=-1)
+    kl_loss = KL_DivergenceLoss()(old_policy, new_policy)
+    
+    # Total Loss
+    total_loss = (
+        ce_loss +
+        entropy_loss +
+        variance_loss +
+        cov_reg +
+        dynamics_loss +
+        thought_loss +
+        pv_loss +
+        action_diversity +
+        etv +
+        exploration +
+        kl_loss
+    )
+    
+    return total_loss
+
+
+def train_epoch_world_model(world_model_components, train_loader, optimizer, scheduler, scaler, args, model_transformer, state_dim, embed_dim, input_dim):
+    representation_network, dynamics_network, prediction_network, action_encoder, ppo_agent = world_model_components
+    representation_network.train()
+    dynamics_network.train()
+    prediction_network.train()
+    action_encoder.train()
+    ppo_agent.policy_network.train()
+
+    mcts = MCTS(prediction_network, dynamics_network, action_encoder, num_iterations=args.mcts_iterations, exploration_constant=args.mcts_exploration_constant)
+
+    total_loss = 0.0
+    optimizer.zero_grad()
+    print(f"Starting World Model training epoch with {len(train_loader)} batches...")
+
+    for i, batch in enumerate(train_loader):
+        print(f"Processing batch {i+1}/{len(train_loader)}...")
+        
+        # Ensure batches are on the appropriate device for the Transformer
+        src_batch = batch['input_ids'].to('cpu')  # Move to CPU for Transformer model
+        tgt_batch = batch['labels'].to('cpu')     # Move to CPU for Transformer model
+
+        with autocast(device_type='cuda'):
+            print("Forward pass through Transformer (frozen)...")
+            with torch.no_grad():
+                transformer_output = model_transformer(src_batch, tgt_batch[:, :-1])
+
+            # Move transformer output to the GPU for further processing
+            transformer_output = transformer_output.to(device)
+
+            # Encode actions directly on the GPU
+            encoded_actions = action_encoder(tgt_batch[:, :-1].to(device))  # Move labels to GPU for encoding
+
+            # World Model - Representation
+            state_representation = representation_network(transformer_output)  # On GPU
+
+            batch_size, seq_len, _ = state_representation.size()
+
+            # Initialize list to collect predicted next states for the batch
+            predicted_next_states = []
+
+            # Iterate over each sample in the batch for MCTS
+            for b in range(batch_size):
+                # Create a State instance for the current sample
+                current_state = State(state_representation[b].unsqueeze(0), dynamics_network, action_encoder)
+
+                # Perform MCTS to find the best action
+                best_action = mcts.search(current_state)
+
+                # Create action sequence filled with best_action
+                action_sequence = torch.full((1, seq_len), best_action, dtype=torch.long, device=device)
+
+                # Get action embedding
+                action_embedding = action_encoder(action_sequence)
+
+                # Apply dynamics network
+                predicted_next_state = dynamics_network(current_state.representation, action_embedding)
+
+                predicted_next_states.append(predicted_next_state)
+
+            # Concatenate predicted next states to form a batch
+            predicted_next_state_batch = torch.cat(predicted_next_states, dim=0)
+
+            # Prediction Network - Policy logits and value
+            policy_logits, value_estimates = prediction_network(predicted_next_state_batch)
+
+
+            # Define true_policy and true_value as placeholders on the GPU
+            true_policy = torch.zeros_like(policy_logits).to(device)
+            true_value = torch.zeros_like(value_estimates).to(device)
+
+            # Compute PPO loss
+            actions = torch.argmax(policy_logits, dim=-1)
+            old_log_probs = torch.zeros_like(actions, dtype=torch.float32).to(device)
+            returns = torch.zeros_like(actions, dtype=torch.float32).to(device)
+            advantages = torch.zeros_like(actions, dtype=torch.float32).to(device)
+
+            # Compute PPO loss using states
+            ppo_loss = ppo_agent.compute_loss(state_representation, old_log_probs, actions, returns, advantages)
+
+            # Compute InfoNCE Loss
+            z_i = state_representation.view(batch_size * seq_len, state_dim)  # Shape: (batch_size * seq_len, state_dim)
+            z_j = F.dropout(z_i, p=0.1, training=True)
+            info_nce = InfoNCE_Loss()(z_i, z_j)
+
+            # Compute other losses
+            covariance = CovarianceRegularization()(predicted_next_state_batch.view(-1, predicted_next_state_batch.size(-1)))
+            dynamics_loss = DynamicsPerformanceLoss()(torch.zeros_like(predicted_next_state_batch).to(device), predicted_next_state_batch)
+            perturbed_next_state = predicted_next_state_batch + torch.randn_like(predicted_next_state_batch) * 0.01
+            thought_loss = ThoughtConsistencyLoss()(torch.zeros_like(predicted_next_state_batch).to(device), perturbed_next_state)
+            pv_loss = PolicyValueJointLoss()(policy_logits, true_policy, value_estimates.squeeze(-1), true_value.squeeze(-1))
+            action_diversity = ActionDiversityReward()(encoded_actions.view(-1, embed_dim))
+            mcts_best_values = torch.zeros(actions.size(0)).to(device)
+            etv = ExpectedThoughtValueLoss()(mcts_best_values)
+            visit_counts = torch.ones(actions.size(0), policy_logits.size(-1)).to(device)
+            exploration = ExplorationRegularization()(visit_counts)
+            old_policy = F.softmax(policy_logits.detach(), dim=-1)
+            new_policy = F.softmax(policy_logits, dim=-1)
+            kl_loss = KL_DivergenceLoss()(old_policy, new_policy)
+
+            # Total Loss
+            loss = (
+                ppo_loss +
+                info_nce +
+                covariance +
+                dynamics_loss +
+                thought_loss +
+                pv_loss +
+                action_diversity +
+                etv +
+                exploration +
+                kl_loss
+            )
+            loss = loss / args.accumulation_steps
+
+        print("Backward pass...")
+        scaler.scale(loss).backward()
+
+        if (i + 1) % args.accumulation_steps == 0:
+            print("Gradient clipping...")
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(
+                [param for group in optimizer.param_groups for param in group['params']],
+                args.max_grad_norm
+            )
+
+            print("Optimizer step...")
+            scaler.step(optimizer)
+            scaler.update()
+
+            print("Zeroing gradients...")
+            optimizer.zero_grad()
+
+            print("Updating learning rate...")
+            scheduler.step()
+
+        total_loss += loss.item() * args.accumulation_steps
+        print(f"Batch {i+1} completed. Current loss: {loss.item():.4f}")
+
+    avg_loss = total_loss / len(train_loader)
+    print(f"World Model training epoch completed. Average loss: {avg_loss:.4f}")
+    return avg_loss
+
+
+
+def evaluate_world_model(world_model_components, model_transformer, eval_loader, args):
+    representation_network, dynamics_network, prediction_network, action_encoder, ppo_agent = world_model_components
+    representation_network.eval()
+    dynamics_network.eval()
+    prediction_network.eval()
+    action_encoder.eval()
+    ppo_agent.policy_network.eval()
+    
+    total_loss = 0.0
+    with torch.no_grad():
+        for batch in eval_loader:
+            src_batch = batch['input_ids'].to(device)
+            tgt_batch = batch['labels'].to(device)
+    
+            # Forward pass through Transformer (on CPU)
+            transformer_output = model_transformer(src_batch, tgt_batch[:, :-1])
+    
+            # Encode actions
+            encoded_actions = action_encoder(tgt_batch[:, :-1].to(device))  # Move to GPU if necessary
+    
+            # World Model - Representation
+            state = representation_network(transformer_output.to(device))
+    
+            # Dynamics Network - Predict next state
+            predicted_next_state = dynamics_network(state, encoded_actions)
+    
+            # Prediction Network - Policy logits and value
+            policy_logits, value_estimates = prediction_network(predicted_next_state)
+    
+            # Placeholder: Define true_policy and true_value
+            # Replace these with actual targets from your environment or dataset
+            true_policy = torch.zeros_like(policy_logits).to(device)
+            true_value = torch.zeros_like(value_estimates).to(device)
+    
+            # Compute PPO loss
+            # Placeholder: Replace with actual old_log_probs, actions, returns, and advantages
+            old_log_probs = torch.zeros_like(policy_logits).to(device)
+            actions = torch.argmax(policy_logits, dim=-1)
+            returns = torch.zeros(actions.size(0)).to(device)
+            advantages = torch.zeros(actions.size(0)).to(device)
+    
+            ppo_loss = ppo_agent.compute_loss(old_log_probs, actions, returns, advantages)
+    
+            # Compute other losses
+            info_nce = InfoNCE_Loss()(state, state)  # Placeholder: replace with actual positive pairs
+            covariance = CovarianceRegularization()(predicted_next_state.view(-1, predicted_next_state.size(-1)))
+            dynamics_loss = DynamicsPerformanceLoss()(torch.zeros_like(predicted_next_state).to(device), predicted_next_state)
+            perturbed_next_state = predicted_next_state + torch.randn_like(predicted_next_state) * 0.01
+            thought_loss = ThoughtConsistencyLoss()(torch.zeros_like(predicted_next_state).to(device), perturbed_next_state)
+            pv_loss = PolicyValueJointLoss()(policy_logits, true_policy, value_estimates.squeeze(-1), true_value.squeeze(-1))
+            action_diversity = ActionDiversityReward()(encoded_actions.view(-1, encoded_actions.size(-1)))
+            mcts_best_values = torch.zeros(actions.size(0)).to(device)  # Placeholder: replace with actual MCTS values
+            etv = ExpectedThoughtValueLoss()(mcts_best_values)
+            visit_counts = torch.ones(actions.size(0), policy_logits.size(-1)).to(device)  # Placeholder: replace with actual visit counts
+            exploration = ExplorationRegularization()(visit_counts)
+            old_policy = F.softmax(policy_logits.detach(), dim=-1)
+            new_policy = F.softmax(policy_logits, dim=-1)
+            kl_loss = KL_DivergenceLoss()(old_policy, new_policy)
+    
+            # Total Loss
+            loss = (
+                ppo_loss +
+                info_nce +
+                covariance +
+                dynamics_loss +
+                thought_loss +
+                pv_loss +
+                action_diversity +
+                etv +
+                exploration +
+                kl_loss
+            )
+    
+            total_loss += loss.item()
+    
+    avg_loss = total_loss / len(eval_loader)
+    print(f"World Model evaluation completed. Average loss: {avg_loss:.4f}")
+    return avg_loss
+
+
+def main():
+    args = parse_args()
+    print("Arguments parsed successfully.")
+
+    # Create save directory
+    if not os.path.exists(args.save_dir):
+        os.makedirs(args.save_dir)
+    print(f"Save directory created: {args.save_dir}")
+
+    # Load tokenizer
+    print("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    print("Tokenizer loaded successfully.")
+
+    # Define padding_idx and input dimension based on tokenizer
+    padding_idx = tokenizer.pad_token_id
+    input_dim = len(tokenizer)
+
+    # Load data
+    print("Loading and preprocessing data...")
+    train_loader, eval_loader = load_data(args, tokenizer)
+    print("Data loaded and preprocessed successfully.")
+
+    # Define model parameters
+    d_model = 512 # half to save space
+    num_heads = 8
+    num_layers = 6
+    d_ff = 2048
+    num_experts = 4
+    output_dim = input_dim
+    dropout = 0.1
+    top_k = 2
+    state_dim = 128
+    action_dim = d_model
+    hidden_dim = 512
+    vocab_dim = len(tokenizer)
+    # Initialize and load the Transformer model (on CPU)
+    print("Initializing and loading Transformer model...")
+    model_transformer = Transformer(input_dim, d_model, num_heads, num_layers, d_ff, num_experts, output_dim, dropout, top_k)
+    model_transformer.load_state_dict(torch.load(args.transformer_model_path, map_location='cpu'))
+    model_transformer.eval()
+    model_transformer.to('cpu')
+    print("Transformer model loaded and moved to CPU.")
+
+    # Define World Model components
+    representation_network = RepresentationNetwork(vocab_dim, d_model, state_dim).to(device)
+    dynamics_network = DynamicsNetwork(state_dim, action_dim, hidden_dim).to(device)
+    prediction_network = PredictionNetwork(state_dim, input_dim, 1).to(device)
+    action_encoder = ActionEncoder(input_dim, action_dim).to(device)
+
+    # Define Optimizers and Schedulers
+    optimizer = optim.AdamW(
+        list(representation_network.parameters()) +
+        list(dynamics_network.parameters()) +
+        list(prediction_network.parameters()) +
+        list(action_encoder.parameters()), 
+        lr=args.learning_rate, weight_decay=args.weight_decay
+    )
+    scheduler = CosineAnnealingLR(optimizer, T_max=args.num_epochs)
+    scaler = GradScaler()
+
+    # Initialize PPO Agent
+    ppo_agent = PPOAgent(
+        policy_network=prediction_network,
+        optimizer=optim.AdamW(prediction_network.parameters(), lr=args.learning_rate),
+        clip_epsilon=0.2,
+        entropy_coef=0.01,
+        value_coef=0.5
+    )
+
+    # Bundle World Model components
+    world_model_components = (representation_network, dynamics_network, prediction_network, action_encoder, ppo_agent)
+
+    print("Setup complete. Starting training...")
+
+    for epoch in range(args.num_epochs):
+        print(f"Epoch {epoch + 1}/{args.num_epochs} started.")
+        
+        # Train World Model
+        avg_train_loss = train_epoch_world_model(
+            world_model_components, 
+            train_loader, 
+            optimizer, 
+            scheduler, 
+            scaler, 
+            args, 
+            model_transformer, 
+            state_dim, 
+            d_model, # this is the embedding dimension
+            input_dim
+        )
+
+        print(f"World Model training epoch {epoch + 1} completed. Average loss: {avg_train_loss:.4f}")
+
+        # Evaluate World Model
+        avg_eval_loss = evaluate_world_model(
+            world_model_components, 
+            model_transformer, 
+            eval_loader, 
+            args
+        )
+        print(f"Evaluation for epoch {epoch + 1} completed. Average loss: {avg_eval_loss:.4f}")
+
+        print(f"Epoch {epoch + 1}/{args.num_epochs}, Train Loss: {avg_train_loss:.4f}, Eval Loss: {avg_eval_loss:.4f}")
+
+        # Save Models
+        save_all_models(model_transformer, representation_network, dynamics_network, prediction_network, action_encoder, args.save_dir, epoch + 1)
+        print(f"Models saved for epoch {epoch + 1}")
+
+    print("Training completed.")
+   
+
+if __name__ == '__main__':
+    main()
+
+