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