test_types.py

"""Tests of `imitation.data.types`."""

import contextlib
import copy
import dataclasses
import os
import pathlib
import pickle
from typing import Any, Callable, Sequence

import gymnasium as gym
import numpy as np
import pytest

from imitation.data import serialize, types
from imitation.util import util


def _check_1d_shape(fn: Callable[[np.ndarray], Any], length: int, expected_msg: str):
    for shape in [(), (length, 1), (length, 2), (length - 1,), (length + 1,)]:
        with pytest.raises(ValueError, match=expected_msg):
            fn(np.zeros(shape))


@pytest.fixture
def transitions_min(
    obs_space: gym.Space,
    act_space: gym.Space,
    length: int,
) -> types.TransitionsMinimal:
    obs = np.array([obs_space.sample() for _ in range(length)])
    acts = np.array([act_space.sample() for _ in range(length)])
    infos = np.array([{i: i} for i in range(length)])
    return types.TransitionsMinimal(obs=obs, acts=acts, infos=infos)


@pytest.fixture
def transitions(
    transitions_min: types.TransitionsMinimal,
    obs_space: gym.Space,
    length: int,
) -> types.Transitions:
    """Fixture to generate transitions of length `length` iid sampled from spaces."""
    next_obs = np.array([obs_space.sample() for _ in range(length)])
    dones = np.zeros(length, dtype=bool)
    return types.Transitions(
        **types.dataclass_quick_asdict(transitions_min),
        next_obs=next_obs,
        dones=dones,
    )


@pytest.fixture
def transitions_rew(
    transitions: types.Transitions,
    length: int,
) -> types.TransitionsWithRew:
    """Like `transitions` but with reward randomly sampled from a Gaussian."""
    rews = np.random.randn(length)
    return types.TransitionsWithRew(
        **types.dataclass_quick_asdict(transitions),
        rews=rews,
    )


def _check_transitions_get_item(trans, key):
    """Check trans[key] by manually indexing/slicing into every `trans` field."""
    item = trans[key]
    for field in dataclasses.fields(trans):
        if isinstance(item, dict):
            observed = item[field.name]  # pytype: disable=unsupported-operands
        else:
            observed = getattr(item, field.name)

        expected = getattr(trans, field.name)[key]
        if isinstance(expected, np.ndarray):
            assert observed.dtype == expected.dtype  # pytype:disable=attribute-error
        np.testing.assert_array_equal(observed, expected)


@contextlib.contextmanager
def pushd(dir_path):
    """Change directory temporarily inside context."""
    orig_dir = pathlib.Path.cwd()
    try:
        os.chdir(dir_path)
        yield
    finally:
        os.chdir(orig_dir)


class TestData:
    """Tests of imitation.util.data.

    Grouped in a class since parametrized over common set of spaces.
    """

    def test_valid_trajectories(
        self,
        trajectory: types.Trajectory,
        trajectory_rew: types.TrajectoryWithRew,
        length: int,
    ) -> None:
        """Checks trajectories can be created for a variety of lengths and spaces."""
        trajs = [trajectory, trajectory_rew]
        trajs += [dataclasses.replace(traj, infos=None) for traj in trajs]
        for traj in trajs:
            assert len(traj) == length

    def test_traj_unequal_to_other_types(
        self,
        trajectory: types.Trajectory,
        trajectory_rew: types.TrajectoryWithRew,
    ) -> None:
        """Test trajectories unequal to objects of different types."""
        for t in [trajectory, trajectory_rew]:
            # Trajectory compare unequal to things that are not trajectories
            assert t != 42
            assert t != "foobar"

        # Trajectory compares unequal to a copy of itself but with reward
        assert trajectory != trajectory_rew

    def test_traj_equal_to_self_and_copies(
        self,
        trajectory: types.Trajectory,
        trajectory_rew: types.TrajectoryWithRew,
    ) -> None:
        """Test that trajectories are equal to themselves and copies."""
        for t in [trajectory, trajectory_rew]:
            # Equal to itself
            assert t == t
            # And to copy
            assert t == copy.copy(t)

    def test_traj_unequal_to_perturbations(
        self,
        trajectory: types.Trajectory,
        trajectory_rew: types.TrajectoryWithRew,
        length: int,
    ) -> None:
        """Test that trajectories unequal to perturbed versions."""
        # Unequal to a copy of itself truncated
        new_length = length - 1
        if new_length > 0:
            assert trajectory != types.Trajectory(
                obs=trajectory.obs[: new_length + 1],
                acts=trajectory.acts[:new_length],
                infos=trajectory.infos[:new_length]
                if trajectory.infos is not None
                else None,
                terminal=trajectory.terminal,
            )

        # Or with contents changed
        for t in [trajectory, trajectory_rew]:
            as_dict = types.dataclass_quick_asdict(t)
            for k in as_dict.keys():
                perturbed = dict(as_dict)
                if k == "infos":
                    perturbed["infos"] = [{"foo": 42}] * len(as_dict["infos"])
                elif isinstance(as_dict[k], types.DictObs):
                    perturbed[k] = as_dict[k].map_arrays(lambda x: x + 1)
                else:
                    perturbed[k] = as_dict[k] + 1
                assert t != type(t)(**perturbed)

    @pytest.mark.parametrize("type_safe", [False, True])
    @pytest.mark.parametrize("use_pickle", [False, True])
    @pytest.mark.parametrize("use_rewards", [False, True])
    @pytest.mark.parametrize("use_chdir", [False, True])
    def test_save_trajectories(
        self,
        trajectory: types.Trajectory,
        trajectory_rew: types.TrajectoryWithRew,
        use_chdir,
        tmpdir,
        use_pickle,
        use_rewards,
        type_safe,
    ):
        if isinstance(trajectory.obs, types.DictObs):
            pytest.xfail("Saving/loading dictobs trajectories not yet supported")

        chdir_context: contextlib.AbstractContextManager
        """Check that trajectories are properly saved."""
        if use_chdir:
            # Test no relative path without directory edge-case.
            chdir_context = pushd(tmpdir)
            save_dir_str = ""
        else:
            chdir_context = contextlib.nullcontext()
            save_dir_str = tmpdir

        with chdir_context:
            save_dir = util.parse_path(save_dir_str)
            trajs = [trajectory_rew if use_rewards else trajectory]
            save_path = save_dir / "trajs"

            if use_pickle:
                # Pickle format
                with open(save_path, "wb") as f:
                    pickle.dump(trajs, f)
            else:
                # HuggingFace Dataset Format
                serialize.save(save_path, trajs)

                # Test that heterogeneous lists of trajectories throw an error
                if use_rewards:
                    with pytest.raises(ValueError):
                        serialize.save(save_path, [trajectory, trajectory_rew])

            loaded_trajs: Sequence[types.Trajectory]
            if type_safe:
                if use_rewards:
                    loaded_trajs = serialize.load_with_rewards(save_path)
                else:
                    with pytest.raises(ValueError):
                        serialize.load_with_rewards(save_path)
                    loaded_trajs = serialize.load(save_path)
            else:
                loaded_trajs = serialize.load(save_path)

            assert len(trajs) == len(loaded_trajs)
            for t1, t2 in zip(trajs, loaded_trajs):
                assert t1 == t2

    def test_invalid_trajectories(
        self,
        trajectory: types.Trajectory,
        trajectory_rew: types.TrajectoryWithRew,
    ) -> None:
        """Checks input validation catches space and dtype related errors."""
        trajs = [trajectory, trajectory_rew]
        for traj in trajs:
            with pytest.raises(
                ValueError,
                match=r"expected one more observations than actions.*",
            ):
                dataclasses.replace(traj, obs=traj.obs[:-1])
            with pytest.raises(
                ValueError,
                match=r"expected one more observations than actions.*",
            ):
                dataclasses.replace(traj, acts=traj.acts[:-1])

            with pytest.raises(
                ValueError,
                match=r"infos when present must be present for each action.*",
            ):
                assert traj.infos is not None
                dataclasses.replace(traj, infos=traj.infos[:-1])
            with pytest.raises(
                ValueError,
                match=r"infos when present must be present for each action.*",
            ):
                dataclasses.replace(traj, obs=traj.obs[:-1], acts=traj.acts[:-1])

        _check_1d_shape(
            fn=lambda rews: dataclasses.replace(trajectory_rew, rews=rews),
            length=len(trajectory_rew),
            expected_msg=r"rewards must be 1D array.*",
        )

        with pytest.raises(ValueError, match=r"rewards dtype.* not a float"):
            dataclasses.replace(
                trajectory_rew,
                rews=np.zeros(len(trajectory_rew), dtype=int),
            )

    def test_valid_transitions(
        self,
        transitions_min: types.TransitionsMinimal,
        transitions: types.Transitions,
        transitions_rew: types.TransitionsWithRew,
        length: int,
        n_checks: int = 20,
    ) -> None:
        """Checks initialization, indexing, and slicing sanity."""
        for trans in [transitions_min, transitions, transitions_rew]:
            assert len(trans) == length

            for _ in range(n_checks):
                # Indexing checks, which require at least one element.
                if length != 0:
                    index = np.random.randint(length)
                    assert isinstance(trans[index], dict)
                    _check_transitions_get_item(trans, index)

                # Slicing checks.
                start = np.random.randint(-2, length)
                stop = np.random.randint(0, length + 2)
                step = np.random.randint(-2, 4)
                if step == 0:  # Illegal. Quick fix that biases tests to ordinary step.
                    step = 1
                s = slice(start, stop, step)
                assert type(trans[s]) is type(trans)
                _check_transitions_get_item(trans, s)

    def test_invalid_transitions(
        self,
        transitions_min: types.Transitions,
        transitions: types.Transitions,
        transitions_rew: types.TransitionsWithRew,
        length: int,
    ) -> None:
        """Checks input validation catches space and dtype related errors."""
        if length == 0:
            pytest.skip()

        for trans in [transitions_min, transitions, transitions_rew]:
            with pytest.raises(
                ValueError,
                match=r"obs and acts must have same number of timesteps:.*",
            ):
                dataclasses.replace(trans, acts=trans.acts[:-1])
            with pytest.raises(
                ValueError,
                match=r"obs and infos must have same number of timesteps:.*",
            ):
                dataclasses.replace(trans, infos=[{}] * (length - 1))

        for trans in [transitions, transitions_rew]:
            with pytest.raises(
                ValueError,
                match=r"obs and next_obs must have same shape:.*",
            ):
                dataclasses.replace(trans, next_obs=np.zeros((len(trans), 4, 2)))

            with pytest.raises(
                ValueError,
                match=r"obs and next_obs must have the same dtype:.*",
            ):
                dataclasses.replace(
                    trans,
                    next_obs=np.zeros_like(trans.next_obs, dtype=bool),
                )

            _check_1d_shape(
                fn=lambda bogus_dones: dataclasses.replace(trans, dones=bogus_dones),
                length=len(trans),
                expected_msg=r"dones must be 1D array.*",
            )

            with pytest.raises(ValueError, match=r"dones must be boolean"):
                dataclasses.replace(trans, dones=np.zeros(len(trans), dtype=int))

        _check_1d_shape(
            fn=lambda bogus_rews: dataclasses.replace(trans, rews=bogus_rews),
            length=len(transitions_rew),
            expected_msg=r"rewards must be 1D array.*",
        )

        with pytest.raises(ValueError, match=r"rewards dtype.* not a float"):
            dataclasses.replace(
                transitions_rew,
                rews=np.zeros(len(transitions_rew), dtype=int),
            )


def test_zero_length_fails():
    """Check zero-length trajectory and transitions fail."""
    empty = np.array([])
    with pytest.raises(ValueError, match=r"Degenerate trajectory.*"):
        types.Trajectory(obs=np.array([42]), acts=empty, infos=None, terminal=True)


def test_parse_path():
    if os.name == "nt":  # pragma: no cover
        pytest.skip(
            "Windows uses path.WindowsPath instead when paths are resolved, which"
            "cannot be compared directly to pathlib.Path objects.",
        )
    # absolute paths
    assert util.parse_path("/foo/bar") == pathlib.Path("/foo/bar")
    assert util.parse_path(pathlib.Path("/foo/bar")) == pathlib.Path("/foo/bar")
    assert util.parse_path(b"/foo/bar") == pathlib.Path("/foo/bar")

    # relative paths. implicit conversion to cwd
    assert util.parse_path("foo/bar") == pathlib.Path.cwd() / "foo/bar"
    assert util.parse_path(pathlib.Path("foo/bar")) == pathlib.Path.cwd() / "foo/bar"
    assert util.parse_path(b"foo/bar") == pathlib.Path.cwd() / "foo/bar"

    # relative paths. conversion using custom base directory
    base_dir = pathlib.Path("/foo/bar")
    assert util.parse_path("baz", base_directory=base_dir) == base_dir / "baz"
    assert (
        util.parse_path(pathlib.Path("baz"), base_directory=base_dir)
        == base_dir / "baz"
    )
    assert util.parse_path(b"baz", base_directory=base_dir) == base_dir / "baz"

    # pass a relative path but disallowing relative paths. should raise error.
    with pytest.raises(ValueError, match="Path .* is not absolute"):
        util.parse_path("foo/bar", allow_relative=False)

    # pass a base direectory but disallowing relative paths. should raise error.
    with pytest.raises(
        ValueError,
        match="If `base_directory` is specified, then `allow_relative` must be True.",
    ):
        util.parse_path(
            "foo/bar",
            base_directory=pathlib.Path("/foo/bar"),
            allow_relative=False,
        )

    # Parse optional path. Works the same way but passes None down the line.
    assert util.parse_optional_path(None) is None
    assert util.parse_optional_path("/foo/bar") == util.parse_path("/foo/bar")


def test_dict_obs():
    A = np.random.rand(3, 4)
    B = np.random.rand(3, 7, 1)
    C = np.random.rand(4)

    ab = types.DictObs({"a": A, "b": B})
    abc = types.DictObs({"a": A, "b": B, "c": C})

    # len
    assert len(ab) == 3
    with pytest.raises(RuntimeError):
        len(abc)
    with pytest.raises(RuntimeError):
        len(types.DictObs({}))

    assert abc.dict_len == 3

    # slicing
    np.testing.assert_equal(abc[0].get("a"), A[0])
    np.testing.assert_equal(abc[0].get("c"), np.array(C[0]))
    np.testing.assert_equal(abc[0:2].get("a"), np.array(A[0:2]))
    np.testing.assert_equal(ab[:, 0].get("a"), np.array(A[:, 0]))
    with pytest.raises(IndexError):
        abc[:, 0]

    # iter
    for i, a_row in enumerate(A):
        np.testing.assert_equal(a_row, ab[i].get("a"))
    assert ab[0] == next(iter(ab))

    # eq
    assert abc == types.DictObs({"a": A, "b": B, "c": C})
    assert abc == types.DictObs({"a": np.array(A), "b": np.array(B), "c": np.array(C)})
    assert abc != types.DictObs({"a": A, "c": B, "b": C})  # diff keys
    assert abc != types.DictObs({"a": A, "b": B + 1, "c": C})  # diff values
    assert abc != {"a": A, "b": B + 1, "c": C}  # diff type
    assert abc != ab  # diff keys

    # shape / dtype
    assert abc.shape == {"a": A.shape, "b": B.shape, "c": C.shape}
    assert abc.dtype == {"a": A.dtype, "b": B.dtype, "c": C.dtype}

    # wrap
    assert types.maybe_wrap_in_dictobs({"a": A, "b": B, "c": C}) == abc
    assert abc.unwrap() == {"a": A, "b": B, "c": C}

    # map, stack, concat
    assert abc.map_arrays(lambda arr: arr + 1) == types.DictObs(
        {"a": A + 1, "b": B + 1, "c": C + 1},
    )
    assert types.DictObs.stack(list(iter(ab))) == ab
    np.testing.assert_equal(
        types.DictObs.concatenate([abc, abc]).get("a"),
        np.concatenate([A, A]),
    )

    with pytest.raises(AssertionError):
        types.assert_not_dictobs(abc)

    with pytest.raises(TypeError):
        types.DictObs({"a": "not an array"})  # type: ignore[wrong-arg-types]