from typing import Type
import pytest
import torch
import torch.nn as nn
from torch import Tensor
import lqsvg.torch.named as nt
from lqsvg.envs import lqr
from lqsvg.envs.lqr import pack_obs, unpack_obs
from lqsvg.envs.lqr.modules.dynamics.gauss import InitStateDynamics
from lqsvg.testing.fixture import standard_fixture
from lqsvg.torch.modules import CholeskyFactor
dim = standard_fixture((2, 4, 8), "Dim")
@pytest.fixture
def module_cls() -> Type[InitStateDynamics]:
return InitStateDynamics
@pytest.fixture
def module(module_cls: Type[InitStateDynamics], dim: int) -> InitStateDynamics:
return module_cls(dim)
@pytest.fixture
def state(dim: int, batch_shape: tuple[int, ...]) -> Tensor:
return nt.vector(torch.randn(batch_shape + (dim, )))
# pylint:disable=invalid-name
from __future__ import annotations
import pytest
import torch
from torch import Tensor
import lqsvg.torch.named as nt
from lqsvg.envs import lqr
from lqsvg.envs.lqr.generators import make_lindynamics, make_linsdynamics
from lqsvg.np_util import RNG
from lqsvg.policy.modules import TVLinearFeedback, TVLinearPolicy
from lqsvg.testing.fixture import standard_fixture
frozen = standard_fixture((True, False), "Frozen")
@pytest.fixture
def dynamics(n_state: int, n_ctrl: int, horizon: int,
rng: RNG) -> lqr.LinSDynamics:
return make_linsdynamics(
make_lindynamics(
n_state,
n_ctrl,
horizon,
stationary=True,
passive_eigval_range=(0.5, 1.5),
rng=rng,
),
stationary=True,
rng=rng,
@pytest.fixture(autouse=True)
def torch_generator_state(seed: int):
with default_generator_seed(seed):
yield
@pytest.fixture
def n_state() -> int:
return 2
@pytest.fixture
def n_ctrl() -> int:
return 2
@pytest.fixture
def horizon() -> int:
return 20
stationary = standard_fixture((True, False), "Stationary")
@pytest.fixture
def lqg_generator(
n_state: int, n_ctrl: int, horizon: int, stationary: bool, seed: int
) -> LQGGenerator:
return LQGGenerator(n_state, n_ctrl, horizon, stationary=stationary, rng=seed)
def test_init(self, n_state: int, n_ctrl: int, horizon: int, stationary: bool):
module = LinearTransitionModel(n_state, n_ctrl, horizon, stationary)
assert isinstance(module, LinearDynamicsModule)
assert all(
list(hasattr(module, attr) for attr in "n_state n_ctrl horizon".split())
)
class TestLinearDiagDynamicsModel(DynamicsModuleTests, LinearParamsTestMixin):
@pytest.fixture
def module(self, n_state: int, n_ctrl: int, horizon: int, stationary: bool):
return LinearDiagDynamicsModel(n_state, n_ctrl, horizon, stationary)
hunits = standard_fixture([(), (10,), (32, 32)], "HiddenUnits")
activation = standard_fixture((None, "ReLU", "ELU", "Tanh"), "Activation")
@pytest.mark.skip(reason="Unimplemented")
class TestMLPTransitionModel(DynamicsModuleTests):
@pytest.fixture
def module(
self,
n_state: int,
n_ctrl: int,
horizon: int,
hunits: tuple[int, ...],
activation: str,
) -> MLPDynamicsModel:
# pylint:disable=too-many-arguments
import pytest
from ray.rllib import RolloutWorker
from torch import Tensor
import lqsvg.torch.named as nt
from lqsvg.experiment.data import DataModuleSpec, TrajectoryData, TransitionData
from lqsvg.experiment.worker import make_worker
from lqsvg.testing.fixture import standard_fixture
n_state = standard_fixture((2, 3), "NState")
n_ctrl = standard_fixture((2, 3), "NCtrl")
horizon = standard_fixture((1, 7, 10), "Horizon")
@pytest.fixture
def worker(n_state: int, n_ctrl: int, horizon: int) -> RolloutWorker:
with nt.suppress_named_tensor_warning():
worker = make_worker(
env_config=dict(n_state=n_state,
n_ctrl=n_ctrl,
horizon=horizon,
num_envs=2),
log_level="WARNING",
)
return worker
def test_trajectory_data(worker: RolloutWorker, n_state: int, n_ctrl: int,
horizon: int):
data_spec = DataModuleSpec(total_trajs=100)
datamodule = TrajectoryData(worker, data_spec)
from __future__ import annotations
import pytest
from torch import Tensor
import lqsvg.torch.named as nt
from lqsvg.testing.fixture import standard_fixture
from lqsvg.torch.utils import make_spd_matrix
n_row = standard_fixture((1, 2, 4), "Rows")
@pytest.fixture
def spdm(n_row: int, seed: int) -> Tensor:
return nt.matrix(make_spd_matrix(n_row, sample_shape=(), rng=seed))
def test_cholesky(spdm: Tensor):
scale_tril = nt.cholesky(spdm)
assert scale_tril.shape == spdm.shape
assert scale_tril.names == spdm.names
assert scale_tril.dtype == spdm.dtype
assert (nt.diagonal(scale_tril) >= 0).all()
assert nt.allclose(scale_tril @ nt.transpose(scale_tril), spdm)
env = RandomLQGEnv(config)
env.reset()
act = env.action_space.sample()
new_obs, rew, done, info = env.step(act)
# pylint:disable=unsupported-membership-test
assert new_obs in env.observation_space
assert isinstance(rew, float)
assert isinstance(done, bool)
assert isinstance(info, dict)
# ==============================================================================
num_envs = standard_fixture((1, 2, 4), "NEnvs")
@pytest.fixture
def vector_config(config: dict, num_envs: int) -> dict:
config["num_envs"] = num_envs
return config
# Test RandomVectorLQG =========================================================
def test_vector_init(vector_config: dict):
env = RandomVectorLQG(vector_config)
assert env.num_envs == vector_config["num_envs"]
assert hasattr(env, "curr_states")
assert env.curr_states is None
from lqsvg.envs import lqr
from lqsvg.envs.lqr.generators import make_lindynamics, make_linsdynamics
from lqsvg.envs.lqr.utils import (
ctrb,
minimal_sample_shape,
random_mat_with_eigval_range,
random_matrix_from_eigs,
wrap_sample_shape_to_size,
)
from lqsvg.np_util import RNG
from lqsvg.testing.fixture import standard_fixture
from .utils import scalar_to_matrix, sort_eigfactors, vector_to_matrix
mat_dim = standard_fixture((2, 3, 4), "MatDim")
eigval_range = standard_fixture([(0, 1), (0.5, 1.5)], "EigvalRange")
n_batch = standard_fixture((None, 1, 4), "NBatch")
vec_dim = standard_fixture((2, 3, 4), "VecDim")
batch_shape = standard_fixture([(), (1, ), (2, ), (2, 1)], "BatchShape")
@pytest.fixture(params=(0, 2))
def dim(request) -> int:
return request.param
@pytest.fixture
def sampler(dim: int) -> Callable[[int], np.ndarray]:
call = normal.rvs if dim == 0 else partial(ortho_group.rvs, dim=3)
def squared(x: np.ndarray) -> np.ndarray:
return x**2
def dot_2d(x: np.ndarray) -> np.ndarray:
return x.dot(np.array([1, 1]))
def dot_3d(x: np.ndarray) -> np.ndarray:
return x.dot(np.array([1, 1, 1]))
FUNCTIONS = {1: squared, 2: dot_2d, 3: dot_3d}
dim = standard_fixture((1, 2, 3), "Dim")
max_scaling = standard_fixture((0.1, 1.0, 3.0), "MaxScaling")
steps = standard_fixture((2, 5, 10), "Steps")
@pytest.fixture
def rng(seed: int) -> Generator:
return np.random.default_rng(seed)
@pytest.fixture
def f_delta(dim: int) -> Callable[[np.ndarray], np.ndarray]:
return FUNCTIONS[dim]
@pytest.fixture
from __future__ import annotations
import pytest
import torch
from torch import Tensor
import lqsvg.torch.named as nt
from lqsvg.envs import lqr
from lqsvg.envs.lqr.modules.reward.quadratic import QuadraticReward
from lqsvg.envs.lqr.utils import pack_obs, unpack_obs
from lqsvg.testing.fixture import standard_fixture
batch_shape = standard_fixture([(), (1,), (4,), (2, 2)], "BatchShape")
class TestQuadraticReward:
@pytest.fixture
def module(self, n_state: int, n_ctrl: int, horizon: int) -> QuadraticReward:
return QuadraticReward(n_state + n_ctrl, horizon)
@pytest.fixture
def state(self, n_state: int, batch_shape: tuple[int, ...]) -> Tensor:
return nt.vector(torch.randn(batch_shape + (n_state,)))
@pytest.fixture
def obs(self, state: Tensor, horizon: int, batch_shape: tuple[int, ...]) -> Tensor:
time = torch.randint(low=0, high=horizon, size=batch_shape + (1,))
return pack_obs(state, nt.vector(time)).requires_grad_(True)
@pytest.fixture
def act(self, n_ctrl: int, batch_shape) -> Tensor:
import pytest
from lqsvg.testing.fixture import standard_fixture
from lqsvg.torch.utils import default_generator_seed
seed = standard_fixture((1, 2, 3), "Seed")
@pytest.fixture(autouse=True)
def torch_generator_state(seed: int):
with default_generator_seed(seed):
yield
LQGGenerator,
box_ddp_random_lqr,
make_lindynamics,
make_linsdynamics,
make_lqr,
make_lqr_linear_navigation,
make_quadcost,
stack_lqs,
)
from lqsvg.envs.lqr.types import LinDynamics, LinSDynamics, QuadCost
from lqsvg.envs.lqr.utils import stationary_dynamics_factors
from lqsvg.testing.fixture import standard_fixture
GeneratorFn = Callable[..., LQGGenerator]
passive_eigval_range = standard_fixture([None, (0.0, 1.0), (0.5, 1.5)],
"PassiveEigvals")
stat_ctrb = standard_fixture([(True, False), (True, True), (False, False)],
"Stationary/Controllable")
transition_bias = standard_fixture((True, False), "TransBias")
sample_covariance = standard_fixture((True, False), "SampleCov")
cost_linear = standard_fixture((True, False), "CostLinear")
cost_cross = standard_fixture((True, False), "CostCross")
n_batch = standard_fixture((None, 1, 4), "NBatch")
@pytest.fixture
def stationary(stat_ctrb: tuple[bool, bool]) -> bool:
stat, _ = stat_ctrb
return stat
from __future__ import annotations
import numpy as np
import torch
from lqsvg.np_util import make_spd_matrix, random_unit_col_matrix, random_unit_vector
from lqsvg.testing.fixture import standard_fixture
seed = standard_fixture((42, 69, 37), "Seed")
batch_shape = standard_fixture([(), (1, ), (2, ), (2, 1)], "BatchShape")
vec_dim = standard_fixture((2, 3, 4), "VecDim")
n_row = standard_fixture((2, 3, 4), "Rows")
n_col = standard_fixture((2, 3, 4), "Columns")
def test_spd_matrix(n_row: int, batch_shape: tuple[int, ...], seed: int):
A = make_spd_matrix(n_row, sample_shape=batch_shape, rng=seed)
assert A.shape == batch_shape + (n_row, n_row)
B = torch.as_tensor(A)
assert torch.allclose(B, B.transpose(-2, -1))
eigval, _ = torch.linalg.eigh(B)
assert eigval.ge(0).all()
def test_random_unit_vector(vec_dim: int, batch_shape: tuple[int, ...],
seed: int):
vector = random_unit_vector(vec_dim, sample_shape=batch_shape, rng=seed)
assert isinstance(vector, np.ndarray)
assert vector.shape == batch_shape + (vec_dim, )
from __future__ import annotations
import numpy as np
import pytest
import torch
from numpy.random import Generator
from torch import Tensor
from lqsvg.envs.lqr.utils import pack_obs, unpack_obs
from lqsvg.testing.fixture import standard_fixture
from lqsvg.torch import named as nt
from lqsvg.torch.utils import default_generator_seed
n_state = standard_fixture((2, 3), "NState")
n_ctrl = standard_fixture((2, 3), "NCtrl")
horizon = standard_fixture((1, 3, 10), "Horizon")
stationary = standard_fixture((True, False), "Stationary")
seed = standard_fixture((1, 2, 3), "Seed")
batch_shape = standard_fixture([(), (1, ), (4, )], "BatchShape")
@pytest.fixture
def rng(seed: int) -> Generator:
return np.random.default_rng(seed)
@pytest.fixture(autouse=True)
def torch_generator_state(seed: int):
with default_generator_seed(seed):
yield
