def make_quadcost(
state_size: int,
ctrl_size: int,
horizon: int,
stationary: bool = True,
n_batch: Optional[int] = None,
linear: bool = False,
cross_terms: bool = False,
rng: RNG = None,
) -> QuadCost:
"""Generate quadratic cost parameters.
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
linear: whether to include a linear term in addition to the quadratic
cross_terms: whether to include state-ctrl cross terms in the quadratic
(C_sa and C_as)
rng: random number generator, seed, or None
"""
# pylint:disable=too-many-arguments,too-many-locals
rng = np.random.default_rng(rng)
n_tau = state_size + ctrl_size
kwargs = dict(horizon=horizon,
stationary=stationary,
n_batch=n_batch,
rng=rng)
C = utils.random_spd_matrix(n_tau, **kwargs)
C_s, C_a = nt.split(C, [state_size, ctrl_size], dim="C")
C_ss, C_sa = nt.split(C_s, [state_size, ctrl_size], dim="R")
C_as, C_aa = nt.split(C_a, [state_size, ctrl_size], dim="R")
if not cross_terms:
C_sa, C_as = torch.zeros_like(C_sa), torch.zeros_like(C_as)
C_s = torch.cat((C_ss, C_sa), dim="R")
C_a = torch.cat((C_as, C_aa), dim="R")
C = torch.cat((C_s, C_a), dim="C")
if linear:
c = utils.random_normal_vector(n_tau, **kwargs)
else:
c = utils.expand_and_refine(nt.vector(torch.zeros(n_tau)),
1,
horizon=horizon,
n_batch=n_batch)
return QuadCost(C, c)
def unpack_obs(obs: Tensor) -> tuple[Tensor, IntTensor]:
"""Unpack observation into state variables and time.
Expects observation as a named 'vector' tensor.
"""
# noinspection PyArgumentList
state, time = nt.split(obs, [obs.size("R") - 1, 1], dim="R")
time = time.int()
return state, time
def last_obs(
n_state: int,
horizon: int,
batch_shape: tuple[int, ...],
batch_names: tuple[str, ...],
) -> Tensor:
state = nt.vector(torch.randn(batch_shape + (n_state, ))).refine_names(
*batch_names, ...)
dummy, _ = nt.split(state, [1, n_state - 1], dim="R")
time = torch.full_like(dummy, fill_value=horizon).int()
return pack_obs(state, time).requires_grad_()
def obs(
n_state: int,
horizon: int,
batch_shape: tuple[int, ...],
batch_names: tuple[str, ...],
) -> Tensor:
state = nt.vector(torch.randn(batch_shape + (n_state, ))).refine_names(
*batch_names, ...)
dummy, _ = nt.split(state, [1, n_state - 1], dim="R")
time = torch.randint_like(nt.unnamed(dummy), low=0, high=horizon)
time = time.refine_names(*dummy.names).int()
return pack_obs(state, time).requires_grad_()
def dynamics_factors(dynamics: AnyDynamics) -> tuple[Tensor, Tensor]:
"""Returns the unactuated and actuaded parts of the transition matrix."""
# pylint:disable=invalid-name
n_state, n_ctrl, _ = dims_from_dynamics(dynamics)
F_s, F_a = nt.split(dynamics.F, [n_state, n_ctrl], dim="C")
return F_s, F_a