Deleting directories, moving files into root

util/common_settings.py

@@ -1,18 +0,0 @@
-
-NUM_INFERENCE_STEPS = 30
-GUIDANCE_SCALE = 7.5
-
-MARIO_HEIGHT = 16
-MARIO_WIDTH = 16
-
-MARIO_TILE_PIXEL_DIM = 16
-MARIO_TILE_COUNT = 13
-
-LR_HEIGHT = 32
-LR_WIDTH = 32
-
-LR_TILE_PIXEL_DIM = 8
-LR_TILE_COUNT = 8
-
-MEGAMAN_HEIGHT = 14
-MEGAMAN_WIDTH = 16

util/naming_conventions.py

@@ -1,29 +0,0 @@
-model_name_map = [
-    ("Mar1and2-conditional-regular", "MLM-regular"),
-    ("Mar1and2-conditional-absence", "MLM-absence"),
-    ("Mar1and2-conditional-negative", "MLM-negative"),
-    ("Mar1and2-conditional-MiniLM-regular", "MiniLM-single-regular"),
-    ("Mar1and2-conditional-MiniLM-absence", "MiniLM-single-absence"),
-    ("Mar1and2-conditional-MiniLM-negative", "MiniLM-single-negative"),
-    ("Mar1and2-conditional-MiniLMsplit-regular", "MiniLM-multiple-regular"),
-    ("Mar1and2-conditional-MiniLMsplit-absence", "MiniLM-multiple-absence"),
-    ("Mar1and2-conditional-MiniLMsplit-negative", "MiniLM-multiple-negative"),
-    ("Mar1and2-conditional-GTE-regular", "GTE-single-regular"),
-    ("Mar1and2-conditional-GTE-absence", "GTE-single-absence"),
-    ("Mar1and2-conditional-GTE-negative", "GTE-single-negative"),
-    ("Mar1and2-conditional-GTEsplit-regular", "GTE-multiple-regular"),
-    ("Mar1and2-conditional-GTEsplit-absence", "GTE-multiple-absence"),
-    ("Mar1and2-conditional-GTEsplit-negative", "GTE-multiple-negative"),
-    ("Mar1and2-fdm-MiniLM-regular", "FDM-MiniLM-regular"),
-    ("Mar1and2-fdm-MiniLM-absence", "FDM-MiniLM-absence"),
-    ("Mar1and2-fdm-GTE-regular", "FDM-GTE-regular"),
-    ("Mar1and2-fdm-GTE-absence", "FDM-GTE-absence"),
-    ("Mar1and2-wgan", "WGAN"),
-    ("Mar1and2-unconditional", "Unconditional"),
-    ("MarioGPT_metrics", "MarioGPT"),
-]
-
-def get_model_name_map_and_order():
-    mapping = dict(model_name_map)
-    order = [v for k, v in model_name_map]
-    return mapping, order

util/plotter.py

@@ -1,173 +0,0 @@
-# Track changes in loss and learning rate during execution
-import argparse
-import matplotlib
-import matplotlib.pyplot as plt
-import os
-import time
-import json
-import tempfile
-import shutil
-from pathlib import Path
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Train a text-conditional diffusion model for tile-based level generation")
-    
-    # Dataset args
-    parser.add_argument("--log_file", type=str, default=None, help="The the filepath of the file to get the data from")
-    parser.add_argument("--left_key", type=str, default=None, help="The key for the left y-axis")
-    parser.add_argument("--right_key", type=str, default=None, help="The key for the right y-axis")
-    parser.add_argument("--left_label", type=str, default=None, help="The label for the left y-axis")
-    parser.add_argument("--right_label", type=str, default=None, help="The label for the right y-axis")
-    parser.add_argument("--output_png", type=str, default="output.png", help="The output png file")
-    parser.add_argument("--update_interval", type=int, default=1.0, help="The update inteval in epochs")
-    parser.add_argument("--start_point", type=int, default=None, help="The start point for the plot")
-
-    return parser.parse_args()
-
-
-def main():
-    args = parse_args()
-
-    log_file = args.log_file
-    left_key = args.left_key
-    right_key = args.right_key
-    left_label = args.left_label
-    right_label = args.right_label
-    output_png = args.output_png
-    update_interval = args.update_interval
-    start_point = args.start_point
-
-    general_update_plot(log_file, left_key, right_key, left_label, right_label, output_png, update_interval=update_interval, startPoint=start_point)
-
-
-def general_update_plot(log_file, left_key, right_key, left_label, right_label, output_png, update_interval=1.0, startPoint=None):
-    log_dir = os.path.dirname(log_file)
-    
-    # Create figure here and ensure it's closed
-    fig = plt.figure(figsize=(10, 6))
-    ax = fig.add_subplot(111)
-    
-    try:
-        if os.path.exists(log_file):
-            with open(log_file, 'r') as f:
-                data = [json.loads(line) for line in f if line.strip()]
-            
-            if not data:
-                return
-            
-            if startPoint is not None:
-                data = [entry for entry in data if entry.get('epoch', 0) >= startPoint]
-            
-            if not data:
-                return
-
-            epochs = [entry.get('epoch', 0) for entry in data]
-            left = [entry.get(left_key, 0) for entry in data]
-
-            # For right axis (e.g., lr), only include points where right_key exists
-            right_points = [(entry.get('epoch', 0), entry.get(right_key))
-                            for entry in data if right_key in entry]
-            if right_points:
-                right_epochs, right_values = zip(*right_points)
-            else:
-                right_epochs, right_values = [], []
-
-            # Clear axis
-            ax.clear()
-            
-            # Plot both metrics on the same axis
-            ax.plot(epochs, left, 'b-', label=left_label)
-            if right_epochs:
-                ax.plot(right_epochs, right_values, 'r-', label=right_label)
-            
-            ax.set_xlabel('Epoch')
-            ax.set_ylabel(left_label) # "Loss" as y-axis label
-            ax.set_title('Training Progress')
-            ax.legend(loc='upper left')
-            #Limit x-axis to startPoint if provided
-            if startPoint is not None:
-                ax.set_xlim(left=startPoint)
-            fig.tight_layout()
-
-            # Use the stored base directory instead of getting it from log_file
-            if os.path.isabs(output_png) or os.path.dirname(output_png):
-                output_path = output_png
-            else:
-                output_path = os.path.join(log_dir, output_png)
-
-            save_figure_safely(fig, output_path)
-    finally:
-        plt.close(fig)  # Ensure figure is closed even if an error occurs
-
-def save_figure_safely(fig, output_path):
-    """Save figure to a temporary file first, then move it to the final location"""
-    output_path = str(Path(output_path))  # Convert to string path
-    
-    # Create temporary file with .png extension
-    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
-        tmp_path = tmp_file.name
-    
-    try:
-        # Save to temporary file
-        fig.savefig(tmp_path)
-        
-        # Create output directory if it doesn't exist
-        os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
-        
-        # Try to move the file to final destination
-        # If move fails, try to copy and then delete
-        try:
-            shutil.move(tmp_path, output_path)
-        except OSError:
-            shutil.copy2(tmp_path, output_path)
-            os.unlink(tmp_path)
-    except Exception as e:
-        # Clean up temporary file if anything goes wrong
-        if os.path.exists(tmp_path):
-            os.unlink(tmp_path)
-        raise e
-
-class Plotter:
-    def __init__(self, log_file, update_interval=1.0, left_key='loss', right_key='lr', 
-                 left_label='Loss', right_label='Learning Rate', output_png='training_progress.png'):
-        self.log_dir = os.path.dirname(log_file)
-        self.log_file = log_file
-        self.update_interval = update_interval
-        self.running = True
-        self.output_png = output_png
-        self.left_key = left_key
-        self.right_key = right_key
-        self.left_label = left_label
-        self.right_label = right_label
-        
-        matplotlib.use('Agg')
-        
-    def __enter__(self):
-        return self
-        
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.stop_plotting()
-        
-    def __del__(self):
-        self.stop_plotting()
-
-    def update_plot(self):
-        general_update_plot(self.log_file, self.left_key, self.right_key, 
-                          self.left_label, self.right_label, self.output_png, 
-                          update_interval=self.update_interval)
-    
-    def start_plotting(self):
-        print("Starting plotting in background")
-        while self.running:
-            self.update_plot()
-            time.sleep(self.update_interval)
-    
-    def stop_plotting(self):
-        if hasattr(self, 'running'):  # Check if already stopped
-            self.running = False
-            self.update_plot()
-            print("Plotting stopped")
-
-if __name__ == "__main__":
-    main()

util/sampler.py

@@ -1,473 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-
-import os
-import subprocess
-import tempfile
-
-import numpy as np
-import torch
-from PIL.Image import Image
-from tqdm import tqdm
-from transformers import LogitsProcessorList, TemperatureLogitsWarper, TopKLogitsWarper
-
-
-from mario_gpt.lm.base import BaseMarioLM
-from mario_gpt.prompter import Prompter
-from mario_gpt.simulator import Simulator
-from mario_gpt.utils import (
-    convert_level_to_png,
-    load_level,
-    save_level,
-    trim_level,
-    view_level,
-)
-
-def scene_to_ascii(scene, id_to_char, shorten: bool = True) -> List[str]:
-    """
-    Convert JSON scene files from a list of lists of ints 
-    to a list of ASCII strings using id_to_char mapping.
-    If shorten is True, only the last 15 rows are kept.
-    Args:
-        scene: List[List[int]] - 2D array of tile IDs
-        id_to_char: Dict[int, str] - mapping from tile ID to ASCII character
-        shorten: bool - If True, will shorten the output to only include the first 15 rows 
-                        so A* Mario (for SNES graphics) to run without glitching
-    Returns:
-        List[str]: List of strings, each representing a row in ASCII
-    """
-    if shorten and len(scene) > 15:
-        scene = scene[-15:]  # Keep only the last 15 rows
-    return ["".join(id_to_char[num] for num in row) for row in scene]
-
-@dataclass
-class SampleOutput:
-    level: Optional[List[str]]
-    prompt: Optional[str] = None
-    img: Optional[Image] = None
-    sample_predictions_str: Optional[List[str]] = None
-    sample_predictions_img: Optional[Image] = None
-    level_tensor: Optional[torch.Tensor] = None
-    sample_predictions_tensor: Optional[torch.Tensor] = None
-    # Uses MarioEval graphics for rendering levels when True
-    use_snes_graphics: bool = False
-
-    @classmethod
-    def create(
-        cls,
-        level_tensor: torch.Tensor,
-        sample_predictions_tensor: torch.Tensor,
-        tokenizer,
-        prompter: Optional[Prompter] = None,
-    ) -> SampleOutput:
-        # batch = 1
-        level = None
-        img = None
-
-        try:
-            level = view_level(level_tensor, tokenizer)
-            img = convert_level_to_png(level)[0]
-        except Exception as e:
-            print(
-                f"Failed to generate string or image representation for full level! Got error {e}"
-            )
-            level = None
-            img = None
-        try:
-            sample_predictions_str = view_level(sample_predictions_tensor, tokenizer)
-            sample_predictions_img = convert_level_to_png(sample_predictions_str)[0]
-        except Exception as e:
-            print(
-                f"Failed to generate string or image representation for sampled predictions! Got error {e}"
-            )
-            sample_predictions_str = None
-            sample_predictions_img = None
-
-        prompt = None
-        if prompter is not None:
-            prompt = prompter(level_tensor)[0]
-
-        return SampleOutput(
-            level,
-            prompt,
-            img,
-            sample_predictions_str,
-            sample_predictions_img,
-            level_tensor,
-            sample_predictions_tensor,
-        )
-
-    @classmethod
-    def from_level_predictions(
-        cls,
-        level: torch.Tensor,
-        sample_predictions: torch.Tensor,
-        tokenizer,
-        prompter: Optional[Prompter] = None,
-    ) -> Union[SampleOutput, List[SampleOutput]]:
-        level_tensor = trim_level(level).squeeze().detach().cpu()
-        sample_predictions_tensor = (
-            trim_level(sample_predictions).squeeze().detach().cpu()
-        )
-
-        if len(level_tensor.shape) == 1:
-            return SampleOutput.create(
-                level_tensor, sample_predictions_tensor, tokenizer, prompter
-            )
-
-        out = []
-        for _level_tensor, _sample_predictions_tensor in zip(
-            level_tensor, sample_predictions_tensor
-        ):
-            sample_output = SampleOutput.create(
-                _level_tensor, _sample_predictions_tensor, tokenizer, prompter
-            )
-            out.append(sample_output)
-        return out
-
-    def save(self, filename: str) -> str:
-        save_level(self.level, filename)
-
-    @classmethod
-    def load(cls, filename: str) -> SampleOutput:
-        level = load_level(filename)
-        return SampleOutput(level=level)
-
-    def play(self, game="mario", level_idx=None, dataset_path=None):
-        """
-        Play the level using the specified game engine.
-        game: "mario" (default) or "loderunner"
-        """
-        if game == "loderunner":
-            import tempfile, json
-            # Convert self.level (list of strings) to Lode Runner JSON format
-            scene = [[c for c in row] for row in self.level]
-            lr_json = [{
-                "scene": scene,
-                "caption": ""
-            }]
-            with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as tmp:
-                json.dump(lr_json, tmp)
-                tmp_path = tmp.name
-            import sys, os
-            #sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
-            from LodeRunner.loderunner import main
-            tmp_path = tmp_path if dataset_path is None else dataset_path
-            print(f"Playing Lode Runner level interactively -- {tmp_path}!")
-            main.play_lr_level(tmp_path, level_index=level_idx if level_idx is not None else 1)
-        else:
-            if self.use_snes_graphics:
-                simulator = CustomSimulator(level=self.level, jar_path="MarioEval.jar")
-            else:
-                simulator = CustomSimulator(level=self.level, jar_path="NESMarioEval.jar")
-            simulator.interactive()
-
-    def run_astar(self, render=True):
-        if self.use_snes_graphics:
-            simulator = CustomSimulator(level=self.level, jar_path="MarioEval.jar")
-        else:
-            simulator = CustomSimulator(level=self.level, jar_path="NESMarioEval.jar")
-        return simulator.astar(render)
-
-class CustomSimulator:
-    """
-        The classic Mario simulator used by MarioGPT is generally,
-        better, but it doesn't return any information about
-        Mario's performance. The main point of this simulator
-        is that information about the performance of the agent
-        is printed to the console (though I still need a way
-        to caption and return that information)
-    """
-
-    def __init__(self, level, jar_path="MarioEval.jar"):
-        while len(level) > 15:
-            level.pop(0)
-        # For some reason, my older A* agent
-        # crashes on Mario levels with 16 rows or more
-
-        self.level = level
-        self.jar_path = jar_path
-
-    def interactive(self):
-        t = tempfile.NamedTemporaryFile(suffix=".txt", delete=False)
-        save_level(self.level, t.name)
-        print(f"Playing level interactively -- {t.name}!")
-        _ = subprocess.run(
-            ["java", "-jar", self.jar_path, "human", t.name, "human"],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-        )
-        t.close()
-        os.unlink(t.name)
-
-    def astar(self, render: bool = True):
-        t = tempfile.NamedTemporaryFile(suffix=".txt", delete=False)
-        save_level(self.level, t.name)
-        print(f"Running Astar agent on level! -- {t.name}")
-        render_str = "human" if render else "norender"
-        result = subprocess.run(
-            ["java", "-jar", self.jar_path, "astar", t.name, render_str],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-        )
-        t.close()
-        os.unlink(t.name)
-        # Combine stdout and stderr, decode to string, and return
-        output = result.stdout.decode("utf-8") + result.stderr.decode("utf-8")
-        return output
-
-def save_level(level: List[str], filename: str):
-    concatenated = "\n".join(level)
-    with open(filename, "w") as f:
-        f.write(concatenated)
-    return filename
-
-class GPTSampler:
-    def __init__(
-        self,
-        mario_lm: BaseMarioLM,
-        temperature: float = 2.0,
-        top_k: int = 16,
-        context_len: int = 700,
-        use_tqdm: bool = False,
-        use_argmax: bool = False,
-    ):
-        self.mario_lm = mario_lm
-        self.temperature = temperature
-        self.top_k = top_k
-        self.context_len = context_len
-        self.use_tqdm = use_tqdm
-        self.use_argmax = use_argmax
-        self.logits_processor = LogitsProcessorList()
-        self.logits_warper = LogitsProcessorList(
-            [
-                TopKLogitsWarper(top_k),  # number of characters
-                TemperatureLogitsWarper(temperature),
-            ]
-        )
-
-    @property
-    def device(self) -> torch.device:
-        return self.mario_lm.device
-
-    def step(
-        self,
-        seed: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        with torch.no_grad():
-            attention_mask = torch.ones_like(seed).to(seed.device)
-            input_ids = seed
-            out = self.mario_lm.lm(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                encoder_hidden_states=encoder_hidden_states,
-                token_type_ids=None,
-            )
-            logits = out.logits.detach()
-            if len(logits.shape) == 2:
-                logits = logits.view(1, 1, -1)
-            next_token_logits = logits[:, -1, :]
-
-            if self.use_argmax:
-                next_tokens = next_token_logits.argmax(-1)
-            else:
-                next_token_scores = self.logits_processor(input_ids, next_token_logits)
-                next_token_scores = self.logits_warper(input_ids, next_token_scores)
-                probs = torch.nn.functional.softmax(next_token_scores, dim=-1)
-                next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
-        return next_tokens, encoder_hidden_states
-
-    def sample(
-        self,
-        seed: Union[Optional[torch.Tensor], Optional[SampleOutput]] = None,
-        prompts: Optional[List[str]] = None,
-        num_steps: int = 1,
-        encoder_hidden_states: torch.Tensor = None,
-        return_tensor: bool = False,
-    ):
-        self.mario_lm.eval()
-        context_len = self.context_len - 28
-        with torch.no_grad():
-            if seed is None:
-                seed = self.mario_lm.generate_seed(1, batch_size=len(prompts)).to(
-                    self.device
-                )
-                out_tensor = seed.to(self.device)
-            elif isinstance(seed, SampleOutput):
-                out_tensor = seed.level_tensor.to(self.device).squeeze()
-            else:
-                out_tensor = seed.to(self.device).squeeze()
-            if len(out_tensor.shape) < 2:
-                # if we pass in a single seed vector, then we repeat for each prompt
-                # Otherwise, we treat inputs as separate seed-prompt pairs
-                out_tensor = out_tensor.view(1, -1).repeat(len(prompts), 1)
-            if encoder_hidden_states is None:
-                if prompts is not None:
-                    encoder_hidden_states = torch.stack(
-                        [
-                            self.mario_lm.prompter.output_hidden(prompt)
-                            for prompt in prompts
-                        ]
-                    )
-                else:
-                    encoder_hidden_states = torch.stack(
-                        [
-                            self.mario_lm.prompter(sample_prompt=True)[1]
-                            for _ in range(seed.shape[0])
-                        ]
-                    )
-            encoder_hidden_states = encoder_hidden_states.to(
-                self.device
-            )  # b x 1 x hidden_dim
-            encoder_hidden_states = encoder_hidden_states.view(
-                out_tensor.shape[0], 1, -1
-            )
-            if not self.use_tqdm:
-                bar = np.arange(num_steps)
-            else:
-                bar = tqdm(np.arange(num_steps))
-            with torch.no_grad():
-                for i in bar:
-                    inp = out_tensor * 1
-                    if len(out_tensor.shape) > 0 and out_tensor.shape[-1] > context_len:
-                        diff = inp.shape[-1] % 14  # height of mario level
-                        ctx = context_len + diff
-                        inp = inp[:, -ctx:] * 1
-                    next_tokens, encoder_hidden_states = self.step(
-                        inp,
-                        encoder_hidden_states=encoder_hidden_states,
-                    )
-                    out_tensor = torch.cat(
-                        [out_tensor, next_tokens.unsqueeze(-1)], dim=-1
-                    )
-                    if self.use_tqdm:
-                        bar.set_description(
-                            f"shape: {inp.shape}, {out_tensor.shape} first: {inp[0][0]}, last: {out_tensor[0][-1]}"
-                        )
-            if self.use_tqdm:
-                bar.close()
-        sample_out = SampleOutput.from_level_predictions(
-            out_tensor,
-            out_tensor[:, -num_steps:],
-            self.mario_lm.tokenizer,
-            self.mario_lm.prompter,
-        )
-        self.mario_lm.train()
-        if return_tensor:
-            return sample_out, out_tensor
-        return sample_out
-
-    def __call__(self, *args, **kwargs):
-        return self.sample(*args, **kwargs)
-
-
-class BertSampler:
-    def __init__(
-        self,
-        mario_lm: BaseMarioLM,
-        temperature: float = 2.0,
-        top_k: int = 16,
-        context_len: int = 448,
-        mask_proportion: float = 0.16,
-    ):
-        self.mario_lm = mario_lm
-        self.temperature = temperature
-        self.top_k = top_k
-        self.logits_processor = LogitsProcessorList()
-        self.logits_warper = LogitsProcessorList(
-            [
-                TopKLogitsWarper(top_k),  # number of characters
-                TemperatureLogitsWarper(temperature),
-            ]
-        )
-        self.context_len = context_len
-        self.mask_proportion = mask_proportion
-        self.mask_portion = int(self.context_len * self.mask_proportion)
-        self.mask_portion = self.mask_portion - self.mask_portion % 14 + 14
-
-    @property
-    def device(self) -> torch.device:
-        return self.mario_lm.device
-
-    def get_context(self, input_ids, mask_indices):
-        start_idx = mask_indices[0]
-        end_idx = mask_indices[-1]
-
-        if input_ids.shape[-1] <= self.context_len:
-            clipped = input_ids.shape[-1] % 14
-            input_ids = input_ids[:clipped]
-
-        portion = (self.context_len - self.mask_portion) / 2
-
-        remainder = 0
-        left = start_idx - portion
-        if left < 0:
-            remainder = -1 * left
-
-        right = end_idx + portion + remainder
-
-        return input_ids[left:right]
-
-    def sample(
-        self,
-        seed: Union[torch.Tensor, SampleOutput],
-        mask: torch.Tensor,
-        return_tensor: bool = False,
-    ):
-        self.mario_lm.eval()
-        mask_indices = mask.nonzero()
-        input_ids = seed
-        if isinstance(seed, SampleOutput):
-            input_ids = seed.level_tensor.to(self.device).squeeze()
-
-        input_id_list = []
-        for i in range(input_ids.shape[0]):
-            input_id = input_ids[i]
-            mask_index = mask_indices[mask_indices[:, 0] == i][:, -1]
-            input_id = self.get_context(input_id, mask_index)
-            input_id_list.append(input_id)
-        input_ids = torch.stack(input_ids, dim=0).to(self.device)
-
-        attention_mask = torch.ones_like(input_ids).to(seed.device)
-
-        if len(input_ids.shape) < 2:
-            # if we pass in a single seed vector, then we repeat for each prompt
-            # Otherwise, we treat inputs as separate seed-prompt pairs
-            input_ids = input_ids.view(1, -1)
-
-        out = self.mario_lm.lm(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            token_type_ids=None,
-        )
-        logits = out.logits.detach()
-        if len(logits.shape) == 2:
-            logits = logits.view(1, 1, -1)
-
-        if self.use_argmax:
-            tokens = logits.argmax(-1)
-        else:
-            tokens_scores = self.logits_processor(input_ids, tokens)
-            tokens_scores = self.logits_warper(input_ids, tokens_scores)
-            probs = torch.nn.functional.softmax(tokens_scores, dim=-1)
-            tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
-
-        out = input_ids.detach()
-
-        for i in range(input_ids.shape[0]):
-            mask_index = mask_indices[mask_indices[:, 0] == i][:, -1]
-            out[i, mask_index] = tokens[i, mask_index].detach()
-
-        sample_out = SampleOutput.from_level_predictions(
-            out,
-            tokens,
-            self.mario_lm.tokenizer,
-            self.mario_lm.prompter,
-        )
-        self.mario_lm.train()
-        if return_tensor:
-            return sample_out, tokens
-        return sample_out