Merge pull request #1 from WLS2002/full-jit

Strange bug! Add {"new_step_api": True} in gym environments.

evox_adaptor/gym_no_distribution.py

@@ -55,9 +55,7 @@ class Gym(Problem):
         return -fitnesses, State(key=key)
 
     def __rollout(self, seeds, pop):
-        observations, infos = zip(
-            *[env.reset(seed=seed) for env, seed in zip(self.envs, seeds)]
-        )
+        observations = [env.reset(seed=seed) for env, seed in zip(self.envs, seeds)]
         terminates, truncates = np.zeros((2, self.pop_size), dtype=bool)
         fitnesses, rewards = np.zeros((2, self.pop_size))
 

examples/evox_/acrobot.py

@@ -39,6 +39,7 @@ if __name__ == '__main__':
     problem = Gym(
         policy=jit(vmap(neat_forward)),
         env_name="Acrobot-v1",
+        env_options={"new_step_api": True},
         pop_size=100,
     )
 

examples/evox_/cartpole.py

@@ -39,6 +39,7 @@ if __name__ == '__main__':
     problem = Gym(
         policy=jit(vmap(neat_forward)),
         env_name="CartPole-v1",
+        env_options={"new_step_api": True},
         pop_size=40,
     )
 

examples/evox_/gym_env_test.py

@@ -0,0 +1,14 @@
+import gym
+
+env = gym.make("CartPole-v1", new_step_api=True)
+print(env.reset())
+obs = env.reset()
+
+print(obs)
+while True:
+    action = env.action_space.sample()
+    obs, reward, terminate, truncate, info = env.step(action)
+    print(obs, info)
+    if terminate | truncate:
+        break
+

examples/evox_/mountain_car.py

@@ -39,6 +39,7 @@ if __name__ == '__main__':
     problem = Gym(
         policy=jit(vmap(neat_forward)),
         env_name="MountainCarContinuous-v0",
+        env_options={"new_step_api": True},
         pop_size=100,
     )