[Feature Request] Natively support training and inference on different PyTorch devices

[ogiarch]

🚀 Feature

Natively allow SB3 models to be set up with two different PyTorch devices: one "inference device" used for methods like predict(...) and the other "training device" used for methods like learn(...). When learn(...) is invoked, e.g., move the model to the "training device" and run training. When predict(...) is invoked, move the model back to the inference device for inference.

Motivation

Disclaimer: I am a professional software engineer, but am completely self-taught when it comes to RL, so my general coding ability is a little better than my deep understanding of these algorithms. Take everything I say with appropriate grains of salt.

For about a year now in my spare time, I have been working on a project using a custom extension to the sb3_contrib Maskable PPO class for a personal board game development project, and have been working on trying to speed up long-running simulations on my macOS laptop. (Note: in case my use case affects my results here, it probably bears mentioning that because I am using it for MARL, with multiple models sharing the same environment, I am abusing the algorithm a little bit by limiting the number of parallel envs to 1. I hope to change this in the future, but it's going to be a heavy lift and will get into some gnarly multiprocessing stuff that is going to be a real slog to build out in Python.)

For a long time I shied away from trying to run anything on my GPU, because the official SB3 documentation on PPO states very clearly that "PPO is meant to be run primarily on the CPU, especially when you are not using a CNN." I figured that the reason for this guidance was that running inference on GPUs tends to dramatically worsen performance due to the fact that 1) inference needs to be run in every single training loop and interfaces heavily with a model's environment, 2) for the time being, most Gymnasium environments and simulations are very CPU-bound, and 3) roundtrips between CPU and GPU memory are way too slow to be running in every simulation loop. Experimentally, I found this to be the case when using py-spy to inspect runtime, with the vast majority of CPU time spent on shipping experiences to and especially collecting selected actions from the GPU when the GPU is used for inference.

However, we don't need to use the same PyTorch device for training and inference in our models, right? Training is particularly slow on CPUs and particularly fast on many GPUs, and because it runs relatively infrequently, I figured it might be a good idea to temporarily ship my whole model to the GPU before a training round and then back to the CPU before inference starts up right afterward. Other than a nasty bug causing model optimizer classes to be loaded on the wrong device, it was relatively straightforward to overload a few methods to build a model class that took separate arguments for an inference_device and a training_device and did what I wanted here.

The results were an absolutely massive speed-up in computation. Training rounds went from taking two seconds to about a half-second on average when shifted to 'mps', and this was particularly true when running multiple simulations at once. Parallel "tournament rounds" of thousands of games I'd set up went from taking 2.5-3 hrs to about 70 minutes to complete. The improvement in efficiency was to the point where I thought that some version of this functionality should be broadly available by default in SB3. If the results remain statistically sound, it seems like by default we're leaving a ton of performance on the table for most computers capable of running PPO and other related algorithms.

Curious to hear others' thoughts, though, especially as I'm new to the space.

Alternatives

1. Given that for many pure RL use cases, it's very often a bad idea to run inference on anything but the CPU, another option could be to just use the inputted device for training and always use the CPU for inference. I don't like this, though, as it would mean a regression in SB3's device flexibility today.
2. We could just... not do any of this. As I've said, though, unless I'm missing something really big when it comes to statistical accuracy when switching off between different devices like this, I think this would mean a missed opportunity for an enormous increase in model performance for many developers.

Checklist

• I have checked that there is no similar issue in the repo
• If I'm requesting a new feature, I have proposed alternatives

[araffin]

Hello,
to understand your problem, you are alternating (maybe with a callback):

model.learn(...)

evaluate(model, eval_env)

model.learn(...)

?

An alternative is to use something similar to what I did in #715 (code: https://github.com/DLR-RM/rl-baselines3-zoo/blob/87001ed8a40f817d46c950e283d1ca29e405ad71/utils/callbacks.py#L95), so literally having two models, one for training, one for inference.

[ogiarch]

Because of the MARL-like setup I have going here in a turn-based strategy setting, for me there's one central game loop with a shared "objective" environment which effectively fills agents' buffers one step at a time with experiences translated on the fly into their own private Gymnasium observations.

That centralization of the simulation code rather than the model code is the biggest difference between most popular RL environments and my use case. Different models don't have the same perspectives on their environments or see things in the same way as each other, so for me the "objective" environment has to live in an external/neutral place. Agents (classes where I've bundled models with their private environments) take actions and fill their rollout buffers in a game based on their current policy, and once those buffers are full, they launch a training round which updates that policy. This training can happen at any point during any game.

Heavy pseudocode for game loop:

while True:
    ...
    active_agent = get_active_agent(...)
    if active_agent.needs_training(...):
        active_agent.train(...)
    active_agent.reset_environment(...)
    action = active_agent.choose_action(...)
    new_state, reward = process_changes(action, ...)
    active_agent.take_step(new_state, reward, ...)
    if game_has_ended(...):
        break
    ...

The code moving the model back and forth between different devices that I've found to work very well currently lies within my train(...) method.

Current game loop execution time averages around 240 µs, including model action selection. It needs to be decently fast and therefore can't support device syncs that take over 100 µs to complete on my MacBook Pro in each loop – so even assuming optimal code, the idea of model inference on the GPU would be a complete non-starter for me unless my models were massive.

The callbacks you've suggested are a cool alternative approach, and something I definitely hadn't considered. I especially like the capacity for multithreading, although I did see in #715 that there was some trouble for some people getting models to behave properly with it. The big limitation I see here is that there is no callback event/hook that quite fits my use case yet. I need to train right before a model takes an action, not right afterward, because since this is a turn-based strategy game I often have to adjust agents' rewards for terminal actions (of choosing to end their turns, or requiring another agent to temporarily assume control to make a choice) retroactively to account for opponents' subsequent triumphs or failures. The idea is that it is probably a better idea to end your turn with the game in such a state that your opponents are unable to accrue high rewards on their subsequent turn, and I need that to be reflected in my reward signal.

Depending on precisely how this is all implemented, too, the multiple-model approach does give me some concern with regard to memory usage in my current setup. I have somewhat large models right now (by pet-project RL standards, not by, say, LLM standards), and there are a lot of weights to keep on hand at any given time, particularly given that in any moment I need many of them in memory at once.

I guess the resolution to this issue may well depend on the ease of defining a new type of callback hook tailored to my specific non-standard one-experience-at-a-time/MARL use case vs the ease of shifting actor and critic networks between the CPU and GPU. In my case, calling e.g. model.policy.to(device='mps') before training and ...to(device='cpu') afterward ended up being a much less intensive refactor. I think this might be even easier to reason about in the domain of single-player games and simulations like CartPole, too: learn could just collect experiences via evaluating the policy on the "inference device" where the code paths are really hot, then make one call to pass the model weights over to the "training device" with Torch before training starts and one call afterward to pass things back to the "inference device." It'd be pointless for simpler tasks or short rollouts, but for complex situations like mine, the ability to ship models to and from different devices easily has seriously led to my simulations running in half the time they used to take. I haven't noticed any detectable dip in quality/accuracy, with the caveat that measuring model quality in MARL can be hard to do objectively.

I have already hacked together a solution that works for me locally (albeit one that has to totally replace a bunch of SB3/SB3-contrib methods and classes to run fast) so I'm not married to any particular solution here. If you think that the callbacks are good enough as a workaround for the average use case, definitely feel free to close. Just wanted to float this idea of using one device for inference and another for training, both as a suggestion and as a potential opportunity to speed up some others' SB3 workflows.

(Edit: I'll add that even if this bug is closed, I do think that the guidance in the SB3 docs could be a little more nuanced on the feasibility of running PPO on devices other than the CPU. Yes, inference/simulation should probably always be handled on the CPU except for very rare occasions, but training can sometimes really benefit from being run on other devices.)