From 4c0dc4cd3b296cf1d9008abf64f698cfac7c9cc8 Mon Sep 17 00:00:00 2001
From: Anri Lombard <anri.m.lombard@gmail.com>
Date: Sun, 11 Jan 2026 05:23:58 +0200
Subject: [PATCH 1/2] Remove unused lazy_import function

This function was introduced for lazy loading ocean environments but was
replaced by direct importlib.import_module() calls. The function definition
was left behind during that refactor.

Fixes #374
---
 pufferlib/ocean/environment.py | 8 --------
 1 file changed, 8 deletions(-)

diff --git a/pufferlib/ocean/environment.py b/pufferlib/ocean/environment.py
index 6c56a4ea2..a9fb90d9a 100644
--- a/pufferlib/ocean/environment.py
+++ b/pufferlib/ocean/environment.py
@@ -1,14 +1,6 @@
 import importlib
 import pufferlib.emulation
 
-def lazy_import(module_path, attr):
-    """
-    Returns a callable that, when called with any arguments, will
-    import the module, retrieve the attribute (usually a class or factory)
-    and then call it with the given arguments.
-    """
-    return lambda *args, **kwargs: getattr(__import__(module_path, fromlist=[attr]), attr)(*args, **kwargs)
-
 def make_foraging(width=1080, height=720, num_agents=4096, horizon=512,
         discretize=True, food_reward=0.1, render_mode='rgb_array'):
     from .grid import grid

From a180ead008b0585ac87d39e20242c1f9e9858ff9 Mon Sep 17 00:00:00 2001
From: Anri Lombard <anri.m.lombard@gmail.com>
Date: Sun, 11 Jan 2026 05:44:15 +0200
Subject: [PATCH 2/2] Remove dead code across codebase

- Remove unused pdb debug imports from 40 files
- Remove dead functions from pytorch.py: compilable_cast, entropy_probs, nativize_observation
- Remove unused Categorical import from pytorch.py
- Remove unused dist_mean function from pufferl.py
---
 pufferlib/emulation.py                        |  1 -
 pufferlib/environments/__init__.py            |  1 -
 pufferlib/environments/atari/environment.py   |  1 -
 pufferlib/environments/box2d/environment.py   |  1 -
 pufferlib/environments/bsuite/environment.py  |  1 -
 .../environments/butterfly/environment.py     |  1 -
 pufferlib/environments/crafter/environment.py |  1 -
 .../environments/dm_control/environment.py    |  1 -
 pufferlib/environments/dm_lab/environment.py  |  1 -
 pufferlib/environments/griddly/environment.py |  1 -
 pufferlib/environments/gvgai/environment.py   |  1 -
 .../environments/links_awaken/environment.py  |  1 -
 pufferlib/environments/magent/environment.py  |  1 -
 .../environments/microrts/environment.py      |  1 -
 pufferlib/environments/minerl/environment.py  |  1 -
 .../environments/minigrid/environment.py      |  1 -
 .../environments/minihack/environment.py      |  1 -
 pufferlib/environments/minihack/torch.py      |  1 -
 pufferlib/environments/mujoco/environment.py  |  1 -
 pufferlib/environments/nethack/environment.py |  1 -
 pufferlib/environments/nethack/torch.py       |  1 -
 pufferlib/environments/nmmo/environment.py    |  1 -
 pufferlib/environments/nmmo/torch.py          |  1 -
 .../environments/open_spiel/environment.py    |  1 -
 .../open_spiel/gymnasium_environment.py       |  1 -
 .../open_spiel/pettingzoo_environment.py      |  1 -
 pufferlib/environments/open_spiel/torch.py    |  1 -
 pufferlib/environments/open_spiel/utils.py    |  1 -
 .../environments/pokemon_red/environment.py   |  1 -
 pufferlib/environments/procgen/environment.py |  1 -
 pufferlib/environments/procgen/torch.py       |  1 -
 .../environments/slimevolley/environment.py   |  1 -
 .../environments/stable_retro/environment.py  |  1 -
 pufferlib/environments/test/environment.py    |  1 -
 .../environments/test/mock_environments.py    |  1 -
 pufferlib/environments/vizdoom/environment.py |  1 -
 pufferlib/models.py                           |  1 -
 pufferlib/ocean/moba/moba.py                  |  1 -
 pufferlib/ocean/nmmo3/nmmo3.py                |  1 -
 pufferlib/pufferl.py                          |  6 ----
 pufferlib/pytorch.py                          | 32 -------------------
 pufferlib/vector.py                           |  1 -
 42 files changed, 78 deletions(-)

diff --git a/pufferlib/emulation.py b/pufferlib/emulation.py
index f576f530f..077581dda 100644
--- a/pufferlib/emulation.py
+++ b/pufferlib/emulation.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 
 import numpy as np
 import warnings
diff --git a/pufferlib/environments/__init__.py b/pufferlib/environments/__init__.py
index d6c969523..062b5e051 100644
--- a/pufferlib/environments/__init__.py
+++ b/pufferlib/environments/__init__.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import pufferlib
 
 def try_import(module_path, package_name=None):
diff --git a/pufferlib/environments/atari/environment.py b/pufferlib/environments/atari/environment.py
index 92295ecaa..b5e91c284 100644
--- a/pufferlib/environments/atari/environment.py
+++ b/pufferlib/environments/atari/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 import functools
 
diff --git a/pufferlib/environments/box2d/environment.py b/pufferlib/environments/box2d/environment.py
index f2fd12b7e..92b2ea075 100644
--- a/pufferlib/environments/box2d/environment.py
+++ b/pufferlib/environments/box2d/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 
 import gymnasium
 import functools
diff --git a/pufferlib/environments/bsuite/environment.py b/pufferlib/environments/bsuite/environment.py
index f39c661e5..ac43c151e 100644
--- a/pufferlib/environments/bsuite/environment.py
+++ b/pufferlib/environments/bsuite/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import gym
 import functools
 
diff --git a/pufferlib/environments/butterfly/environment.py b/pufferlib/environments/butterfly/environment.py
index 994e53e9e..a41652fb9 100644
--- a/pufferlib/environments/butterfly/environment.py
+++ b/pufferlib/environments/butterfly/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 from pettingzoo.utils.conversions import aec_to_parallel_wrapper
 import functools
 
diff --git a/pufferlib/environments/crafter/environment.py b/pufferlib/environments/crafter/environment.py
index e320e347d..8b172d7d9 100644
--- a/pufferlib/environments/crafter/environment.py
+++ b/pufferlib/environments/crafter/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 
 import gym
 import gymnasium
diff --git a/pufferlib/environments/dm_control/environment.py b/pufferlib/environments/dm_control/environment.py
index f492fa659..c00314ad1 100644
--- a/pufferlib/environments/dm_control/environment.py
+++ b/pufferlib/environments/dm_control/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 
 import gym
 import shimmy
diff --git a/pufferlib/environments/dm_lab/environment.py b/pufferlib/environments/dm_lab/environment.py
index 17bbde770..1ef7ce1e6 100644
--- a/pufferlib/environments/dm_lab/environment.py
+++ b/pufferlib/environments/dm_lab/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 
 import gym
 import shimmy
diff --git a/pufferlib/environments/griddly/environment.py b/pufferlib/environments/griddly/environment.py
index fdf8c2430..6b50223f6 100644
--- a/pufferlib/environments/griddly/environment.py
+++ b/pufferlib/environments/griddly/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 
 import gym
 import shimmy
diff --git a/pufferlib/environments/gvgai/environment.py b/pufferlib/environments/gvgai/environment.py
index fa4da86fd..a20293cf2 100644
--- a/pufferlib/environments/gvgai/environment.py
+++ b/pufferlib/environments/gvgai/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 import functools
 
diff --git a/pufferlib/environments/links_awaken/environment.py b/pufferlib/environments/links_awaken/environment.py
index 8016c3cf5..de78e2f32 100644
--- a/pufferlib/environments/links_awaken/environment.py
+++ b/pufferlib/environments/links_awaken/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 
 import gymnasium
 
diff --git a/pufferlib/environments/magent/environment.py b/pufferlib/environments/magent/environment.py
index 10fbfc2ad..5428008b9 100644
--- a/pufferlib/environments/magent/environment.py
+++ b/pufferlib/environments/magent/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 from pettingzoo.utils.conversions import aec_to_parallel_wrapper
 import functools
 
diff --git a/pufferlib/environments/microrts/environment.py b/pufferlib/environments/microrts/environment.py
index 6c00f4d52..6db53efea 100644
--- a/pufferlib/environments/microrts/environment.py
+++ b/pufferlib/environments/microrts/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 
 import warnings
diff --git a/pufferlib/environments/minerl/environment.py b/pufferlib/environments/minerl/environment.py
index 771c3c600..7cb6c26b9 100644
--- a/pufferlib/environments/minerl/environment.py
+++ b/pufferlib/environments/minerl/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 
 import gym
 import shimmy
diff --git a/pufferlib/environments/minigrid/environment.py b/pufferlib/environments/minigrid/environment.py
index 2bf373b7e..f8045716b 100644
--- a/pufferlib/environments/minigrid/environment.py
+++ b/pufferlib/environments/minigrid/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 
 import gymnasium
 import functools
diff --git a/pufferlib/environments/minihack/environment.py b/pufferlib/environments/minihack/environment.py
index d70ea03d5..4daab21fd 100644
--- a/pufferlib/environments/minihack/environment.py
+++ b/pufferlib/environments/minihack/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 
 import gym
 import shimmy
diff --git a/pufferlib/environments/minihack/torch.py b/pufferlib/environments/minihack/torch.py
index 7781cb677..89d07c5c7 100644
--- a/pufferlib/environments/minihack/torch.py
+++ b/pufferlib/environments/minihack/torch.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 
 import pufferlib.pytorch
 from pufferlib.environments.nethack import Policy
diff --git a/pufferlib/environments/mujoco/environment.py b/pufferlib/environments/mujoco/environment.py
index f9adad9f8..a9b4b37a3 100644
--- a/pufferlib/environments/mujoco/environment.py
+++ b/pufferlib/environments/mujoco/environment.py
@@ -1,5 +1,4 @@
 
-from pdb import set_trace as T
 
 import functools
 
diff --git a/pufferlib/environments/nethack/environment.py b/pufferlib/environments/nethack/environment.py
index 033b1d4e0..c003128ed 100644
--- a/pufferlib/environments/nethack/environment.py
+++ b/pufferlib/environments/nethack/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 
 import shimmy
 import gym
diff --git a/pufferlib/environments/nethack/torch.py b/pufferlib/environments/nethack/torch.py
index 64af81e54..a6c0f697a 100644
--- a/pufferlib/environments/nethack/torch.py
+++ b/pufferlib/environments/nethack/torch.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 
 import torch
 import torch.nn as nn
diff --git a/pufferlib/environments/nmmo/environment.py b/pufferlib/environments/nmmo/environment.py
index f4ef1c998..e885bf6d4 100644
--- a/pufferlib/environments/nmmo/environment.py
+++ b/pufferlib/environments/nmmo/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 import functools
 
diff --git a/pufferlib/environments/nmmo/torch.py b/pufferlib/environments/nmmo/torch.py
index e98134f3a..8615629a7 100644
--- a/pufferlib/environments/nmmo/torch.py
+++ b/pufferlib/environments/nmmo/torch.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 
 import torch
 import torch.nn.functional as F
diff --git a/pufferlib/environments/open_spiel/environment.py b/pufferlib/environments/open_spiel/environment.py
index a5679b759..5fbfecf63 100644
--- a/pufferlib/environments/open_spiel/environment.py
+++ b/pufferlib/environments/open_spiel/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 import functools
 
diff --git a/pufferlib/environments/open_spiel/gymnasium_environment.py b/pufferlib/environments/open_spiel/gymnasium_environment.py
index ae5bfba59..c236702e1 100644
--- a/pufferlib/environments/open_spiel/gymnasium_environment.py
+++ b/pufferlib/environments/open_spiel/gymnasium_environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 
 from open_spiel.python.algorithms import mcts
diff --git a/pufferlib/environments/open_spiel/pettingzoo_environment.py b/pufferlib/environments/open_spiel/pettingzoo_environment.py
index cbfca0af4..956873e3e 100644
--- a/pufferlib/environments/open_spiel/pettingzoo_environment.py
+++ b/pufferlib/environments/open_spiel/pettingzoo_environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 
 import pufferlib
diff --git a/pufferlib/environments/open_spiel/torch.py b/pufferlib/environments/open_spiel/torch.py
index 4d1845b6b..81705bd5d 100644
--- a/pufferlib/environments/open_spiel/torch.py
+++ b/pufferlib/environments/open_spiel/torch.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 
 import torch
diff --git a/pufferlib/environments/open_spiel/utils.py b/pufferlib/environments/open_spiel/utils.py
index 8eab105c2..35b667e80 100644
--- a/pufferlib/environments/open_spiel/utils.py
+++ b/pufferlib/environments/open_spiel/utils.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 
 import gymnasium
diff --git a/pufferlib/environments/pokemon_red/environment.py b/pufferlib/environments/pokemon_red/environment.py
index fe0bd85e7..5c89362a6 100644
--- a/pufferlib/environments/pokemon_red/environment.py
+++ b/pufferlib/environments/pokemon_red/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 
 import gymnasium
 import functools
diff --git a/pufferlib/environments/procgen/environment.py b/pufferlib/environments/procgen/environment.py
index 3d889b73f..71f8c1ac6 100644
--- a/pufferlib/environments/procgen/environment.py
+++ b/pufferlib/environments/procgen/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 
 import gym
diff --git a/pufferlib/environments/procgen/torch.py b/pufferlib/environments/procgen/torch.py
index 924926c34..8a7f00ada 100644
--- a/pufferlib/environments/procgen/torch.py
+++ b/pufferlib/environments/procgen/torch.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 from torch import nn
 import pufferlib.models
 
diff --git a/pufferlib/environments/slimevolley/environment.py b/pufferlib/environments/slimevolley/environment.py
index a46d79562..9936a6df6 100644
--- a/pufferlib/environments/slimevolley/environment.py
+++ b/pufferlib/environments/slimevolley/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 import functools
 
diff --git a/pufferlib/environments/stable_retro/environment.py b/pufferlib/environments/stable_retro/environment.py
index 092708613..bfeaa3b2b 100644
--- a/pufferlib/environments/stable_retro/environment.py
+++ b/pufferlib/environments/stable_retro/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 
 import gymnasium as gym
diff --git a/pufferlib/environments/test/environment.py b/pufferlib/environments/test/environment.py
index ff82c6e47..a81784dcd 100644
--- a/pufferlib/environments/test/environment.py
+++ b/pufferlib/environments/test/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 
 import time
diff --git a/pufferlib/environments/test/mock_environments.py b/pufferlib/environments/test/mock_environments.py
index 8e8db330e..e062c150c 100644
--- a/pufferlib/environments/test/mock_environments.py
+++ b/pufferlib/environments/test/mock_environments.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 
 import time
diff --git a/pufferlib/environments/vizdoom/environment.py b/pufferlib/environments/vizdoom/environment.py
index d98a263f9..779cc9256 100644
--- a/pufferlib/environments/vizdoom/environment.py
+++ b/pufferlib/environments/vizdoom/environment.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 import functools
 
diff --git a/pufferlib/models.py b/pufferlib/models.py
index fa43d7071..c32aed032 100644
--- a/pufferlib/models.py
+++ b/pufferlib/models.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 
 import torch
diff --git a/pufferlib/ocean/moba/moba.py b/pufferlib/ocean/moba/moba.py
index 4b4dd0be8..a14b5a892 100644
--- a/pufferlib/ocean/moba/moba.py
+++ b/pufferlib/ocean/moba/moba.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 import os
 
diff --git a/pufferlib/ocean/nmmo3/nmmo3.py b/pufferlib/ocean/nmmo3/nmmo3.py
index 32623501b..9e6c97657 100644
--- a/pufferlib/ocean/nmmo3/nmmo3.py
+++ b/pufferlib/ocean/nmmo3/nmmo3.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as T
 import numpy as np
 from types import SimpleNamespace
 import gymnasium
diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py
index 7132a19f9..d8da14387 100644
--- a/pufferlib/pufferl.py
+++ b/pufferlib/pufferl.py
@@ -713,12 +713,6 @@ def dist_sum(value, device):
     torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM)
     return tensor.item()
 
-def dist_mean(value, device):
-    if not torch.distributed.is_initialized():
-        return value
-
-    return dist_sum(value, device) / torch.distributed.get_world_size()
-
 class Profile:
     def __init__(self, frequency=5):
         self.profiles = defaultdict(lambda: defaultdict(float))
diff --git a/pufferlib/pytorch.py b/pufferlib/pytorch.py
index e18bea40c..0dc134c46 100644
--- a/pufferlib/pytorch.py
+++ b/pufferlib/pytorch.py
@@ -1,11 +1,9 @@
 import sys
-from pdb import set_trace as T
 from typing import Dict, List, Tuple, Union
 
 import numpy as np
 import torch
 from torch import nn
-from torch.distributions import Categorical
 from torch.distributions.utils import logits_to_probs
 
 import pufferlib
@@ -103,23 +101,6 @@ def nativize_tensor(observation: torch.Tensor, native_dtype: NativeDType):
     return _nativize_tensor(observation, native_dtype)
 
 
-# torch.view(dtype) does not compile
-# This is a workaround hack
-# @thatguy - can you figure out a more robust way to handle cast?
-# I think it may screw up for non-uint data... so I put a hard .view
-# fallback that breaks compile
-def compilable_cast(u8, dtype):
-    if dtype in (torch.uint8, torch.uint16, torch.uint32, torch.uint64):
-        n = dtype.itemsize
-        bytes = [u8[..., i::n].to(dtype) for i in range(n)]
-        if not LITTLE_BYTE_ORDER:
-            bytes = bytes[::-1]
-
-        bytes = sum(bytes[i] << (i * 8) for i in range(n))
-        return bytes.view(dtype)
-    return u8.view(dtype)  # breaking cast
-
-
 def _nativize_tensor(observation: torch.Tensor, native_dtype: NativeDType):
     if isinstance(native_dtype, tuple):
         dtype, shape, offset, delta = native_dtype
@@ -141,15 +122,6 @@ def _nativize_tensor(observation: torch.Tensor, native_dtype: NativeDType):
         return subviews
 
 
-def nativize_observation(observation, emulated):
-    # TODO: Any way to check that user has not accidentally cast data to float?
-    # float is natively supported, but only if that is the actual correct type
-    return nativize_tensor(
-        observation,
-        emulated['observation_dtype'],
-        emulated['emulated_observation_dtype'],
-    )
-
 def flattened_tensor_size(native_dtype):
     return _flattened_tensor_size(native_dtype)
 
@@ -182,10 +154,6 @@ def entropy(logits):
     p_log_p = logits * logits_to_probs(logits)
     return -p_log_p.sum(-1)
 
-def entropy_probs(logits, probs):
-    p_log_p = logits * probs
-    return -p_log_p.sum(-1)
-
 def sample_logits(logits, action=None):
     is_discrete = isinstance(logits, torch.Tensor)
     if isinstance(logits, torch.distributions.Normal):
diff --git a/pufferlib/vector.py b/pufferlib/vector.py
index 78614f4d6..f47e14d12 100644
--- a/pufferlib/vector.py
+++ b/pufferlib/vector.py
@@ -1,6 +1,5 @@
 # TODO: Check actions passed to envs are right shape? On first call at least
 
-from pdb import set_trace as T
 
 import numpy as np
 import time