60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
import pickle
|
|
import imitation.data.rollout as rollout
|
|
from options import LLOptions, OptionsEnv
|
|
from intersim.envs import NRasterized
|
|
import itertools
|
|
import stable_baselines3
|
|
from policy import OptionsCnnPolicy
|
|
from train import flatten_transitions
|
|
import numpy as np
|
|
|
|
def test_ll_expert_data():
|
|
with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
|
|
expert_trajectories = pickle.load(f)
|
|
expert_transitions = rollout.flatten_trajectories(expert_trajectories)
|
|
|
|
env = LLOptions(NRasterized(agent=51, width=36, height=36, m_per_px=2))
|
|
|
|
gen_transitions = list(itertools.islice(env.sample_ll(
|
|
policy=stable_baselines3.PPO(
|
|
OptionsCnnPolicy,
|
|
OptionsEnv(env),
|
|
verbose=1,
|
|
)
|
|
), 10))
|
|
gen_transitions = flatten_transitions(gen_transitions)
|
|
|
|
assert expert_transitions[:10].obs.shape == gen_transitions['obs'].shape
|
|
assert expert_transitions[:10].next_obs.shape == gen_transitions['next_obs'].shape
|
|
assert expert_transitions[:10].acts.shape == gen_transitions['acts'].shape
|
|
assert expert_transitions[:10].dones.shape == gen_transitions['dones'].shape
|
|
|
|
def test_ll_states():
|
|
env = NRasterized()
|
|
policy = stable_baselines3.PPO(
|
|
OptionsCnnPolicy,
|
|
OptionsEnv(env),
|
|
verbose=1,
|
|
)
|
|
llenv = LLOptions(env)
|
|
transitions = list(itertools.islice(llenv.sample_ll(policy=policy), 100))
|
|
|
|
env2 = NRasterized()
|
|
s2 = env2.reset()
|
|
for i, t in enumerate(transitions):
|
|
assert i == 0 or np.array_equal(t['obs'], transitions[i-1]['next_obs'])
|
|
assert np.array_equal(t['obs'], s2)
|
|
assert t['acts'].shape == (1,)
|
|
|
|
nexts2, _, done2, _ = env2.step(t['acts'])
|
|
assert np.array_equal(t['next_obs'], nexts2)
|
|
assert np.array_equal(t['dones'], done2)
|
|
|
|
if done2:
|
|
break
|
|
|
|
s2 = nexts2
|
|
|
|
def test_hl_transitions():
|
|
pass
|