def policy(n_state: int, n_ctrl: int, horizon: int) -> lqr.Linear:
K = torch.Tensor(horizon, n_ctrl, n_state)
k = torch.Tensor(horizon, n_ctrl)
nn.init.xavier_uniform_(K)
nn.init.constant_(k, 0)
K, k = nt.horizon(nt.matrix(K), nt.vector(k))
return K, k
def make_gaussinit(
state_size: int,
n_batch: Optional[int] = None,
sample_covariance: bool = False,
rng: RNG = None,
) -> GaussInit:
"""Generate parameters for Gaussian initial state distribution.
Args:
state_size: size of state vector
n_batch: batch size, if any
sample_covariance: whether to sample a random SPD matrix for the
Gaussian covariance or use the identity matrix.
rng: random number generator, seed, or None
"""
# pylint:disable=invalid-name
vec_shape = (state_size, )
batch_shape = () if n_batch is None else (n_batch, )
mu = torch.zeros(batch_shape + vec_shape)
if sample_covariance:
sig = as_float_tensor(
make_spd_matrix(state_size, sample_shape=batch_shape, rng=rng))
else:
sig = torch.eye(state_size)
return GaussInit(
mu=utils.expand_and_refine(nt.vector(mu), 1, n_batch=n_batch),
sig=utils.expand_and_refine(nt.matrix(sig), 2, n_batch=n_batch),
)
def scale_tril(self) -> Tensor:
"""Compute scale tril from pre-diagonal parameters.
Output is differentiable w.r.t. pre-diagonal parameters.
"""
diag = softplus(self.pre_diag, beta=self._softplus_beta)
return nt.matrix(torch.diag_embed(diag))
def _gains_at(
self,
index: Union[IntTensor, LongTensor,
None] = None) -> tuple[Tensor, Tensor]:
K, k = nt.horizon(nt.matrix(self.K), nt.vector(self.k))
if index is not None:
index = torch.clamp(index, max=self.horizon - 1)
# Assumes index is a named scalar tensor
# noinspection PyTypeChecker
K, k = (nt.index_by(x, dim="H", index=index) for x in (K, k))
return K, k
def test_init(self, module: CholeskyFactor, shape: tuple[int, ...]):
assert hasattr(module, "beta")
assert hasattr(module, "ltril")
assert hasattr(module, "pre_diag")
assert isinstance(module.ltril, nn.Parameter)
assert isinstance(module.pre_diag, nn.Parameter)
assert module.ltril.shape == shape
assert module.pre_diag.shape == shape[:-1]
cholesky = module()
assert nt.allclose(cholesky, nt.matrix(torch.eye(shape[-1])))
def _transition_factors(
self,
index: Optional[IntTensor] = None) -> (Tensor, Tensor, Tensor):
F, f, L = nt.horizon(nt.matrix(self.F), nt.vector(self.f),
self.scale_tril())
if index is not None:
if self.stationary:
idx = torch.zeros_like(index)
else:
# Timesteps after termination use last parameters
idx = torch.clamp(index, max=self.horizon - 1).int()
F, f, L = (nt.index_by(x, dim="H", index=idx) for x in (F, f, L))
return F, f, L
def index_quadratic_parameters(
quad: nn.Parameter,
linear: nn.Parameter,
const: nn.Parameter,
index: IntTensor,
max_idx: int,
) -> tuple[Tensor, Tensor, Tensor]:
# pylint:disable=missing-function-docstring
quad, linear, const = nt.horizon(nt.matrix(quad), nt.vector(linear),
nt.scalar(const))
index = torch.clamp(index, max=max_idx)
quad, linear, const = map(lambda x: nt.index_by(x, dim="H", index=index),
(quad, linear, const))
return quad, linear, const
def check_dynamics_covariance(W: Tensor, n_state: int, horizon: int,
stationary: int, sample_covariance: bool):
assert_horizon_len(W, horizon)
assert_row_size(W, n_state)
assert_col_size(W, n_state)
assert nt.allclose(W, nt.transpose(W))
eigval, _ = torch.linalg.eigh(nt.unnamed(W))
assert eigval.gt(0).all()
assert sample_covariance != nt.allclose(W, nt.matrix(torch.eye(n_state)))
# noinspection PyTypeChecker
assert (horizon == 1 or not sample_covariance
or stationary == nt.allclose(W, W.select("H", 0)))
def test_factorize_(
self,
module: SPDMatrix,
size: int,
sample_shape: tuple[int, ...],
use_sample_shape: bool,
seed: int,
):
# pylint:disable=too-many-arguments
sample_shape = sample_shape if use_sample_shape else ()
A = make_spd_matrix(size, sample_shape=sample_shape, rng=seed)
module.factorize_(nt.matrix(A))
B = nt.unnamed(module())
A, B = torch.broadcast_tensors(A, B)
isclose = torch.isclose(A, B, atol=1e-6)
assert isclose.all(), (A[~isclose].tolist(), B[~isclose].tolist())
def refine_lqr(dynamics: LinDynamics,
cost: QuadCost) -> Tuple[LinDynamics, QuadCost]:
"""Add dimension names to LQR parameters.
Args:
dynamics: transition matrix and vector
cost: quadratic cost matrix and vector
Returns:
A tuple with named dynamics and cost parameters
"""
F, f = dynamics
C, c = cost
F, C = nt.matrix(F, C)
f, c = nt.vector(f, c)
F, f, C, c = nt.horizon(F, f, C, c)
return LinDynamics(F, f), QuadCost(C, c)
def test_factorize_(
self,
module: CholeskyFactor,
size: int,
sample_shape: tuple[int, ...],
use_sample_shape: bool,
seed: int,
):
# pylint:disable=too-many-arguments
sample_shape = sample_shape if use_sample_shape else ()
A = make_spd_matrix(size, sample_shape=sample_shape, rng=seed)
module.factorize_(nt.matrix(A))
L = nt.unnamed(module())
C = nt.cholesky(A)
C, L = torch.broadcast_tensors(C, L)
isclose = torch.isclose(C, L, atol=1e-6)
assert isclose.all(), (C[~isclose].tolist(), L[~isclose].tolist())
def random_spd_matrix(
size: int,
horizon: int,
stationary: bool = False,
n_batch: Optional[int] = None,
rng: RNG = None,
) -> Tensor:
# pylint:disable=missing-function-docstring
mat = make_spd_matrix(
size,
sample_shape=minimal_sample_shape(
horizon, stationary=stationary, n_batch=n_batch
),
rng=rng,
)
mat = nt.matrix(as_float_tensor(mat))
mat = expand_and_refine(mat, 2, horizon=horizon, n_batch=n_batch)
return mat
def random_uniform_matrix(
row_size: int,
col_size: int,
horizon: int,
stationary: bool = False,
low: float = 0.0,
high: float = 1.0,
n_batch: Optional[int] = None,
rng: RNG = None,
) -> Tensor:
"""Matrix with Uniform i.i.d. entries."""
# pylint:disable=too-many-arguments
mat_shape = (row_size, col_size)
shape = (
minimal_sample_shape(horizon, stationary=stationary, n_batch=n_batch)
+ mat_shape
)
mat = nt.matrix(as_float_tensor(rng.uniform(low=low, high=high, size=shape)))
mat = expand_and_refine(mat, 2, horizon=horizon, n_batch=n_batch)
return mat
def make_linsdynamics(
dynamics: LinDynamics,
stationary: bool = False,
n_batch: Optional[int] = None,
sample_covariance: bool = True,
rng: RNG = None,
) -> LinSDynamics:
"""Generate stochastic linear dynamics from linear deterministic dynamics.
Warning:
This function does not check if `stationary`, and `nbatch` are
compatible with `dynamics` (i.e., if the dynamics satisfy these
parameters), but passing incompatible dynamics may lead to errors
downstream.
Args:
dynamics: linear deterministic transition dynamics
stationary: whether dynamics vary with time
n_batch: batch size, if any
sample_covariance: whether to sample a random SPD matrix for the
Gaussian covariance or use the identity matrix.
rng: random number generator, seed, or None
"""
# pylint:disable=too-many-arguments
rng = np.random.default_rng(rng)
state_size, _, horizon = utils.dims_from_dynamics(dynamics)
if sample_covariance:
W = utils.random_spd_matrix(state_size,
horizon=horizon,
stationary=stationary,
n_batch=n_batch,
rng=rng)
else:
W = utils.expand_and_refine(nt.matrix(torch.eye(state_size)),
2,
horizon=horizon,
n_batch=n_batch)
F, f = dynamics
return LinSDynamics(F, f, W)
def linear(self, n_state: int, n_ctrl: int, horizon: int) -> lqr.Linear:
K, k = torch.randn(horizon, n_ctrl,
n_state), torch.randn(horizon, n_ctrl)
K, k = nt.horizon(nt.matrix(K), nt.vector(k))
return K, k
def policy(n_state: int, n_ctrl: int, horizon: int) -> Linear:
K = torch.rand((horizon, n_ctrl, n_state))
k = torch.rand((horizon, n_ctrl))
K, k = nt.horizon(nt.matrix(K), nt.vector(k))
return K, k
def _refined_parameters(self) -> tuple[Tensor, Tensor]:
C, c = nt.horizon(nt.matrix(self.C), nt.vector(self.c))
return C, c
def spdm(n_row: int, seed: int) -> Tensor:
return nt.matrix(make_spd_matrix(n_row, sample_shape=(), rng=seed))
def refine_dynamics_input(dynamics: LinDynamics):
F, f = dynamics
F, f = nt.horizon(nt.matrix(F), nt.vector_to_matrix(f))
return LinDynamics(F, f)
def make_lindynamics(
state_size: int,
ctrl_size: int,
horizon: int,
stationary: bool = False,
n_batch: Optional[int] = None,
passive_eigval_range: Optional[tuple[float, float]] = None,
controllable: bool = False,
bias: bool = True,
rng: RNG = None,
) -> LinDynamics:
"""Generate linear transition matrices.
Args:
state_size: size of state vector
ctrl_size: size of control vector
horizon: length of the horizon
stationary: whether dynamics vary with time
n_batch: batch size, if any
passive_eigval_range: range of eigenvalues for the unnactuated system.
If None, samples the F_s matrix entries independently from a
standard normal distribution
controllable: whether to ensure the actuator dynamics (the B matrix of
the (A,B) pair) make the system controllable
bias: whether to use a non-zero bias vector for transition dynamics
rng: random number generator, seed, or None
Raises:
ValueError: if `controllable` is True but not `stationary`
"""
# pylint:disable=too-many-arguments
if controllable and not stationary:
raise ValueError(
"Controllable non-stationary dynamics are unsupported.")
rng = np.random.default_rng(rng)
Fs, _, eigvec = generate_passive(
state_size,
eigval_range=passive_eigval_range,
horizon=horizon,
stationary=stationary,
n_batch=n_batch,
rng=rng,
)
Fa = generate_active(Fs,
ctrl_size,
eigvec=eigvec,
controllable=controllable,
rng=rng)
Fs = utils.expand_and_refine(nt.matrix(as_float_tensor(Fs)),
2,
horizon=horizon,
n_batch=n_batch)
Fa = utils.expand_and_refine(nt.matrix(as_float_tensor(Fa)),
2,
horizon=horizon,
n_batch=n_batch)
F = torch.cat((Fs, Fa), dim="C")
if bias:
f = random_unit_vector(
state_size,
sample_shape=utils.minimal_sample_shape(horizon, stationary,
n_batch),
rng=rng,
)
else:
f = np.zeros(state_size)
f = nt.vector(as_float_tensor(f))
f = utils.expand_and_refine(f, 1, horizon=horizon, n_batch=n_batch)
return LinDynamics(F, f)
def refine_cost_ouput(cost: QuadCost) -> QuadCost:
C, c = cost
C, c = nt.horizon(nt.matrix(C), nt.matrix_to_vector(c))
return QuadCost(C, c)
def forward(self) -> Tensor:
"""Compute the Cholesky factor from parameters."""
ltril = nt.matrix(self.ltril)
pre_diag = nt.vector(self.pre_diag)
return assemble_cholesky(ltril, pre_diag, beta=self.beta)
def gains(self) -> lqr.Linear:
"""Return current parameters as linear parameters."""
K, k = nt.horizon(nt.matrix(self.K), nt.vector(self.k))
K.grad, k.grad = self.K.grad, self.k.grad
return K, k
def refine_cost_input(cost: QuadCost):
C, c = cost
C, c = nt.horizon(nt.matrix(C), nt.vector_to_matrix(c))
return QuadCost(C, c)
def refine_sdynamics_input(dynamics: LinSDynamics):
F, f, W = dynamics
F, f = refine_dynamics_input((F, f))
W = nt.horizon(nt.matrix(W))
return LinSDynamics(F, f, W)
def forward(self) -> Tensor:
"""Compute the symmetric positive definite matrix from parameters."""
ltril = nt.matrix(self.ltril)
pre_diag = nt.vector(self.pre_diag)
cholesky = assemble_cholesky(ltril, pre_diag, beta=self.beta)
return cholesky @ nt.transpose(cholesky)
def refine_linear_output(linear: Linear):
K, k = linear
K, k = nt.horizon(nt.matrix(K), nt.matrix_to_vector(k))
return K, k
def refine_linear_input(linear: Linear):
K, k = linear
K, k = nt.horizon(nt.matrix(K), nt.vector_to_matrix(k))
return K, k
def init(dim: int) -> lqr.GaussInit:
return lqr.GaussInit(nt.vector(torch.randn(dim)),
nt.matrix(torch.eye(dim)))
def refine_quadratic_output(quadratic: Quadratic):
A, b, c = quadratic
A, b, c = nt.horizon(nt.matrix(A), nt.matrix_to_vector(b),
nt.matrix_to_scalar(c))
return A, b, c
