Stable-Baselines3 (SB3) Improvements

[araffin]

Hello,

I would like to discuss and propose some improvements for SB3 integration with isaac lab (as a follow up of #528).

I've been playing around IsaacLab with SB3/SBX (SBX = SB3 + jax) and I have some improvements to suggests/questions:

• with respect to the speed, most of the slow down comes from the adapter (SB3 wrapper), and by only treating envs that are terminated, I managed to improve speed by 3X (my fork is here, trying out things: https://github.com/araffin/IsaacLab/tree/feat/sb3-update), so if you think it is interesting, I can contribute the fast variant (https://github.com/araffin/IsaacLab/blob/feat/sb3-update/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/utils/wrappers/sb3.py)
• why not use python script to define hyperparameter config for sb3 too? this would be consistent with other frameworks and would allow to support multiple algorithms?
• why not have default hyperparameters for each algo? (currently the process of adding support for a new task is a bit annoying and require to dig deep in the folder structure :/)
• would you be interested in tuned hyperparameters for more isaac lab tasks?

To reproduce, using my fork (https://github.com/araffin/IsaacLab, branch /feat/sb3-update):

Default, runs at 10 000 steps/s:

./isaaclab.sh -p source/standalone/workflows/sb3/train.py --task Isaac-Velocity-Flat-Unitree-A1-v0 --num_envs 1024 --algo ppo

Fast variant, runs at 35 000 steps/s (SBX) and 30 000 steps/s:

./isaaclab.sh -p source/standalone/workflows/sb3/train.py --task Isaac-Velocity-Flat-Unitree-A1-v0 --num_envs 1024 --algo ppo --fast

fastest variant but incorrect (truncation not handled, lower performance): --no--info flag, runs at 40 000 steps/s (similar to rsl_rl)

Side remark: while playing around, the rsl_rl ppo hyperparameters don't seem to be tuned for the A1 task/ I managed to get much higher rewards with SB3/SBX implementation (mean reward ~30 vs ~10 for rsl_rl, probably because they don't handle truncation?)

[RandomOakForest]

Thanks for posting this. I will move this to our issues and tag it as a proposal of enhancements for the team to consider.