Gymnasium migration

[Rohan138]

Types of changes

• [X] Docs change / refactoring / dependency upgrade
• [ ] Bug fix (non-breaking change which fixes an issue)
• [ ] New feature (non-breaking change which adds functionality)
• [ ] Breaking change (fix or feature that would cause existing functionality to change)

Motivation and Context / Related issue

Migrating mbrl-lib from gym=0.17 all the way to a combination of gym=0.26.2 and gymnasium=0.26.3.

How Has This Been Tested (if it applies)

Will need thorough testing to ensure I didn't break anything, and also to make sure that performance is retained after refactoring the step API to use terminated and truncated.

Checklist

• [X] The documentation is up-to-date with the changes I made.
• [X] I have read the CONTRIBUTING document and completed the CLA (see CONTRIBUTING).
• [x] All tests passed, and additional code has been covered with new tests.

[Rohan138]

Update: Most files have been changed and fixed, tests with tests/core, tests/pybullet, tests/mujoco, tests/dmcontrol. Still fixing tests/algorithms.

[Rohan138]

@luisenp @natolambert this should be ready for review.

• Note: I decided to add the --follow-imports=skip flag to mypy for now, typing the whole codebase correctly with all the gymnasium and gym wrappers around dmcontrol and pybullet envs is too large a refactor, hope that's okay?
• @raghavauppuluri13 is currently running learning experiments with this and the current main branch to verify that the migration does not affect performance; we'll put up results in a day or two.

[raghavauppuluri13]

Benchmarking the PR against main. Here are the reward curves for some select envs. The planet algorithms are still training, but you can see they are generally similar. - EDIT: All are finished training

[natolambert]

@raghavauppuluri13 lmk if you're compute limited, shouldn't be too hard for me to run some too.

[raghavauppuluri13]

@raghavauppuluri13 lmk if you're compute limited, shouldn't be too hard for me to run some too.

Do we want to benchmark all the env/algo combos in the examples? Or is this good enough? If so, that would be great. I've got a helper script.

#!/bin/bash

main_envs=("pets_reacher" "pets_pusher" "planet_cartpole_swingup" "planet_finger_spin" "mbpo_halfcheetah" "mbpo_inv_pendulum")
diff_envs=("pets_reacher" "pets_pusher" "planet_cartpole_swingup" "planet_finger_spin" "mbpo_half_cheetah_v4" "mbpo_inv_pendulum_v4")
algos=("pets" "pets" "planet" "planet" "mbpo" "mbpo")
device=("cuda:0" "cuda:0" "cuda:0" "cuda:0" "cuda:0" "cuda:0")
envs=$main_envs
tsp -S 6

for i in ${!envs[@]}; do
    if [[ "${algos[$i]}" == "planet" ]]; then
        tsp python -m mbrl.examples.main algorithm=${algos[$i]} dynamics_model=${algos[$i]} overrides=${envs[$i]} device=${device[$i]}
    else
        tsp python -m mbrl.examples.main algorithm=${algos[$i]} overrides=${envs[$i]} device=${device[$i]}
    fi
done

[raghavauppuluri13]

@natolambert @luisenp Any other validation steps before the pr is ready to merge? (I updated the benchmarks with the finished training curves)

[natolambert]

Awesome work all, have been following along happily.

[Rohan138]

@luisenp Could I get another review? Should be good to merge now, let me know if there's anything else I can change