Port TRPO, PPO, GAIL

scratch/etienne/trpo/options/test_options.py

@@ -0,0 +1,54 @@
+from intersim.envs import IntersimpleLidarFlat
+from options import OptionsEnv
+import gym
+import numpy as np
+
+def test_obs_shape():
+    options = [(0, 5), (5, 5), (10, 5)]
+    env = OptionsEnv(IntersimpleLidarFlat(n_rays=5), options)
+    assert env.reset().shape == (36,)
+
+def test_act_space():
+    options = [(0, 5), (5, 5), (10, 5)]
+    env = OptionsEnv(IntersimpleLidarFlat(n_rays=5), options)
+    assert env.action_space == gym.spaces.Discrete(3)
+
+def test_plan():
+    options = [(0, 5), (5, 5), (10, 5)]
+    env = OptionsEnv(IntersimpleLidarFlat(n_rays=5), options)
+    env.reset()
+    plan = env.plan(options[0])
+    assert np.allclose(plan, -13.998268127441406 * np.ones((5,)))
+
+def test_plan2():
+    options = [(0, 5), (5, 5), (10, 5)]
+    env = OptionsEnv(IntersimpleLidarFlat(n_rays=5), options)
+    obs = env.reset()
+    states, actions, rewards, dones, plan_done, infos, n_steps = env.execute_plan(obs, options[0])
+    assert states.shape == (6, 36)
+    assert rewards.shape == (6,)
+    assert dones.shape == (6,)
+    assert len(infos) == 5
+
+def test_step():
+    options = [(0, 5), (5, 5), (10, 5)]
+    env = OptionsEnv(IntersimpleLidarFlat(n_rays=5), options)
+    env.reset()
+    obs, reward, done, _ = env.step(0)
+    assert obs.shape == (36,)
+    assert reward == 5.0
+    assert done == False
+
+def test_ll_step():
+    options = [(0, 5), (5, 5), (10, 5)]
+    env = OptionsEnv(IntersimpleLidarFlat(n_rays=5), options)
+    env.reset()
+    _, _, _, info = env.step(0)
+    assert info['ll']['observations'].shape == (6, 36)
+    assert info['ll']['actions'].shape == (6, 1)
+    assert info['ll']['rewards'].shape == (6,)
+    assert info['ll']['env_done'].shape == (6,)
+    assert info['ll']['plan_done'].shape == (6,)
+    assert info['ll']['plan_done'][5] == True
+    assert info['ll']['steps'] == 5
+    assert len(info['ll']['infos']) == 5