imitation/tests/algorithms/conftest.py at a55ff9e823d1604f7af3e75102f37bc4d1827b96 · HumanCompatibleAI/imitation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
"""Fixtures common across algorithm tests."""
from typing import Sequence

import pytest
from stable_baselines3.common import envs
from stable_baselines3.common.policies import BasePolicy
from stable_baselines3.common.vec_env import DummyVecEnv, VecEnv

from imitation.algorithms import bc
from imitation.data.types import TrajectoryWithRew
from imitation.data.wrappers import RolloutInfoWrapper
from imitation.policies import serialize
from imitation.testing.expert_trajectories import (
    lazy_generate_expert_trajectories,
    make_expert_transition_loader,
)
from imitation.util import util

CARTPOLE_ENV_NAME = "seals/CartPole-v0"


@pytest.fixture
def cartpole_expert_policy(cartpole_venv: VecEnv) -> BasePolicy:
    return serialize.load_policy(
        "ppo-huggingface",
        cartpole_venv,
        env_name=CARTPOLE_ENV_NAME,
    )


@pytest.fixture
def cartpole_expert_trajectories(
    cartpole_expert_policy,
    cartpole_venv,
    pytestconfig,
    rng,
) -> Sequence[TrajectoryWithRew]:
    return lazy_generate_expert_trajectories(
        pytestconfig.cache.makedir("experts"),
        CARTPOLE_ENV_NAME,
        60,
        rng,
    )


PENDULUM_ENV_NAME = "Pendulum-v1"


@pytest.fixture
def cartpole_bc_trainer(
    pytestconfig,
    cartpole_venv,
    cartpole_expert_trajectories,
    rng,
):
    return bc.BC(
        observation_space=cartpole_venv.observation_space,
        action_space=cartpole_venv.action_space,
        batch_size=50,
        demonstrations=make_expert_transition_loader(
            cache_dir=pytestconfig.cache.makedir("experts"),
            batch_size=50,
            expert_data_type="transitions",
            env_name="seals/CartPole-v0",
            rng=rng,
            num_trajectories=60,
        ),
        custom_logger=None,
        rng=rng,
    )


@pytest.fixture
def pendulum_expert_trajectories(
    pytestconfig,
    rng,
) -> Sequence[TrajectoryWithRew]:
    return lazy_generate_expert_trajectories(
        pytestconfig.cache.makedir("experts"),
        PENDULUM_ENV_NAME,
        60,
        rng,
    )


@pytest.fixture
def pendulum_expert_policy(pendulum_venv) -> BasePolicy:
    return serialize.load_policy(
        "ppo-huggingface",
        pendulum_venv,
        env_name=PENDULUM_ENV_NAME,
    )


@pytest.fixture
def pendulum_venv(rng) -> VecEnv:
    return util.make_vec_env(
        PENDULUM_ENV_NAME,
        n_envs=8,
        post_wrappers=[lambda env, _: RolloutInfoWrapper(env)],
        rng=rng,
    )


@pytest.fixture
def pendulum_single_venv(rng) -> VecEnv:
    return util.make_vec_env(
        PENDULUM_ENV_NAME,
        n_envs=1,
        post_wrappers=[lambda env, _: RolloutInfoWrapper(env)],
        rng=rng,
    )


@pytest.fixture
def multi_obs_venv() -> VecEnv:
    def make_env():
        env = envs.SimpleMultiObsEnv(channel_last=False)
        return RolloutInfoWrapper(env)

    return DummyVecEnv([make_env, make_env])