[Question] Significant Performance Disparity Between Maskable PPO and PPO

[gemelom]

❓ Question

I tried to run PPO and Maskable PPO on my custom environment with the same configuration, but I found that Maskable PPO(~5fps) is mush slower than PPO(~140fps).

Here's my configurations:

• environment setup

  env.action_space = MultiBinary(339)

• reproduction code

  config = {
      "env_name": "my_env_name",
      "n_envs": 16,
      "policy_type": "MlpPolicy",
      "total_timesteps": 100000,
  }
  
  # DummyVecEnv
  vec_env = make_vec_env(config["env_name"], n_envs=config["n_envs"])
  
  model = MaskablePPO("MlpPolicy", vec_env, n_steps=128, verbose=1")
  # model = PPO("MlpPolicy", vec_env, n_steps=128, verbose=1")
  
  model.learn(
      total_timesteps=config["total_timesteps"],
      callback=WandbCallback(
          gradient_save_freq=100,
          model_save_path=f"models/{experiment_name}",
          verbose=2,
      ),
      progress_bar=True,
  )

I also tried to profile my code with py-spy, and I found that MaskablePPO spent many extra time in these lines

while PPO spends much less time in train and most of its time in collect_rollouts just as expected.

I wonder if this extreme decline in training efficiency is a normal situation because of the large action_space or if there are other bugs in the implementation.

Checklist

• I have checked that there is no similar issue in the repo
• I have read the documentation
• If code there is, it is minimal and working
• If code there is, it is formatted using the markdown code blocks for both code and stack traces.

[araffin]

Hello,

I also tried to profile my code with py-spy, and I found that MaskablePPO spent many extra time in these lines

at least the slow down is where it would be expected.
I'm a bit surprised by how much slow down it, but the code was never optimized for speed, so there is probably room for improvement.

Related code:

stable-baselines3-contrib/sb3_contrib/common/maskable/distributions.py

Lines 245 to 252 in 55a0332

	masks_tensor = th.as_tensor(masks)
	# Restructure shape to align with logits
	masks_tensor = masks_tensor.view(-1, sum(self.action_dims))
	# Then split columnwise for each discrete action
	split_masks = th.split(masks_tensor, list(self.action_dims), dim=1) # type: ignore[assignment]
	for distribution, mask in zip(self.distributions, split_masks):
	distribution.apply_masking(mask)

and

stable-baselines3-contrib/sb3_contrib/common/maskable/distributions.py

Lines 58 to 62 in 55a0332

	device = self.logits.device
	self.masks = th.as_tensor(masks, dtype=th.bool, device=device).reshape(self.logits.shape)
	HUGE_NEG = th.tensor(-1e8, dtype=self.logits.dtype, device=device)
	logits = th.where(self.masks, self._original_logits, HUGE_NEG)

[ogiarch]

Yeah – I've been bitten by this issue too. Out of my entire personal simulation + model training and prediction, 20 percent of my runtime is being eaten by this step alone. There is no good reason for things to be this slow – we are sampling from a probability distribution where the logits are all known because we've already done a forward pass through the model.

For at least the MaskableCategorical distribution, also, this problem is also stacked on top of shockingly inefficient code in pytorch that has sat unchanged for years (see pytorch/pytorch#30968) with seemingly no one caring to do anything about it. The whole sampling process needs rewriting from the ground up.

[ogiarch]

Here's an example quick-and-dirty workaround for a custom Maskable PPO variant to be used in the sample(...) method of a custom distribution class that does not initialize a new object every time you want to sample actions. This code has been informed by multiple sources, including pytorch/pytorch#30968 (comment), pytorch/pytorch#16706, and #81. It's not optimal, but I have found it to be adequately accurate for my use case while enabling a big performance boost in model inference.

cum_sum = self._prbs.cumsum(dim=-1)
# Address https://github.com/Stable-Baselines-Team/stable-baselines3-contrib/issues/81
cum_sum[..., -1] = 1.

# Circumvent https://github.com/pytorch/pytorch/issues/16706
rand_num = rand(self._prbs.shape[:-1], device=self._logits.device) \
    .clamp(min=1e-24)

sampled_actions = searchsorted(cum_sum, rand_num[..., None])

[araffin]

@ogiarch would you mind doing a PR that fixes this issue? (maybe we could keep the old behavior commented out or in a branch for reference)

[araffin]

20 percent of my runtime is being eaten by this step alone.

btw, do you know which lines exactly are slow?
and could you give a minimal env example to reproduce the slowness?

We should consider integrating the fixes from @Sarthak-Dayal (PR #302 ) if possible too.

[ogiarch]

I'm afraid my environment/repo is set up in such a different way from most, and I'm tied up with enough at my day job at the moment, that it might take some time for me to repro in a minimal env. I can remind myself to revisit this issue once I'm a little freer if you would prefer that.

In the meantime, I can shed some more light on lines within sb3-contrib I identified as particularly slow:

stable-baselines3-contrib/sb3_contrib/common/maskable/distributions.py

Line 137 in e03aa1e

self.distribution = MaskableCategorical(logits=reshaped_logits)

stable-baselines3-contrib/sb3_contrib/common/maskable/distributions.py

Line 163 in e03aa1e

self.distribution.apply_masking(masks)

stable-baselines3-contrib/sb3_contrib/common/maskable/distributions.py

Line 68 in e03aa1e

super().__init__(logits=logits)

stable-baselines3-contrib/sb3_contrib/common/maskable/distributions.py

Line 147 in e03aa1e

return self.distribution.sample()

Thoughts:

• To begin with, PyTorch's distribution classes are often quite poorly optimized. By default they introduce slow constraint checks into what can be very hot code paths, and on the backend they opt for slow matrix ops like multinomial when much faster ones such as searchsorted (as seen above, the central insight for which I am not smart enough to be able to take credit) can be used. I'd personally recommend ditching their concrete implementations entirely wherever possible, rather than wrapping them or inheriting from them, and just using the Distribution/MaskableDistribution abstract base class instead.
• On top of that, sb3_contrib code is creating a ton of unnecessary objects. _get_action_dist_from_critic in MaskableActorCriticPolicy is called in every forward pass as well as every get_distribution and evaluate_actions call, and it calls MaskableCategoricalDistribution.proba_distribution. Despite the fact that it can very handily reuse one or several internal MaskableCategorical objects, proba_distribution creates a new one for MaskableCategoricalDistribution every single time it is called.
• When a new MaskableCategorical object is created, apply_masking, a pretty slow function in concept even if written perfectly due to the need to normalize probabilities after masking is applied, is immediately called. We then always call apply_masking on that same object again one more time in our policy's forward method. I don't believe this needs to happen twice.
• What's more is that every time MaskableCategorical.apply_masking is called, rather than having it simply normalize its own internal logits/probabilities after the mask is applied, we call super().__init__ to force the object to also call all the other bloated Categorical PyTorch __init__ code all over again. This adds even more overhead to every forward pass, and unfortunately that overhead really adds up over time.

Locally I've found that when I fix all of these issues, I achieve a 5-6x speedup in the analogous code paths. This does require basically rewriting the entire distribution class and getting rid of as much PyTorch bloat as possible, though; otherwise the full benefits cannot be realized. That's why in my opinion a proper PR to put the code I shared above in place will take more time and care than may initially seem required. I have a solution working well locally, but it's pretty specific to my use case and I'd need to change and add a fair bit to meet SB3 standards. I would like to help out more here; it just may take me a while. Meanwhile I hope that my deep-dive into the existing code above helps at least a little.