Upgrade gym to gymnasium

lerobot/common/envs/aloha/env.py

@@ -31,7 +31,7 @@ from lerobot.common.envs.aloha.tasks.sim_end_effector import (
 from lerobot.common.envs.aloha.utils import sample_box_pose, sample_insertion_pose
 from lerobot.common.utils import set_seed
 
-_has_gym = importlib.util.find_spec("gym") is not None
+_has_gym = importlib.util.find_spec("gymnasium") is not None
 
 
 class AlohaEnv(AbstractEnv):

lerobot/common/envs/pusht/env.py

@@ -18,7 +18,7 @@ from torchrl.envs.libs.gym import _gym_to_torchrl_spec_transform
 from lerobot.common.envs.abstract import AbstractEnv
 from lerobot.common.utils import set_seed
 
-_has_gym = importlib.util.find_spec("gym") is not None
+_has_gym = importlib.util.find_spec("gymnasium") is not None
 
 
 class PushtEnv(AbstractEnv):

lerobot/common/envs/pusht/pusht_env.py

@@ -1,14 +1,14 @@
 import collections
 
 import cv2
-import gym
+import gymnasium as gym
 import numpy as np
 import pygame
 import pymunk
 import pymunk.pygame_util
 import shapely.geometry as sg
 import skimage.transform as st
-from gym import spaces
+from gymnasium import spaces
 from pymunk.vec2d import Vec2d
 
 from lerobot.common.envs.pusht.pymunk_override import DrawOptions

lerobot/common/envs/pusht/pusht_image_env.py

@@ -1,5 +1,5 @@
 import numpy as np
-from gym import spaces
+from gymnasium import spaces
 
 from lerobot.common.envs.pusht.pusht_env import PushTEnv
 

lerobot/common/envs/simxarm/env.py

@@ -20,7 +20,6 @@ from lerobot.common.utils import set_seed
 MAX_NUM_ACTIONS = 4
 
 _has_gym = importlib.util.find_spec("gymnasium") is not None
-# _has_simxarm = importlib.util.find_spec("simxarm") is not None and _has_gym
 
 
 class SimxarmEnv(AbstractEnv):

lerobot/common/envs/simxarm/simxarm/__init__.py

@@ -1,8 +1,8 @@
 from collections import OrderedDict, deque
 
-import gym
+import gymnasium as gym
 import numpy as np
-from gym.wrappers import TimeLimit
+from gymnasium.wrappers import TimeLimit
 
 from lerobot.common.envs.simxarm.simxarm.task.base import Base as Base
 from lerobot.common.envs.simxarm.simxarm.task.lift import Lift