Onnx scripts with verification of model structure by mpragnay · Pull Request #305 · Emerge-Lab/PufferDrive

mpragnay · 2026-02-19T14:09:28Z

Onnx converter
Onnx I/O and structure Verifier
Creates .pt for I/O to verify

greptile-apps · 2026-02-19T14:12:38Z

Greptile Summary

Added three utility scripts for exporting and verifying ONNX models, plus 18MB of model files directly to the repository.

The scripts provide functionality to:

Export PyTorch model weights to .bin format (export_model.py)
Convert models to ONNX format with verification (export_onnx.py)
Verify ONNX model structure and run test inference (verify_onnx.py)

Critical Issue:

The hardcoded path /scratch/pm3881/PufferDrive/... in export_model.py:62 will fail for all other users

Major Concerns:

18MB of binary files (ONNX model + weights) committed to repository root without Git LFS
8.8MB .bin file in pufferlib/resources/drive/ also committed without Git LFS
These files will permanently bloat the repository and should be excluded via .gitignore or managed with Git LFS

Code Quality Issues:

Bare except clauses in both export scripts should catch specific exceptions
Duplicate print statement in export_onnx.py (lines 197 and 203)

Confidence Score: 2/5

This PR has a critical hardcoded path bug and commits large binary files without Git LFS
Score reflects one blocking bug (hardcoded user-specific path will fail for others) and significant repository management concern (18MB+ of binaries committed directly to git). The export scripts are functional but need the path fix before merging.
scripts/export_model.py requires path fix; all binary files (.onnx, .onnx.data, .bin) should be removed or managed with Git LFS

Important Files Changed

Filename	Overview
scripts/export_model.py	Adds script to export PyTorch model weights to binary format; hardcoded default path could be problematic
scripts/export_onnx.py	Adds ONNX export with verification; has bare except clause and duplicate print statement
scripts/verify_onnx.py	Adds ONNX verification script with structure check and inference test; clean implementation
model_puffer_drive_000100.onnx	Adds 79KB ONNX model file to repository root; should not be committed to git
model_puffer_drive_000100.onnx.data	Adds 8.9MB ONNX weights file to repository root; should not be committed to git
pufferlib/resources/drive/puffer_drive_ntski86s.bin	Adds 8.8MB binary model weights; should not be committed without Git LFS

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[PyTorch Checkpoint .pt] --> B[export_model.py]
    A --> C[export_onnx.py]
    
    B --> D[Binary Weights .bin]
    D --> E[pufferlib/resources/drive/]
    
    C --> F[ONNX Model .onnx]
    C --> G[ONNX Data .onnx.data]
    C --> H[Test I/O .pt files]
    
    F --> I[verify_onnx.py]
    G --> I
    
    I --> J{Validation}
    J -->|Structure Check| K[onnx.checker]
    J -->|Runtime Test| L[onnxruntime]
    
    C -->|if verify=True| M[Compare Outputs]
    M -->|PyTorch vs ONNX| N[Numerical Verification]

_{Last reviewed commit: 6f1e654}

greptile-apps

_{6 files reviewed, 4 comments}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps · 2026-02-19T14:12:41Z

scripts/export_model_bin.py

+    def puffer_type(value):
+        try:
+            return ast.literal_eval(value)
+        except:
+            return value


bare except catches all exceptions including KeyboardInterrupt and SystemExit

Suggested change

def puffer_type(value):

try:

return ast.literal_eval(value)

except:

return value

def puffer_type(value):

try:

return ast.literal_eval(value)

except (ValueError, SyntaxError):

return value

greptile-apps · 2026-02-19T14:12:43Z

scripts/export_model.py

+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        default="/scratch/pm3881/PufferDrive/experiments/puffer_drive_ntski86s/models/model_puffer_drive_000100.pt",


hardcoded absolute path specific to user's scratch directory will fail for other users

scripts/export_onnx.py

Copilot

Pull request overview

Adds helper scripts to export PufferDrive PyTorch checkpoints to ONNX / flat weight binaries and to validate exported ONNX models via structural checks and runtime inference.

Changes:

Added ONNX export script that can optionally run an ORT vs PyTorch output comparison and save .pt I/O artifacts.
Added ONNX verification script that runs onnx.checker and a dummy-input ORT inference pass.
Added weight export script to dump model parameters into a single .bin file.

Reviewed changes

Copilot reviewed 3 out of 6 changed files in this pull request and generated 6 comments.

File	Description
scripts/verify_onnx.py	New CLI to validate ONNX structure and do a basic ORT execution smoke test with generated inputs.
scripts/export_onnx.py	New CLI to export a Drive(+optional LSTM wrapper) policy to ONNX and optionally compare ORT outputs vs PyTorch.
scripts/export_model.py	New CLI to flatten and export model weights from a checkpoint into a `.bin` file.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

scripts/export_onnx.py

Copilot · 2026-02-19T14:14:36Z

scripts/verify_onnx.py

+            dtype = input_meta.type
+
+            # Handle dynamic axes (often represented as strings or -1)
+            processed_shape = []
+            for dim in shape:
+                if isinstance(dim, str) or dim is None or dim < 0:
+                    processed_shape.append(batch_size)
+                else:
+                    processed_shape.append(dim)
+
+            print(f"  Input: {name}, Shape: {shape} -> Using: {processed_shape}, Type: {dtype}")
+
+            # Create random input data
+            if "float" in dtype:
+                data = np.random.randn(*processed_shape).astype(np.float32)
+            else:
+                data = np.zeros(processed_shape).astype(np.int64)  # Fallback
+


Input dtype handling is too coarse: input_meta.type is typically a string like tensor(float16), tensor(double), tensor(int32), tensor(bool), etc. Always using float32 for any float type (and int64 for everything else) can make ORT inference fail purely due to dtype mismatch, even when the model is valid. Map ONNX Runtime input types to the correct NumPy dtype (e.g., float16/float32/float64/int32/int64/bool) and generate matching dummy arrays.

Copilot · 2026-02-19T14:14:36Z

scripts/export_model_bin.py

+        "--checkpoint",
+        type=str,
+        default="/scratch/pm3881/PufferDrive/experiments/puffer_drive_ntski86s/models/model_puffer_drive_000100.pt",
+        help="Path to .pt checkpoint",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="pufferlib/resources/drive/model_puffer_drive_000100.bin",
+        help="Output .bin file path",
+    )


The default --checkpoint path is machine/user-specific (/scratch/...) and will break for most users running the script. Prefer no default (require the flag) or use a repo-relative default similar to scripts/export_onnx.py so the script is portable.

Copilot · 2026-02-19T14:14:37Z

scripts/export_model.py

+
+    # Use valid dummy env to initialize policy
+    # Ensure env args/kwargs are correctly passed as expected by make()
+    env_args = []


env_args is defined but never used in this script. Removing it will reduce confusion about whether positional env args are expected/required here.

Suggested change

env_args = []

Copilot · 2026-02-19T14:14:37Z

scripts/export_onnx.py

+
+
+def load_config(env_name, config_dir=None):
+    # Minimal config loader based on pufferl.py
+    import configparser
+    import glob
+    from collections import defaultdict
+    import ast
+
+    if config_dir is None:
+        puffer_dir = os.path.dirname(os.path.realpath(pufferlib.__file__))
+    else:
+        puffer_dir = config_dir
+
+    puffer_config_dir = os.path.join(puffer_dir, "config/**/*.ini")
+    puffer_default_config = os.path.join(puffer_dir, "config/default.ini")
+
+    found = False
+    for path in glob.glob(puffer_config_dir, recursive=True):
+        p = configparser.ConfigParser()
+        p.read([puffer_default_config, path])
+        if env_name in p["base"]["env_name"].split():
+            found = True
+            break
+
+    if not found:
+        raise ValueError(f"No config for env_name {env_name}")
+
+    def puffer_type(value):
+        try:
+            return ast.literal_eval(value)
+        except:
+            return value
+
+    args = defaultdict(dict)
+    for section in p.sections():
+        for key in p[section]:
+            value = puffer_type(p[section][key])
+            args[section][key] = value
+
+    return args
+
+


load_config is duplicated (also present in scripts/export_model.py). Duplicated config parsing logic tends to drift over time; consider moving it into a shared helper module under scripts/ (or an existing internal utility location) and importing it from both scripts.

Suggested change

def load_config(env_name, config_dir=None):

# Minimal config loader based on pufferl.py

import configparser

import glob

from collections import defaultdict

import ast

if config_dir is None:

puffer_dir = os.path.dirname(os.path.realpath(pufferlib.__file__))

else:

puffer_dir = config_dir

puffer_config_dir = os.path.join(puffer_dir, "config/**/*.ini")

puffer_default_config = os.path.join(puffer_dir, "config/default.ini")

found = False

for path in glob.glob(puffer_config_dir, recursive=True):

p = configparser.ConfigParser()

p.read([puffer_default_config, path])

if env_name in p["base"]["env_name"].split():

found = True

break

if not found:

raise ValueError(f"No config for env_name {env_name}")

def puffer_type(value):

try:

return ast.literal_eval(value)

except:

return value

args = defaultdict(dict)

for section in p.sections():

for key in p[section]:

value = puffer_type(p[section][key])

args[section][key] = value

return args

from scripts.export_model import load_config

I agree with the copilot review here?

eugenevinitsky · 2026-02-19T14:35:39Z

Nice! I think the folder structure of this could be a bit more organized? This should probably be a folder within scripts and the test files in there also loaded into that folder

Onnx scripts with verification of model structure

6f1e654

Copilot AI review requested due to automatic review settings February 19, 2026 14:09

Copilot started reviewing on behalf of mpragnay February 19, 2026 14:10 View session

greptile-apps bot reviewed Feb 19, 2026

View reviewed changes

Copilot AI reviewed Feb 19, 2026

View reviewed changes

mpragnay added 2 commits February 19, 2026 13:12

Code cleanup

e8bafad

Added documentation for model export scripts

7cb93db

Conversation

mpragnay commented Feb 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

greptile-apps bot commented Feb 19, 2026

Greptile Summary

Confidence Score: 2/5

Important Files Changed

Flowchart

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Feb 19, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Feb 19, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Copilot AI Feb 19, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 19, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 19, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 19, 2026

Choose a reason for hiding this comment

Uh oh!

eugenevinitsky Feb 19, 2026

Choose a reason for hiding this comment

Uh oh!

eugenevinitsky commented Feb 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Comments

mpragnay commented Feb 19, 2026 •

edited

Loading