[Bug]: Manually setting net_arch=None causes crash when loading model

[jak3122]

🐛 Bug

I've run into this issue a few times, since I use command line args in a wrapper script where I default net_arch to None, and I pass that to policy_kwargs.

I was able to fix it locally by changing this line from

if "net_arch" in data["policy_kwargs"] and len(data["policy_kwargs"]["net_arch"]) > 0:

to

if data.get("policy_kwargs", {}).get("net_arch"):

To Reproduce

import gymnasium as gym

from stable_baselines3 import PPO

env = gym.make("CartPole-v1")

model = PPO("MlpPolicy", env, policy_kwargs=dict(net_arch=None))
model.learn(total_timesteps=5000)
model.save("ppo_cartpole")

del model

model = PPO.load("ppo_cartpole")

Relevant log output / Error message

Traceback (most recent call last):
  File "/Users/hyzer/stable-baselines3/min_example.py", line 13, in <module>
    model = PPO.load("ppo_cartpole")
  File "/Users/hyzer/stable-baselines3/stable_baselines3/common/base_class.py", line 695, in load
    if "net_arch" in data["policy_kwargs"] and len(data["policy_kwargs"]["net_arch"]) > 0:
TypeError: object of type 'NoneType' has no len()

System Info

• OS: macOS-14.4.1-arm64-arm-64bit Darwin Kernel Version 23.4.0: Fri Mar 15 00:19:22 PDT 2024; root:xnu-10063.101.17~1/RELEASE_ARM64_T8112
• Python: 3.10.14
• Stable-Baselines3: 2.4.0a1
• PyTorch: 2.2.2
• GPU Enabled: False
• Numpy: 1.26.4
• Cloudpickle: 3.0.0
• Gymnasium: 0.29.1

Checklist

• [X] My issue does not relate to a custom gym environment. (Use the custom gym env template instead)
• [X] I have checked that there is no similar issue in the repo
• [X] I have read the documentation
• [X] I have provided a minimal and working example to reproduce the bug
• [X] I've used the markdown code blocks for both code and stack traces.

[araffin]

Hello, why do you want to set net_arch=None?

[jak3122]

In my training script I have CLI args, including an optional net_arch arg:

net_arch: Tuple[int, ...] | None = None

Which I then pass to the sb3 model:

model = PPO(
    policy_type,
    env,
    policy_kwargs=dict(
        net_arch=args.net_arch,
        ...
    )
)

Since the documentation says that net_arch can be None, I would expect this to work, and have None indicate the default net_arch, which it does, except for loading the saved model.

I also realize it's possible to just omit net_arch from policy_kwargs instead, like this:

policy_kwargs = dict()
if args.net_arch is not None:
    policy_kwargs["net_arch"] = args.net_arch
model = PPO(
    policy_type,
    env,
    policy_kwargs=policy_kwargs
)

But then for anyone like me who happens to save a model with net_arch manually set to None, then the model will not load.

[araffin]

Since the documentation says that net_arch can be None, I would expect this to work, and have None indicate the default net_arch, which it does, except for loading the saved mode

I would be happy to receive a PR that solves this issue =)

Although in your case, I would indeed recommend not populating the net_arch instead of passing None.