import matplotlib
matplotlib.use("MacOSX")
import metaworld
import random
import matplotlib.pyplot as plt


mt50 = metaworld.MT50() # Construct the benchmark, sampling tasks

training_envs = []
for name, env_cls in mt50.train_classes.items():
  # env = env_cls(render_mode="rgb_array")
  env = env_cls(render_mode="human")
  # env = env_cls()
  task = random.choice([task for task in mt50.train_tasks
                        if task.env_name == name])
  env.set_task(task)
  env.camera_name="corner3"
  training_envs.append(env)

# 1. works (image 1)
plt.imshow(training_envs[34].render()) # pick-place-v2
plt.show()

# 2. works (image 2)
plt.imshow(training_envs[1].render()) # basketball-v2
plt.show()

# 3. does not work anymore (image 3)
plt.imshow(training_envs[34].render()) # pick-place-v2
plt.show()