Monkey Patch失效排查：修改Tianshou库类方法失败问题

问题背景 #

在conda环境中通过pip安装Tianshou库后，尝试用Monkey Patch修改BaseTrainer类的test_step方法逻辑，实际操作中选择替换tianshou.trainer.utils路径下的test_episode函数。主程序优先导入自定义的monkey模块，再导入Tianshou相关库，调用继承自BaseTrainer的OffpolicyTrainer类，但两种patch写法均未生效，运行后看不到预期的print输出。

现有代码 #

主程序（~/projects/testing_ground/main.py） #

import testing_ground.monkey
from tianshou.data import Collector
from tianshou.env import DummyVectorEnv
from tianshou.utils import TensorboardLogger

def auxiliar_func():
    #do_something
def main_func():
    #do_main_thing
if __name__ == '__main__':
    auxiliar_func()
    main_func()

monkey.py的两种写法 #

写法一

import tianshou.trainer.utils
def my_test_episode(*args, **kwargs):
    result = tianshou.trainer.utils.test_episode(*args, **kwargs)
    print("check comment in execution to know it works")
    return result
setattr(tianshou.trainer.utils, 'test_episode', my_test_episode)

写法二

import tianshou.trainer.utils
def my_test_episode(*args, **kwargs):
    result = tianshou.trainer.utils.test_episode(*args, **kwargs)
    print("check comment in execution to know it works")
    return result
tianshou.trainer.utils.test_episode = my_test_episode

排查思路 #

• 确认导入顺序优先级：检查是否有其他代码（比如auxiliar_func内部、其他依赖模块）在monkey模块之前就导入了tianshou.trainer.utils，如果有，原函数会被缓存，后续patch无法覆盖这些提前导入的引用。
• 验证patch是否实际执行：在monkey.py末尾添加print(tianshou.trainer.utils.test_episode)，运行后查看输出是否为自定义的my_test_episode函数对象，确认patch操作本身是否成功。
• 追踪test_episode的实际调用路径：查看BaseTrainer的test_step方法中，是否是直接引用tianshou.trainer.utils.test_episode，还是已经通过from tianshou.trainer.utils import test_episode将函数绑定到类内部，后者会导致patch模块中的函数无法影响已绑定的引用。
• 检查函数引用缓存：如果其他模块（比如OffpolicyTrainer所在模块）提前导入了test_episode，它们会持有原函数的副本，不会使用你patch后的函数。

解决建议 #

直接patch目标方法（推荐） #

既然核心需求是修改BaseTrainer的test_step逻辑，直接针对类方法进行patch，绕开test_episode的引用问题：

from tianshou.trainer.utils import BaseTrainer

# 保存原方法引用
original_test_step = BaseTrainer.test_step

def my_test_step(self, *args, **kwargs):
    # 插入自定义逻辑，比如执行原方法后添加打印
    result = original_test_step(self, *args, **kwargs)
    print("check comment in execution to know it works")
    return result

# 替换类的test_step方法
BaseTrainer.test_step = my_test_step

确保覆盖所有test_episode的引用 #

如果坚持要patchtest_episode函数，需确保所有调用它的地方都使用模块级别的引用：

1. 先确认BaseTrainer或OffpolicyTrainer内部是否直接导入了test_episode，如果是，需要同时替换类内绑定的该函数。
2. 在monkey.py中，patch后可以尝试重新加载相关模块（需谨慎使用，可能破坏现有状态）：

import importlib
import tianshou.trainer.utils
import tianshou.trainer.offpolicy  # 假设OffpolicyTrainer在这里

# 先patch test_episode
def my_test_episode(*args, **kwargs):
    result = tianshou.trainer.utils.test_episode(*args, **kwargs)
    print("check comment in execution to know it works")
    return result
tianshou.trainer.utils.test_episode = my_test_episode

# 重新加载使用test_episode的模块
importlib.reload(tianshou.trainer.utils)
importlib.reload(tianshou.trainer.offpolicy)

验证patch生效状态 #

在main_func中添加验证代码，确认patch后的函数是否被正确引用：

def main_func():
    import tianshou.trainer.utils
    print("当前test_episode函数对象:", tianshou.trainer.utils.test_episode)
    # 后续原有逻辑

内容的提问来源于stack exchange，提问作者Aurelie Navir