methods/sim_cot_rl_vae_end.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
import math
import re

from torch.distributions import Normal
from collections import defaultdict
import torch
from torch.nn.utils import clip_grad_norm_
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader
from collections import namedtuple
from transformers.models.gpt2 import GPT2LMHeadModel

from modules import grpo
from modules.projector import LatentPolicy
from modules.utils import get_position_ids_from_attention_mask
import copy

Outputs = namedtuple("Outputs", ["loss", "loss_explain_all", "inputs_embeds", "logits"])
Outputs_withkl = namedtuple("Outputs", ["loss", "loss_explain_all", "loss_kl", "inputs_embeds", "logits"])
Outputs_withmask = namedtuple("Outputs_withmask",
                              ["loss", "attention_mask", "loss_explain_all", "inputs_embeds", "logits"])
MAX_N_LATENT = 8

from dataclasses import dataclass
from typing import Optional
from modules.projector import STOPPolicy


@dataclass
class LatentForwardResult:
    inputs_embeds: torch.Tensor  # (B, L_prompt, H)
    attention_mask: torch.Tensor  # (B, L_prompt)
    latent_inputs_embeds: torch.Tensor  # (B, L_latent, H)


@dataclass
class GenerationOutputs:
    question_input_ids: torch.Tensor
    question_attention_mask: torch.Tensor
    latent_inputs_embeds: torch.Tensor
    latent_attention_mask: torch.Tensor
    pred_ids: torch.Tensor


class Coconut_RL_End_VAE(nn.Module):
    def __init__(
            self,
            base_causallm,
            expainable_llm,
            latent_token_id,
            start_latent_id,
            end_latent_id,
            eos_token_id,
            tokenizer=None,
            rl_config=None,
            is_training=True,
    ):
        super(Coconut_RL_End_VAE, self).__init__()
        self.gen_forward_cnt = 0
        self.base_causallm = base_causallm
        self.latent_token_id = latent_token_id
        self.eos_token_id = eos_token_id
        self.pad_token_id = eos_token_id
        self.start_latent_id = start_latent_id
        self.end_latent_id = end_latent_id
        self.expainable_llm = expainable_llm
        self.base_causallm.config.use_cache = True
        self.expainable_llm.config.use_cache = True
        self.end_head = STOPPolicy(
            feature_size=self.base_causallm.config.hidden_size,
            intermediate_size=self.base_causallm.config.hidden_size
        )
        self.latent_head = LatentPolicy(
            feature_size=self.base_causallm.config.hidden_size,
            intermediate_size=self.base_causallm.config.hidden_size,
            deterministic=False
        )
        self.tokenizer = tokenizer
        self.rl_config = rl_config
        self.latent_sigma = rl_config.latent_sigma
        self.rl_step_counter = 0
        self.is_training = is_training
        if self.expainable_llm is not None:
            for param in self.expainable_llm.parameters():
                param.requires_grad = False

        if isinstance(self.base_causallm, GPT2LMHeadModel):
            self.embedding = self.base_causallm.transformer.get_input_embeddings()
        else:
            self.embedding = self.base_causallm.get_input_embeddings()

        self.replay_buffer = grpo.ReplayBuffer()
        self.grpo_loss = grpo.GRPOLoss(rl_config=self.rl_config)

    @property
    def device(self):
        return next(self.parameters()).device

    def latent_forward_with_cache(self, input_ids, attention_mask, position_ids=None):
        """
        只负责：对 batch 中所有 latent token 做并行采样，
        复用 KV cache，并返回 latent embedding/logprob 等。
        """
        bs, seq_len = input_ids.size()
        device = input_ids.device
        latent_mask = (input_ids == self.latent_token_id)
        latent_positions = latent_mask.nonzero(as_tuple=False)
        latent_lists = [
            [idx[1].item() for idx in latent_positions if idx[0] == b]
            for b in range(bs)
        ]
        bol_position = latent_lists[0][0] - 1
        eol_position = latent_lists[0][-1] + 1
        max_n_latents = max((len(lst) for lst in latent_lists), default=0)

        next_compute_range = (0, seq_len)
        inputs_embeds = self.embedding(input_ids)

        if max_n_latents > 0:
            next_compute_range = (0, latent_positions[:, 1].min().item())

        # stop_logprob = []
        kv_cache = None
        latent_samples_all = [[] for _ in range(bs)]

        for pass_idx in range(max_n_latents if max_n_latents > 0 else 1):
            slice_start, slice_end = next_compute_range

            if kv_cache is None:
                outputs = self.base_causallm(
                    inputs_embeds=inputs_embeds[:, slice_start:slice_end, :],
                    attention_mask=attention_mask[:, slice_start:slice_end],
                    position_ids=(
                        position_ids[:, slice_start:slice_end]
                        if position_ids is not None else None
                    ),
                    output_hidden_states=True,
                )
                hidden_offset = 0
            else:
                past_key_values = [
                    (k[:, :, :slice_start, :], v[:, :, :slice_start, :])
                    for k, v in kv_cache
                ]
                outputs = self.base_causallm(
                    inputs_embeds=inputs_embeds[:, slice_start:slice_end, :],
                    attention_mask=attention_mask[:, :slice_end],
                    position_ids=(
                        position_ids[:, slice_start:slice_end]
                        if position_ids is not None else None
                    ),
                    past_key_values=past_key_values,
                    output_hidden_states=True,
                )
                hidden_offset = slice_start

            kv_cache = outputs.past_key_values
            hidden_states = outputs.hidden_states[-1]

            next_compute_range = (
                slice_end,
                seq_len if pass_idx + 1 >= max_n_latents else slice_end + 1,
            )

            tensor_list = [
                [inputs_embeds[b, pos, :] for pos in range(seq_len)]
                for b in range(bs)
            ]

            for b, latent_positions in enumerate(latent_lists):
                if pass_idx >= len(latent_positions):
                    continue
                token_idx = latent_positions[pass_idx]
                vec = hidden_states[b, token_idx - 1 - hidden_offset, :]

                mu, logvar = self.latent_head.forward(vec)  # [dim*2]
                # print(mu, logvar)
                std = torch.exp(0.5 * logvar)
                eps = torch.randn_like(std)
                z = mu + std * eps
                # print(self.is_training)
                # dist = torch.distributions.Normal(mu, self.latent_sigma)
                # logprob = dist.log_prob(sample).sum()

                tensor_list[b][token_idx] = z
                latent_samples_all[b].append(z)

            inputs_embeds = torch.stack([torch.stack(tensor_list[b]) for b in range(bs)], dim=0)

            pred_token_id = self.end_head(inputs_embeds[:, next_compute_range[0], :])
            # print(torch.softmax(pred_token_id, dim=-1))
            if self.is_training == True:
                is_end = torch.multinomial(torch.softmax(pred_token_id, dim=-1), num_samples=1).squeeze(
                    -1).bool()  # [batch_size]
            else:
                is_end = (pred_token_id[:, 1] > pred_token_id[:, 0])
            # stop_logprob.append(torch.log_softmax(pred_token_id, dim=-1))
            if pass_idx % self.rl_config.c_thought == 0:
                attention_mask[is_end, bol_position + 1 + pass_idx:eol_position] = 0
            if pass_idx >= max_n_latents:  # or (attention_mask[:, bol_position + 1 + pass_idx] == 0).all():
                break

        self.gen_forward_cnt += max_n_latents + 1

        latent_dim = inputs_embeds.size(-1)
        latent_lengths = torch.tensor([len(lst) for lst in latent_lists],
                                      device=device, dtype=torch.long)
        max_latent_len = latent_lengths.max().item() if latent_lengths.numel() > 0 else 0

        if max_latent_len > 0:
            latent_inputs_embeds = torch.zeros(bs, max_latent_len, latent_dim, device=device)
            for b in range(bs):
                if latent_samples_all[b]:
                    latent_inputs_embeds[b, :len(latent_samples_all[b])] = torch.stack(latent_samples_all[b])
        else:
            latent_inputs_embeds = torch.zeros(bs, 0, latent_dim, device=device)

        return LatentForwardResult(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            latent_inputs_embeds=latent_inputs_embeds,
        )

    @torch.no_grad()
    def generate_with_latent(
            self,
            questions,
            questions_mask,
            max_new_tokens=16,
            temperature=1.0,
            synced_gpus=False,
    ):
        self.eval()
        device = questions.device
        batch_size = questions.size(0)

        latent_mask = (questions == self.latent_token_id)
        latent_positions = latent_mask.nonzero(as_tuple=False)
        latent_lists = [
            [idx[1].item() for idx in latent_positions if idx[0] == b]
            for b in range(batch_size)
        ]
        bol_position = latent_lists[0][0] - 1
        eol_position = latent_lists[0][-1] + 1
        questions = questions[:, :eol_position + 1]
        questions_mask = questions_mask[:, :eol_position + 1]
        position_ids = get_position_ids_from_attention_mask(questions_mask)

        latent_out = self.latent_forward_with_cache(
            input_ids=questions,
            attention_mask=questions_mask,
            position_ids=position_ids,
        )
        attention_mask = latent_out.attention_mask
        prompt_embeds = latent_out.inputs_embeds
        latent_inputs_embeds = latent_out.latent_inputs_embeds
        latent_attention_mask = attention_mask[:, bol_position + 1:eol_position]

        # 直接用 prompt_embeds 做一次全量前向，拿到最后 logits
        position_ids = get_position_ids_from_attention_mask(attention_mask)
        outputs = self.base_causallm(
            inputs_embeds=prompt_embeds,
            attention_mask=attention_mask,
            position_ids=position_ids,
            use_cache=False,
        )
        logits = outputs.logits[:, -1, :]

        tokens = [[questions[i, -1].item()] for i in range(batch_size)]
        unfinished = torch.ones(batch_size, dtype=torch.bool, device=device)
        cur_inputs_embeds = prompt_embeds
        cur_attention_mask = attention_mask

        for _ in range(max_new_tokens):
            if temperature != 1.0:
                logits = logits / temperature
            if self.is_training:
                probs = torch.softmax(logits, dim=-1)
                next_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1)
            else:
                next_tokens = torch.argmax(logits, dim=-1)

            next_tokens = torch.where(
                unfinished, next_tokens, torch.full_like(next_tokens, self.eos_token_id)
            )
            for i, tok in enumerate(next_tokens.tolist()):
                tokens[i].append(tok)

            next_token_embeds = self.embedding(next_tokens).unsqueeze(1)
            cur_inputs_embeds = torch.cat([cur_inputs_embeds, next_token_embeds], dim=1)
            cur_attention_mask = torch.cat(
                [cur_attention_mask, torch.ones(batch_size, 1, dtype=cur_attention_mask.dtype, device=device)],
                dim=1,
            )

            position_ids = get_position_ids_from_attention_mask(cur_attention_mask)
            outputs = self.base_causallm(
                inputs_embeds=cur_inputs_embeds,
                attention_mask=cur_attention_mask,
                position_ids=position_ids,
                use_cache=False,
            )
            self.gen_forward_cnt += 1
            logits = outputs.logits[:, -1, :]

            unfinished &= next_tokens.ne(self.eos_token_id)
            if not unfinished.any():
                break

        if synced_gpus:
            while self.gen_forward_cnt < max_new_tokens + MAX_N_LATENT:
                self.gen_forward_cnt += 1
                _ = self.base_causallm(
                    inputs_embeds=cur_inputs_embeds[:, cur_pos - 1:cur_pos, :],
                    attention_mask=cur_attention_mask[:, :cur_pos],
                    past_key_values=cur_kv_cache,
                    use_cache=True,
                )

        max_len = max(len(seq) for seq in tokens)
        pred_ids = torch.full(
            (batch_size, max_len),
            fill_value=self.pad_token_id,
            dtype=torch.long,
            device=device,
        )
        for i, seq in enumerate(tokens):
            pred_ids[i, : len(seq)] = torch.tensor(seq, device=device)

        return GenerationOutputs(
            question_input_ids=questions[:, :bol_position + 1],
            question_attention_mask=questions_mask[:, :bol_position + 1],
            latent_inputs_embeds=latent_inputs_embeds,
            latent_attention_mask=latent_attention_mask,
            pred_ids=pred_ids,
        )

    def rl_step(self, batch, optimizer, writer=None, global_step=None):
        rl_config = self.rl_config
        device = self.device
        questions = batch["input_ids"]
        questions_mask = batch["attention_mask"]
        answers = batch["labels"]
        explainable_ids_list = batch["explainable_ids_list"]

        self.replay_buffer.clear()

        experience = self.rollout(
            questions=questions,
            questions_mask=questions_mask,
            group_size=rl_config.group_size,
            gt_answers=answers,
            explainable_ids_list=explainable_ids_list,
        )
        self.replay_buffer.append(experience.to("cpu"))

        log_dict = {
            "train/rewards": experience.rewards.mean().item(),
            "train/accuracies": experience.accuracies.mean().item(),
            "train/n_latent_forward": experience.n_latent_forward.float().mean().item(),
        }
        if getattr(experience, "explain_logprob", None) is not None:
            log_dict["train/explain_logprob"] = experience.explain_logprob.mean().item()

        torch.cuda.empty_cache()

        dataloader = DataLoader(
            dataset=self.replay_buffer,
            batch_size=rl_config.exp_batch_size,
            shuffle=True,
            collate_fn=grpo.join_experience_batch,
        )

        gradient_accumulation_steps = max(1, rl_config.gradient_accumulation_steps)
        total_inner_steps = len(dataloader)

        loop_sums = defaultdict(float)
        inner_steps = 0
        accum_counter = 0
        optimizer.zero_grad()
        last_grad_norm = 0.0

        for exp in dataloader:
            inner_steps += 1
            accum_counter += 1
            exp: grpo.Experience = exp.to(device)

            latent_logprobs, answer_logprobs = self.get_logprobs(e=exp)
            loss_dict = self.grpo_loss(
                latent_logprobs=latent_logprobs,
                answer_logprobs=answer_logprobs,
                experience=exp,
            )

            raw_total_loss = loss_dict["total_loss"]
            (raw_total_loss / gradient_accumulation_steps).backward()

            for k, v in loss_dict.items():
                value = v.item() if torch.is_tensor(v) else float(v)
                loop_sums[f"train/{k}"] += value

            need_update = (accum_counter % gradient_accumulation_steps == 0) or (inner_steps == total_inner_steps)
            if need_update:
                grad_norm = clip_grad_norm_(self.parameters(), max_norm=1.0)
                optimizer.step()
                optimizer.zero_grad()
                last_grad_norm = grad_norm.item() if torch.is_tensor(grad_norm) else float(grad_norm)

        if inner_steps > 0:
            for key in list(loop_sums.keys()):
                log_dict[key] = loop_sums[key] / inner_steps
            log_dict["train/grad_norm"] = last_grad_norm

        return log_dict

    def extract_equal_result(self, expression):
        match = re.search(r'=(.*?)>>', expression)
        if match:
            return match.group(1).strip()
        else:
            return ''

    @torch.no_grad()
    def compute_explain_logprob(
            self,
            latent_inputs_embeds: torch.Tensor,  # (B, L_latent, H)
            n_latent_forward: torch.Tensor,  # (B,)
            pred_strings,
            max_gen_len=16
    ):
        if self.expainable_llm is None:
            return None

        device = self.expainable_llm.device
        eos_token_id = self.eos_token_id
        thought = self.rl_config.c_thought
        max_stage = self.rl_config.max_latent_stage
        tok = self.tokenizer  # 或 self.expainable_llm 的 tokenizer

        B, _, hidden = latent_inputs_embeds.size()

        # 1) 收集所有可用的 latent prefix
        prefix_list, owner = [], []
        for b in range(B):
            latent_valid = int(n_latent_forward[b].item())
            stage_num = min(max_stage, latent_valid // thought)
            for s in range(stage_num):
                start = s * thought
                end = start + thought
                prefix = latent_inputs_embeds[b, start:end, :].to(device)
                prefix_list.append(prefix)
                owner.append(b)

        N = len(prefix_list)
        if N == 0:
            return torch.zeros(B, device=device), torch.ones(B, device=device)

        seq_embeds = torch.stack(prefix_list, dim=0).clone()  # (N, thought, H)
        attention_mask = torch.ones(N, thought, dtype=torch.long, device=device)
        position_ids = torch.arange(thought, dtype=torch.long, device=device).unsqueeze(0).expand(N, -1)

        total_logprob = torch.zeros(N, device=device)
        total_tokens = torch.zeros(N, device=device)

        gen_token_ids = [[] for _ in range(N)]
        active_idx = torch.arange(N, device=device)
        # embed_layer = self.expainable_llm.get_input_embeddings()

        # 2) batched greedy decoding
        for _ in range(max_gen_len):
            outputs = self.expainable_llm(
                inputs_embeds=seq_embeds,
                attention_mask=attention_mask,
                position_ids=position_ids,
            )
            logits = outputs.logits[:, -1, :]
            log_probs = torch.log_softmax(logits, dim=-1)

            next_token = torch.argmax(logits, dim=-1)
            step_logprob = log_probs.gather(-1, next_token.unsqueeze(-1)).squeeze(-1)

            total_logprob[active_idx] += step_logprob
            total_tokens[active_idx] += 1

            for aidx, tok_id in zip(active_idx.tolist(), next_token.tolist()):
                gen_token_ids[aidx].append(tok_id)

            eos_mask = next_token.eq(eos_token_id)
            keep_mask = ~eos_mask
            if keep_mask.any():
                kept_idx = active_idx[keep_mask]
                next_embed = self.embedding(next_token[keep_mask]).unsqueeze(1)
                seq_embeds = torch.cat([seq_embeds[keep_mask], next_embed], dim=1)

                attention_mask = torch.cat(
                    [attention_mask[keep_mask],
                     torch.ones(len(kept_idx), 1, dtype=torch.long, device=device)],
                    dim=1
                )
                next_pos = position_ids[keep_mask][:, -1:] + 1
                position_ids = torch.cat([position_ids[keep_mask], next_pos], dim=1)

                active_idx = kept_idx
            else:
                break

        # 3) 汇总到 batch 维度
        explain_logprob_b = torch.zeros(B, device=device)
        explain_coherence_b = torch.ones(B, device=device)
        explain_count_b = torch.zeros(B, dtype=torch.long, device=device)
        owner_t = torch.tensor(owner, dtype=torch.long, device=device)

        valid_mask = total_tokens > 0
        avg_logprob = torch.zeros_like(total_logprob)
        avg_logprob[valid_mask] = total_logprob[valid_mask]

        for idx in range(N):
            b = owner_t[idx].item()
            explain_logprob_b[b] += avg_logprob[idx]
            explain_count_b[b] += 1

        explain_logprob_b = torch.where(
            explain_count_b > 0,
            explain_logprob_b / explain_count_b.clamp_min(1),
            torch.zeros_like(explain_logprob_b)
        )

        # 4) 解码并打印
        preview_by_batch = [[] for _ in range(B)]
        result_by_batch = [[] for _ in range(B)]
        for idx, tlist in enumerate(gen_token_ids):
            txt = tok.decode(tlist, skip_special_tokens=True) if tlist else ""
            b = owner_t[idx].item()
            result_of_eq = self.extract_equal_result(txt)
            preview_by_batch[b].append(txt)
            result_by_batch[b].append(result_of_eq)

        for b in range(B):
            total_steps = len(result_by_batch[b])  # eg. ["4", "3", "3"]
            if total_steps == 0:
                explain_coherence_b[b] = 1.0
                continue

            matched = 0
            for i in range(len(result_by_batch[b]) - 1):
                for j in range(i + 1, len(preview_by_batch[b])):
                    if result_by_batch[b][i] in preview_by_batch[b][j]:
                        matched += 1
                        break
            last_result = result_by_batch[b][-1]
            if last_result in pred_strings[b]:
                matched += 1
            explain_coherence_b[b] = matched / total_steps
        # for b, seqs in enumerate(preview_by_batch):
        #     print(f"[sample {b}]", end=' ')
        #     for j, s in enumerate(seqs):
        #         print(f"step {j}: {s}", end=' ')
        #     print()
        #
        # for b, seqs in enumerate(result_by_batch):
        #     print(f"[sample {b}]", end=' ')
        #     for j, s in enumerate(seqs):
        #         print(f"step {j}: {s}", end=' ')
        #     print()
        # 只返回 logprob
        return explain_logprob_b, explain_coherence_b

    @torch.no_grad()
    def rollout(self, questions, questions_mask, gt_answers, group_size, explainable_ids_list=None):
        rl_cfg = self.rl_config
        group_size = group_size

        # 扩展 group
        group_questions = questions.repeat_interleave(group_size, dim=0)
        group_masks = questions_mask.repeat_interleave(group_size, dim=0)
        if explainable_ids_list is not None:
            group_explain_ids = [copy.deepcopy(ids) for ids in explainable_ids_list for _ in range(group_size)]
        else:
            group_explain_ids = None

        gen_out = self.generate_with_latent(
            questions=group_questions.to(self.device),
            questions_mask=group_masks.to(self.device),
            max_new_tokens=rl_cfg.max_new_tokens,
            temperature=rl_cfg.temperature,
        )
        # print(gen_out.pred_ids.tolist())
        pred_strings = self.tokenizer.batch_decode(gen_out.pred_ids, skip_special_tokens=True)
        n_latent_forward = gen_out.latent_attention_mask.sum(dim=1)

        explain_logprob, explain_coherence = None, None
        if self.expainable_llm is not None and group_explain_ids and rl_cfg.reason_reward_weight != 0:
            explain_logprob, explain_coherence = self.compute_explain_logprob(
                latent_inputs_embeds=gen_out.latent_inputs_embeds,
                n_latent_forward=n_latent_forward,
                pred_strings=pred_strings,
            )
            # print(explain_logprob.size())

        rewards_list, acc_list, adv_list = [], [], []
        for i in range(len(questions)):
            start = i * group_size
            end = (i + 1) * group_size
            rewards, accuracies = self.get_group_rewards_and_acc(
                pred_answers=pred_strings[start:end],
                gt_answer=gt_answers[i],
                explain_coherence=None if explain_coherence is None else explain_coherence[start:end],
                n_latent_forward=n_latent_forward[start:end],
                explain_logprob_chunk=None if explain_logprob is None else explain_logprob[start:end],
            )
            rewards_list.append(rewards)
            acc_list.append(accuracies)
            adv_list.append(grpo.group_advantages(rewards))

        rewards = torch.cat(rewards_list, dim=0)
        accuracies = torch.cat(acc_list, dim=0)
        advantages = torch.cat(adv_list, dim=0)
        # print(gen_out.latent_inputs_embeds.size())
        # print(gen_out.latent_attention_mask.size())
        # print(gen_out.question_input_ids.size())
        # print(gen_out.question_attention_mask.size())
        experience = grpo.Experience(
            question_input_ids=gen_out.question_input_ids,
            question_attention_mask=gen_out.question_attention_mask,
            latent_inputs_embeds=gen_out.latent_inputs_embeds,
            latent_attention_mask=gen_out.latent_attention_mask,
            answer_input_ids=gen_out.pred_ids,
            answer_attention_mask=gen_out.pred_ids.ne(self.pad_token_id).long(),
            n_latent_forward=n_latent_forward.unsqueeze(1),
            rewards=rewards,
            accuracies=accuracies,
            advantages=advantages,
        )

        if explain_logprob is not None:
            experience.explain_logprob = explain_logprob

        latent_logprobs, answer_logprobs = self.get_logprobs(experience)
        # print(latent_logprobs.size())
        # print(advantages.size())
        experience.latent_logprobs = latent_logprobs
        experience.answer_logprobs = answer_logprobs
        return experience

    def get_logprobs(self, e: grpo.Experience):
        question_length = e.question_input_ids.shape[1]
        latent_length = e.latent_inputs_embeds.shape[1]
        answer_length = e.answer_input_ids.shape[1]

        question_inputs_embeds = self.embedding(e.question_input_ids)
        answer_inputs_embeds = self.embedding(e.answer_input_ids)

        all_inputs_embeds = torch.cat(
            [question_inputs_embeds, e.latent_inputs_embeds, answer_inputs_embeds], dim=1
        )
        all_attention_mask = torch.cat(
            [e.question_attention_mask, e.latent_attention_mask, e.answer_attention_mask], dim=1
        )
        all_position_ids = get_position_ids_from_attention_mask(all_attention_mask)

        all_outputs = self.base_causallm.forward(
            inputs_embeds=all_inputs_embeds,
            attention_mask=all_attention_mask,
            position_ids=all_position_ids,
            output_hidden_states=True,
        )

        # ===== Latent log_prob Gaussian =====
        hidden = all_outputs.hidden_states[-1]
        vec = hidden[:, question_length - 1: question_length + latent_length - 1, :]

        mu, logvar = self.latent_head.forward(vec)  # [dim*2]
        std = torch.exp(0.5 * logvar)
        dist = torch.distributions.Normal(mu, std)
        latent_samples = e.latent_inputs_embeds[:, :latent_length, :]
        pred_token_id = self.end_head(latent_samples)
        # print(torch.softmax(pred_token_id, dim=-1))
        for pass_idx in range(e.latent_attention_mask.size(1)):
            if pass_idx % self.rl_config.c_thought:
                pred_token_id[:, pass_idx, 1] = float('-inf')
        pred_logprob = torch.log_softmax(pred_token_id, dim=-1)
        eos_place = e.latent_attention_mask.sum(-1)[:, None]
        # print(eos_place)
        log_prob_fill = torch.cat((pred_logprob[:, :, 1], torch.zeros_like(pred_logprob[:, 0:1, 1])), dim=1)
        latent_logprobs = dist.log_prob(latent_samples).sum(dim=-1) + pred_logprob[:, :, 0]
        # print(latent_logprobs.tolist())
        logits = all_outputs.logits
        answer_logits = logits[:, question_length + latent_length - 1: -1, :]
        answer_logprobs = F.log_softmax(answer_logits, dim=-1)
        answer_logprobs = answer_logprobs.gather(dim=-1, index=e.answer_input_ids.unsqueeze(-1)).squeeze(-1)
        answer_logprobs[:, 0] = log_prob_fill.gather(dim=1, index=eos_place).squeeze(-1)
        # print(answer_logprobs.tolist())

        return latent_logprobs, answer_logprobs

    def get_group_rewards_and_acc(
            self,
            pred_answers,
            gt_answer,
            explain_coherence,
            n_latent_forward,
            explain_logprob_chunk=None,
    ):
        rl_config = self.rl_config
        group_size = len(pred_answers)
        format = torch.zeros(group_size, 1, device=self.device)
        accuracies = torch.zeros(group_size, 1, device=self.device)
        for i, ans in enumerate(pred_answers):
            gt_answer[gt_answer == -100] = self.pad_token_id
            gt_string = self.tokenizer.decode(gt_answer, skip_special_tokens=True)
            marker = "<|end-latent|>###"
            gt_answer_string = gt_string[gt_string.rfind("###") + len("###"):]
            idx = ans.rfind(marker)
            if idx == -1:
                pred_string = ""
            else:
                pred_string = ans[idx + len(marker):]
                if len(pred_string.strip()) != 0:
                    format[i] = 1
            accuracies[i] = self.verify_answer(gt_answer=gt_answer_string, pred_answer=pred_string)
            print(
                f"Ground_truth: {gt_answer_string} Answer {ans} Accuracy: {accuracies[i]} Format: {format[i]}")
            if explain_logprob_chunk is not None:
                print(f"Explain_prob: {explain_logprob_chunk[i]}")
            if explain_coherence is not None:
                print(f"Explain_coherence: {explain_coherence[i]}")

        rewards = rl_config.accuracy_reward_weight * accuracies.detach().clone() + rl_config.format_reward_weight * format.detach().clone()

        # if explain_logprob_chunk is not None:
        #     rewards = rewards + rl_config.reason_reward_weight * explain_logprob_chunk[:, None]
        if explain_coherence is not None:
            if rl_config.accuracy_reward_weight ==0:
                rewards = rewards + rl_config.reason_reward_weight * explain_coherence[:, None]
            else:
                rewards = rewards + rl_config.reason_reward_weight * explain_coherence[:, None] * (accuracies.detach().clone())

        return rewards, accuracies

    def verify_answer(self, gt_answer: str, pred_answer: str) -> float:
        # gt_answer = gt_answer.strip("\n ").rstrip(".").replace(",", "")
        # pred_answer = pred_answer.strip("\n ").rstrip(".").rstrip(".\n").replace(",", "")
        try:  # some answers may be like '10.0' but predicted as '10'
            gt_answer = float(gt_answer)
            pred_answer = float(pred_answer)
            return float(gt_answer == pred_answer)
        except ValueError:
            return float(self.is_equiv(gt_answer, pred_answer))

    @torch.no_grad()
    def eval_generation(self, dataloader, val='val'):
        loop_sums = defaultdict(float)
        inner_steps = 0
        self.is_training = False
        for step, batch in enumerate(dataloader):
            inner_steps += 1
            questions = batch["input_ids"]
            questions_mask = batch["attention_mask"]
            answers = batch["labels"]
            explainable_ids_list = batch["explainable_ids_list"]

            experience = self.rollout(
                questions=questions,
                questions_mask=questions_mask,
                gt_answers=answers,
                group_size=1,
                explainable_ids_list=explainable_ids_list,
            )

            loop_sums[f"{val}/rewards"] += experience.rewards.mean().item()
            loop_sums[f"{val}/accuracies"] += experience.accuracies.mean().item()
            loop_sums[f"{val}/n_latent_forward"] += experience.n_latent_forward.float().mean().item()
            if getattr(experience, "explain_logprob", None) is not None:
                loop_sums[f"{val}/explain_logprob"] += experience.explain_logprob.mean().item()

            torch.cuda.empty_cache()
        print(f"Accuracy in Validation {loop_sums[f'{val}/accuracies'] * 100 / inner_steps}")
        log_dict = {}
        if inner_steps > 0:
            for key in loop_sums:
                log_dict[key] = loop_sums[key] / inner_steps

        self.is_training = True
        return log_dict

        # -- evaluation ends --#

    def is_equiv(self, str1, str2, verbose=False):
        if str1 is None and str2 is None:
            print("WARNING: Both None")
            return True
        if str1 is None or str2 is None:
            return False

        try:
            ss1 = self.strip_string(str1)
            ss2 = self.strip_string(str2)
            if verbose:
                print(ss1, ss2)
            return ss1 == ss2
        except Exception:
            return str1 == str2

    def remove_boxed(self, s):
        if "\\boxed " in s:
            left = "\\boxed "
            assert s[:len(left)] == left
            return s[len(left):]

        left = "\\boxed{"

        assert s[:len(left)] == left
        assert s[-1] == "}"

        return s[len(left):-1]

    def last_boxed_only_string(self, string):
        idx = string.rfind("\\boxed")
        if "\\boxed " in string:
            return "\\boxed " + string.split("\\boxed ")[-1].split("$")[0]
        if idx < 0:
            idx = string.rfind("\\fbox")
            if idx < 0:
                return None

        i = idx
        right_brace_idx = None
        num_left_braces_open = 0
        while i < len(string):
            if string[i] == "{":
                num_left_braces_open += 1
            if string[i] == "}":
                num_left_braces_open -= 1
                if num_left_braces_open == 0:
                    right_brace_idx = i
                    break
            i += 1

        if right_brace_idx is None:
            retval = None
        else:
            retval = string[idx:right_brace_idx + 1]

        return retval

    def fix_fracs(self, string):
        substrs = string.split("\\frac")
        new_str = substrs[0]
        if len(substrs) > 1:
            substrs = substrs[1:]
            for substr in substrs:
                new_str += "\\frac"
                if substr[0] == "{":
                    new_str += substr
                else:
                    try:
                        assert len(substr) >= 2
                    except AssertionError:
                        return string
                    a = substr[0]
                    b = substr[1]
                    if b != "{":
                        if len(substr) > 2:
                            post_substr = substr[2:]
                            new_str += "{" + a + "}{" + b + "}" + post_substr
                        else:
                            new_str += "{" + a + "}{" + b + "}"
                    else:
                        if len(substr) > 2:
                            post_substr = substr[2:]
                            new_str += "{" + a + "}" + b + post_substr
                        else:
                            new_str += "{" + a + "}" + b
        string = new_str
        return string

    def fix_a_slash_b(self, string):
        if len(string.split("/")) != 2:
            return string
        a = string.split("/")[0]
        b = string.split("/")[1]
        try:
            a = int(a)
            b = int(b)
            assert string == "{}/{}".format(a, b)
            new_string = "\\frac{" + str(a) + "}{" + str(b) + "}"
            return new_string
        except AssertionError:
            return string

    def remove_right_units(self, string):
        # "\\text{ " only ever occurs (at least in the val set) when describing units
        if "\\text{ " in string:
            splits = string.split("\\text{ ")
            assert len(splits) == 2
            return splits[0]
        else:
            return string

    def fix_sqrt(self, string):
        if "\\sqrt" not in string:
            return string
        splits = string.split("\\sqrt")
        new_string = splits[0]
        for split in splits[1:]:
            if split[0] != "{":
                a = split[0]
                new_substr = "\\sqrt{" + a + "}" + split[1:]
            else:
                new_substr = "\\sqrt" + split
            new_string += new_substr
        return new_string

    def strip_string(self, string):
        # linebreaks
        string = string.replace("\n", "")

        # remove inverse spaces
        string = string.replace("\\!", "")

        # replace \\ with \
        string = string.replace("\\\\", "\\")

        # replace tfrac and dfrac with frac
        string = string.replace("tfrac", "frac")
        string = string.replace("dfrac", "frac")

        # remove \left and \right
        string = string.replace("\\left", "")
        string = string.replace("\\right", "")

        # Remove circ (degrees)
        string = string.replace("^{\\circ}", "")
        string = string.replace("^\\circ", "")

        # remove dollar signs
        string = string.replace("\\$", "")

        # remove units (on the right)
        string = self.remove_right_units(string)

        # remove percentage
        string = string.replace("\\%", "")
        string = string.replace("\%", "")  # noqa: W605

        # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string
        string = string.replace(" .", " 0.")
        string = string.replace("{.", "{0.")
        # if empty, return empty string
        if len(string) == 0:
            return string
        if string[0] == ".":
            string = "0" + string

        # to consider: get rid of e.g. "k = " or "q = " at beginning
        if len(string.split("=")) == 2:
            if len(string.split("=")[0]) <= 2:
                string = string.split("=")[1]

        # fix sqrt3 --> sqrt{3}
        string = self.fix_sqrt(string)

        # remove spaces
        string = string.replace(" ", "")

        # \frac1b or \frac12 --> \frac{1}{b} and \frac{1}{2}, etc. Even works with \frac1{72} (but not \frac{72}1). Also does a/b --> \\frac{a}{b}
        string = self.fix_fracs(string)

        # manually change 0.5 --> \frac{1}{2}
        if string == "0.5":
            string = "\\frac{1}{2}"

        # NOTE: X/Y changed to \frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y
        string = self.fix_a_slash_b(string)

        return string