NitroGen Documentation

NitroGen is an open foundation model for generalist gaming agents. It is a 500M parameter Diffusion Transformer (DiT) that takes pixel input and predicts gamepad actions.

Key Characteristics:

• Training: Trained via behavior cloning on a large video-action gameplay dataset from internet videos.
• Capabilities: Acts as a fast-reacting 'system-1' sensory model. It can be adapted via post-training to unseen games.
• Limitations: Because it only sees the last frame, it currently lacks the ability to plan over long horizons, play games end-to-end, self-improve, or play completely unseen games without adaptation.

To install NitroGen, clone the repository and install it in editable mode using pip. Ensure you have Python ≥ 3.12 installed.

Note on Environments: NitroGen does not distribute game environments. You must provide your own copies of the games. This repository is designed to run agents on Windows games. While you can serve the model from a Linux machine for inference, the game itself must run on Windows.

git clone https://github.com/MineDojo/NitroGen.git
cd NitroGen
pip install -e .

Download the NitroGen checkpoint (ng.pt) from HuggingFace using the hf command line tool.

hf download nvidia/NitroGen ng.pt

To use NitroGen, you must first start an inference server and then run the agent against a specific game process.

1. Start the inference server

Run the serve.py script pointing to your downloaded checkpoint:

python scripts/serve.py <path_to_ng.pt>

2. Run the agent on a game

Run the play.py script using the --process flag to specify the game's executable name.

Finding the executable name:

1. Open Windows Task Manager (Ctrl+Shift+Esc).
2. Right-click the running game process and select Properties.
3. The name in the General tab (ending in .exe) is the value for the --process parameter.

The GamepadEnv uses one of two backend classes to capture the game window:

1. DxcamScreenshotBackend: Uses the dxcam library for high-performance, high-FPS screen capture. Requires a bbox (bounding box) and fps.
2. PyautoguiScreenshotBackend: Uses pyautogui for standard screen capture. Requires a bbox.

Both backends implement a .screenshot() method that returns a PIL.Image object.

The InferenceSession uses a sliding window approach to provide temporal context to the model. It maintains two internal collections.deque buffers:

1. obs_buffer: Stores processed image frames. Its maximum size is determined by context_length or the frame_per_sample value in the checkpoint's modality_cfg.
2. action_buffer: Stores previously predicted actions. This is used when action_interleaving is enabled in the modality configuration, allowing the model to condition on its own previous outputs.

When predict() is called, the session concatenates the current buffer contents to form the input tensor. For Flow Matching models, it also manages dropped_frames masks to account for cases where the buffer is not yet full.

Use NitrogenTokenizerConfig to define the hyperparameters for the tokenization process.

Key fields:

• tokenizer_id: Fixed to 'nitrogen'.
• training (bool): If True, applies transformations like action packing. Defaults to True.
• num_visual_tokens_per_frame (int): Number of visual tokens per frame. Defaults to 256.
• max_action_dim (int): Maximum dimension for the action vector. Defaults to 25.
• max_sequence_length (int): Maximum sequence length for the vision-language tokens. Defaults to 300.
• action_horizon (int): The number of steps in the action sequence. Defaults to 16.
• game_mapping_cfg (GameMappingConfig | None): Configuration for building game-to-ID mappings from parquet files.
• old_layout (bool): If True, action layout is [buttons, j_left, j_right]. If False, layout is [j_left, j_right, buttons]. Defaults to False.

Use the predict(obs) method to perform inference on a single observation.

Workflow:

1. The method processes the input obs using the internal img_proc.
2. The processed frame is appended to an internal obs_buffer (a deque of size max_buffer_size).
3. It automatically handles history management, including interleaving actions if configured.
4. It selects the appropriate inference method based on whether the model is a Flow Matching model (is_flowmatching) or an Autoregressive model.
5. It returns a dictionary containing the decoded actions.

Returns: A dictionary with the following keys:

• j_left: NumPy array of left joystick values.
• j_right: NumPy array of right joystick values.
• buttons: NumPy array of button states.

An InferenceSession maintains internal state to support temporal context (history) during inference.

• reset(): Clears both the obs_buffer and action_buffer. Call this when starting a new episode or if the environment state has changed significantly to prevent history leakage from previous runs.
• info(): Returns a dictionary containing the current session configuration, including ckpt_path, selected_game, cfg_scale, context_length, and whether the model uses flow matching.

The NitrogenTokenizer is a concrete implementation of the Tokenizer base class designed to prepare video, language, state, and action data into a tokenized format suitable for multi-modal models.

It handles:

• Action Packing: Combines buttons, j_left, and j_right into a single action tensor.
• Token ID Generation: Creates vl_token_ids (vision-language) and sa_token_ids (state-action) arrays.
• Attention Masking: Generates vl_attn_mask for vision-language tokens, padding them to max_sequence_length using _PAD_TOKEN.
• Game Mapping: Optionally maps game names to integer IDs using a provided GameMappingConfig.
• Action Layouts: Supports both old_layout ([buttons, j_left, j_right]) and the default layout ([j_left, j_right, buttons]).

The draw_button_grid function renders a grid of squares representing button states on an image.

Parameters:

• img (np.ndarray): The image to draw on.
• x (int): Starting X-coordinate.
• y (int): Starting Y-coordinate.
• button_size (int): The size of each button square.
• buttons (np.ndarray): A 2D array of boolean values representing button states.
• current_row (int): The index of the row to highlight with a red border.
• token_set (list, optional): A list of button names used to generate a legend below the grid.

Visual Behavior:

• Pressed buttons: Rendered as green squares.
• Unpressed buttons: Rendered as black squares.
• Current row: Highlighted with a red rectangle.
• Legend: If token_set is provided, a legend is drawn below the grid mapping column indices to names.