dair_pll.state_space

Classes and utilities for operation on Lie group/algebra state spaces.

This module implements a StateSpace abstract type which defines fundamental operations on states in spaces that are not necessarily Euclidean. In general, we assume configurations spaces are Lie groups, which together with their associated Lie algebra form a state space. We implement associated operations on these spaces.

Of note is the FloatingBaseSpace, with configurations in the Cartesian product of \(SE(3)\) and \(\mathbb{R}^m\). This space receives the particular implementation of representing the \(SO(3)\) configuration as a quaternion and the Lie algebra velocity/tangent space as the body-axes angular velocity / rotation vector.

This is also the place where batching dimensions are defined for states. By convention, the state element index is always the last dimension of the tensor, and when states are batched in time, time is the second-to-last index.

dair_pll.state_space.partial_sum_batch(summands, keep_batch=False)[source]

Sums over a batch, possibly keeping the first batching dimension.

Parameters:

summands (Tensor) – (b_1, ..., b_j, n_1, ..., n_k) tensor batch of tensors
keep_batch (bool) – whether to keep the first batching dimension

Return type:

Tensor

Returns:

sum of x as scalar tensor, or (b_1,) tensor if keep_batch == True.

class dair_pll.state_space.StateSpace(n_q, n_v)[source]

Bases: ABC

Mathematical model of a state space.

Each state space is modeled as the Cartesian product of a connected Lie group \(G\) and its associated Lie algebra \(\mathfrak g\) ( equivalently up to diffeomorphism, the tangent bundle \(TG\)). The Lie group element may be given non-minimal coordinates, e.g. representing SO(3) with quaternions. As \(\mathfrak g\) is a vector space, \(G \times \mathfrak g\) itself is also a Lie group.

The following assumptions about the group \(G\) and algebra g are made:

The Lie group exponential map \(\exp: \mathfrak g \to G\) is surjective/onto, such that a left inverse \(\log: G \to \mathfrak g\) can be defined, i.e. \(\exp(\log(g)) = g\).

The Lie group exponential map coincides with the underlying manifold’s Riemannian geometric exponential map, such that the geodesic distance from \(g_1\) to \(g_1\) is \(|\log(g_2 \cdot g_1^{-1})|\)

These conditions are met if and only if \(G\) is the Cartesian product of a compact group and an Abelian group [Mil76] – For example, \(SO(3)\times\mathbb{R}^n\).

For each concrete class inheriting from StateState, a few fundamental mathematical operators associated with Lie groups must be defined on these coordinates. StateState defines several other group operations from these units.

Parameters:

n_q (int) – number of Lie group (configuration) coordinates (\(>= 0\))
n_v (int) – number of Lie algebra (velocity) coordinates (\(>= 0\))

n_q: int

n_v: int

n_x: int

comparisons: typing.Dict[str, typing.Callable[[Tensor, Tensor], Tensor]]

abstract configuration_difference(q_1, q_2)[source]

Returns the relative transformation between q_1 and q_2.

Specifically, as \(G\) is a Lie group, it has a well-defined inverse operator. This function returns \(dq = \log(q_2 \cdot q_1^{-1})\), i.e. the Lie algebra element such that \(q_1 \exp(dq) = q_2\).

This method has a corresponding “inverse” function exponential().

Parameters:

q_1 (Tensor) – (*, n_q) “starting” configuration, element(s) of Lie group \(G\).
q_2 (Tensor) – (*, n_q) “ending” configuration, element(s) of Lie group \(G\).

Return type:

Tensor

Returns:

(*, n_v) element of Lie algebra g defining the transformation: from q_1 to q_2

abstract exponential(q, dq)[source]

Evolves q along the Lie group G in the direction dq.

This function implements the inverse of configuration_difference() by returning q * exp(dq).

Parameters:

q (Tensor) – (*, n_q) “starting” configuration, element(s) of Lie group G
dq (Tensor) – (*, n_v) perturbation, element(s) of Lie algebra g

Return type:

Tensor

Returns:

(*, n_q) group product of q and exp(dq)

abstract project_configuration(q)[source]

Projects a tensor of size (*, n_q) onto the Lie group G.

This function is used, mostly for numerical stability, to ensure a (*, n_q) tensor corresponds to Lie group elements. While not necessarily a Euclidean projection, this function should be:

The identity on G, i.e. q = projection_configuration(q)

Continuous

(Piecewise) differentiable near G

Parameters:: q (Tensor) – (*, n_q) vectors to project onto G.
Return type:: Tensor
Returns:: (*, n_q) projection of q onto G.

abstract zero_state()[source]

Identity element of the Lie group G.

Entitled “zero state” as the group operation is typically thought of as addition.

Return type:: Tensor
Returns:: (n_x,) tensor group identity

q(x)[source]

Selects configuration indices from state(s) x

Return type:: Tensor

v(x)[source]

Selects velocity indices from state(s) x

Return type:: Tensor

q_v(x)[source]

Separates state(s) x into configuration and velocity

Return type:: Tuple[Tensor, Tensor]

x(q, v)[source]

Concatenates configuration q and velocity v into a state

Return type:: Tensor

config_square_error(q_1, q_2, keep_batch=False)[source]

Returns squared distance between two Lie group elements/configurations.

Interprets an \(l_2\)-like error between two configurations as the square of the geodesic distance between them. This is simply equal to \(|\log(q_2 \mathrm{inv}(q_1))|^2\) under the assumptions about G.

Parameters:

q_1 (Tensor) – (b_1, ..., b_k, n_q) “starting” configuration
q_2 (Tensor) – (b_1, ..., b_k, n_q) “ending” configuration
keep_batch (bool) – whether to keep the outermost batch

Return type:

Tensor

Returns:

(b_1,) or scalar tensor of squared geodesic distances

velocity_square_error(v_1, v_2, keep_batch=False)[source]

Returns squared distance between two Lie algebra elements/velocities.

As the Lie algebra is a vector space, the squared error is interpreted as the geodesic/Euclidean distance

Parameters:

v_1 (Tensor) – (b_1, ..., b_k, n_v) “starting” velocity
v_2 (Tensor) – (b_1, ..., b_k, n_v) “ending” velocity
keep_batch (bool) – whether to keep the outermost batch

Return type:

Tensor

Returns:

(b_1,) or scalar tensor of squared geodesic distances.

state_square_error(x_1, x_2, keep_batch=False)[source]

Returns squared distance between two states, which are in the cartesian product G x g.

As g is a vector space, it is Abelian, and thus G x g is the product of a compact group and an Abelian group. We can then define the geodesic distance as:

dist(x_1, x_2)^2 == dist(q(x_1), q(x_2))^2 + dist(v(x_1), v(x_2))^2

Parameters:

x_1 (Tensor) – (b_1, ..., b_k, n_x) “starting” state
x_2 (Tensor) – (b_1, ..., b_k, n_x) “ending” state
keep_batch (bool) – whether to keep the outermost batch

Return type:

Tensor

Returns:

(b_1,) or scalar tensor of squared geodesic distances

auxiliary_comparisons()[source]

Any additional useful comparisons between pairs of states

Return type:: Dict[str, Callable[[Tensor, Tensor], Tensor]]

finite_difference(q, q_plus, dt)[source]

Rate of change of configuration

Interprets the rate of change of q as an element of the Lie algebra v, such that q_plus == q * exp(v * dt).

finite_difference() has a corresponding “inverse” function euler_step().

Parameters:

q (Tensor) – (*, n_q) “starting” configuration, element(s) of Lie group G
q_plus (Tensor) – (*, n_q) “ending” configuration, element(s) of Lie group G
dt (float) – time difference in [s]

Return type:

Tensor

Returns:

(*, n_v) finite-difference velocity, element(s) of Lie algebra g

euler_step(q, v, dt)[source]

Integrates q forward in time given derivative v.

Implements the inverse of finite_difference() by returning q * exp(v * dt), a geodesic forward Euler step.

Parameters:

q (Tensor) – (*, n_q) “starting” configuration, element(s) of Lie group G
v (Tensor) – (*, n_v) “starting” velocity, element(s) of Lie algebra g
dt (float) – time difference in [s]

Return type:

Tensor

Returns:

(*, n_q) configuration after Euler step.

state_difference(x_1, x_2)[source]

Returns the relative transformation between x_1 and x_2

As G x g is a Lie group, we can interpret the difference between two states via its corresponding Lie algebra, just as in configuration_difference(), as log(x_2 / x_1).

state_difference() has a corresponding “inverse” function shift_state().

Parameters:

x_1 (Tensor) – (*, n_x) “starting” state, element(s) of Lie group G x g
x_2 (Tensor) – (*, n_x) “ending” state, element(s) of Lie group G x g

Return type:

Tensor

Returns:

(*, n_x) element of Lie algebra g x R^n_v defining the transformation from x_1 to x_2

shift_state(x, dx)[source]

Evolves x along the Lie group G in the direction dx.

This function implements the inverse of state_difference() by returning q * exp(dq).

Parameters:

x (Tensor) – (*, n_x) “starting” state, element(s) of Lie group G x g
dx (Tensor) – (*, 2 * n_v) perturbation, element(s) of Lie algebra g x R^n_v

Return type:

Tensor

Returns:

(*, n_q) group product of q and exp(dq).

project_state(x)[source]

Projects a tensor of size (*, n_x) onto the state space G x g.

This function has the same basic requirements as project_configuration() translated to the lie group G x g.

Parameters:: x (Tensor) – (*, n_x) vectors to project onto G x g.
Return type:: Tensor
Returns:: (*, n_x) tensor, projection of x onto G x g.

project_derivative(x, dt)[source]

Changes velocity sequence in x to a finite difference.

Extracts configurations q_t from trajectory(ies) and replaces velocities v_i with finite_difference(q_{i-1}, q_i, dt).

Parameters:

x (Tensor) – (*, T, n_x) trajectories
dt (float) – time-step

Return type:

Tensor

Returns:

(*, T, n_x) trajectories with finite-difference velocities.

class dair_pll.state_space.FloatingBaseSpace(n_joints)[source]

Bases: StateSpace

State space with configurations in SE(3) x R^n_joints.

Called FloatingBaseSpace as it is the state space of an open kinematic chain with a free-floating base body.

Coordinates for SO(3) are unit quaternions, with remaining states represented as R^{3 + n_joints}.

Inits :py:class:`FloatingBaseSpace of prescribed size.

The floating base has configurations in SE(3), with 4 quaternion + 3 world-axes position configuration coordinates and 3 body-axes angular velocity + 3 world-axes linear velocity. Each joint is represented as a single real number.

Parameters:: n_joints (int) – number of joints in chain (>= 0)

quat(q_or_x)[source]

select quaternion elements from configuration/state

Return type:: Tensor

base(q_or_x)[source]

select floating base position elements from configuration/state

Return type:: Tensor

configuration_difference(q_1, q_2)[source]

Implements configuration offset for a floating-base rigid chain.

exp() map of SO(3) corresponds to the space of rotation vectors, or equivalently the matrix group so(3); therefore, the first 3 elements of the return value are body-axes rotation vectors.

Parameters:

q_1 (Tensor) – (*, n_q) “starting” configuration in SE(3) x R^n_joints
q_2 (Tensor) – (*, n_q) “ending” configuration, SE(3) x R^n_joints

Return type:

Tensor

Returns:

(*, n_v) body-axes rotation vector, world-axes linear displacement, and joint offsets.

exponential(q, dq)[source]

Implements exponential perturbation for a floating-base rigid chain.

This function implements the inverse of configuration_difference() by rotating quat(q) around the body-axis rotation vector in dq, and adding a linear offset to the remaining coordinates.

Parameters:

q (Tensor) – (*, n_q) “starting” configuration in SE(3) x R^n_joints
dq (Tensor) – (*, n_v) perturbation in se(3) x R^n_joints

Return type:

Tensor

Returns:

(*, n_q) perturbed quaternion, world-axes floating base origin

project_configuration(q)[source]

Implements projection onto the floating-base rigid chain configuration space.

This function projects a (*, n_q) tensor onto SE(3) x R^n_joints by simply normalizing the quaternion elements.

Parameters:: q (Tensor) – (*, n_q) vectors to project onto SE(3) x R^n_joints.
Return type:: Tensor
Returns:: (*, n_q) tensor, projection of q onto SE(3) x R^n_joints.

zero_state()[source]

Identity element of SE(3) x R^n_joints.

Return type:: Tensor
Returns:: Concatenation of identity quaternion [1, 0, 0, 0] with (n_joints + 3) zeros.

quaternion_error(x_1, x_2)[source]

Auxiliary comparison that returns floating base orientation geodesic distance.

Returns a scalar comparison of two floating base rigid chain states by the angle of rotation between their base orientations.

Parameters:

x_1 (Tensor) – (*, n_x) “starting” state
x_2 (Tensor) – (*, n_x) “ending” state

Return type:

Tensor

Returns:

scalar tensor, average angle of rotation in batch.

base_error(x_1, x_2)[source]

Auxiliary comparison that returns floating base translation geodesic distance.

Returns a scalar comparison of two floating base rigid chain states by the Euclidean between their bases.

Parameters:

x_1 (Tensor) – (*, n_x) “starting” state
x_2 (Tensor) – (*, n_x) “ending” state

Return type:

Tensor

Returns:

scalar tensor, average translation in batch.

class dair_pll.state_space.FixedBaseSpace(n_joints)[source]

Bases: StateSpace

State space with configurations in R^n_joints.

Called FixedBaseSpace as it is the state space of an open kinematic chain with fixed base body.

As the Lie group R^n_joints is equivalent to its own algebra, the state space is simply R^{2 * n_joints}, and the group operation coincides with addition on this vector space. Thus:

n_q == n_v == n_x/2

Inits FixedBaseSpace of prescribed size.

Parameters:: n_joints (int) – number of joints in chain (>= 0)

configuration_difference(q_1, q_2)[source]

Implements configuration offset for a fixed-base rigid chain.

In R^n_joints, this is simply vector subtraction.

Parameters:

q_1 (Tensor) – (*, n_q) “starting” configuration in R^n_joints
q_2 (Tensor) – (*, n_q) “ending” configuration in R^n_joints

Return type:

Tensor

Returns:

(*, n_v) difference of configurations

exponential(q, dq)[source]

Implements exponential perturbation for a fixed-base rigid chain.

In R^n_joints, this is simply vector addition.

Parameters:

q (Tensor) – (*, n_q) “starting” configuration in R^n_joints
dq (Tensor) – (*, n_v) perturbation in R^n_joints

Return type:

Tensor

Returns:

(*, n_q) perturbed configuration

project_configuration(q)[source]

Implements projection onto the fixed-base rigid chain configuration space.

In R^n_joints, this is simply the identity function.

Parameters:: q (Tensor) – (*, n_q) vectors in R^n_joints.
Return type:: Tensor
Returns:: (*, n_q) tensor, q.

zero_state()[source]

Zero element of R^n_x

Return type:: Tensor

class dair_pll.state_space.ProductSpace(spaces)[source]

Bases: StateSpace

State space constructed as the Cartesian product of subspaces.

The product space conforms with the required properties of our Lie group as long as each constituent subspace does as well.

Inits ProductSpace from given factors.

The coordinates of each space in spaces are concatenated to construct the product space’s coordinates, and similar for the velocities.

q_split(q)[source]

Splits configuration into list of subspace configurations.

Return type:: List[Tensor]

v_split(v)[source]

Splits velocity into list of subspace velocities.

Return type:: List[Tensor]

x_split(x)[source]

Splits state into list of subspace states.

Return type:: List[Tensor]

configuration_difference(q_1, q_2)[source]

Constructs configuration difference as concatenation of subspace configuration differences.

Return type:: Tensor

exponential(q, dq)[source]

Constructs perturbed configuration as concatenation of perturbed subspace configurations

Return type:: Tensor

project_configuration(q)[source]

Projects configuration onto Lie group by projecting each subspace’s configuration onto its respective subgroup.

Return type:: Tensor

zero_state()[source]

Constructs zero state as concatenation of subspace zeros

Return type:: Tensor

dair_pll.state_space.centered_uniform(size)[source]

Uniform distribution on zero-centered box [-1, 1]^size

Return type:: Tensor

class dair_pll.state_space.WhiteNoiser(space, unit_noise, variance_factor=1)[source]

Bases: object

Helper class for adding artificial noise to state batches.

Defines an interface for noise distortion of a batch of states. Noise is modeled as a zero-mean distribution on the Lie algebra of the state space, \(\mathbb{R}^{2 n_v}\). Note that this means that velocities receive noise independent to the configuration, and thus may break the finite-difference relationship in a trajectory.

Inits a WhiteNoiser of specified distribution.

Parameters:

space (StateSpace) – State space upon which
unit_noise (Callable[[Size], Tensor]) – Callback, returns coordinate-independent noise of nominal unit size.
variance_factor (float) – Variance of a single coordinate’s unit-scale noise.

ranges: torch.Tensor

space: dair_pll.state_space.StateSpace

variance_factor: float

noise(x, ranges, independent=True)[source]

Adds noise to a given batch of states.

Uses the unit_noise() to get. Optionally, adds identical distortion to each state in the batch, or i.i.d. noise to each state.

Parameters:

x (Tensor) – (*, space.n_x) batch of states to distort with noise.
ranges (Tensor) – (2 * space.n_v,) multiplicative scale of noise.
independent (bool) – whether to independently distort each state.

Return type:

Tensor

Returns:

(*, space.n_x) distorted batch of states.

covariance(ranges)[source]

State covariance matrix associated with noise scale.

Parameters:: ranges (Tensor) – (2 * space.n_v,) multiplicative scale of noise.
Return type:: Tensor
Returns:: (2 * space.n_v, 2 * space.n_v) covariance matrix on state space Lie algebra.

class dair_pll.state_space.UniformWhiteNoiser(space)[source]

Bases: WhiteNoiser

Convenience WhiteNoiser class for uniform noise.

Inits a WhiteNoiser of specified distribution.

Parameters:

space (StateSpace) – State space upon which
unit_noise – Callback, returns coordinate-independent noise of nominal unit size.
variance_factor – Variance of a single coordinate’s unit-scale noise.

class dair_pll.state_space.GaussianWhiteNoiser(space)[source]

Bases: WhiteNoiser

Convenience WhiteNoiser class for Gaussian noise.

Inits a WhiteNoiser of specified distribution.

Parameters:

space (StateSpace) – State space upon which
unit_noise – Callback, returns coordinate-independent noise of nominal unit size.
variance_factor – Variance of a single coordinate’s unit-scale noise.

class dair_pll.state_space.StateSpaceSampler(space)[source]

Bases: ABC

Abstract utility class for sampling on a state space.

Inits StateSpaceSampler on prescribed space.

Parameters:: space (StateSpace) – State space of sampler.

space: dair_pll.state_space.StateSpace

abstract get_sample()[source]

Get sample from state distribution.

Return type:: Tensor
Returns:: (space.n_x,) state sample.

abstract covariance()[source]

Returns covariance of state space distribution.

Interprets the distribution in logarithmic coordinates (the Lie algebra of the state space), and returns a covariance matrix in \(\mathbb{R}^{2 n_v \times 2 n_v}\).

Return type:: Tensor
Returns:: (2 * space.n_v, 2 * space.n_v) distribution covariance.

class dair_pll.state_space.ConstantSampler(space, x_0)[source]

Bases: StateSpaceSampler

Convenience StateSpaceSampler for returning constant state.

Inits ConstantSampler with specified constant state.

Parameters:

space (StateSpace) – Sampler’s state space.
x_0 (Tensor) – (space.n_x,) singleton support of underlying probability distribution.

space: dair_pll.state_space.StateSpace

x_0: torch.Tensor

get_sample()[source]

Returns copy of constant x_0.

Return type:: Tensor

covariance()[source]

Returns zero covariance.

Return type:: Tensor

class dair_pll.state_space.ZeroSampler(space)[source]

Bases: ConstantSampler

Convenience ConstantSampler for returning zero state.

Inits ConstantSampler with specified constant state.

Parameters:

space (StateSpace) – Sampler’s state space.
x_0 – (space.n_x,) singleton support of underlying probability distribution.

class dair_pll.state_space.CenteredSampler(space, ranges, unit_noise=<built-in method randn of type object>, x_0=None)[source]

Bases: StateSpaceSampler

State space sampling distribution centered around specified state.

Implemented by sampling the state, and perturbing it with specified white noise.

Inits CenteredSampler with specified distribution

Parameters:

space (StateSpace) – Sampler’s state space.
ranges (Tensor) – (2 * space.n_v,) multiplicative scale on noise distribution standard deviation.
unit_noise (Callable[[Size], Tensor]) – Callback, returns coordinate-independent noise of nominal unit size.
x_0 (Tensor) – (space.n_x,) center of distribution, around which Lie algebra perturbation is applied by underlying WhiteNoiser.

variance_factor: float

x_0: torch.Tensor

ranges: torch.Tensor

get_sample()[source]

Returns x_0 distorted by white noise.

Return type:: Tensor

covariance()[source]

Returns covariance of underlying noiser.

Return type:: Tensor

class dair_pll.state_space.UniformSampler(space, ranges, x_0=None)[source]

Bases: CenteredSampler

Convenience CenteredSampler for uniform noise.

Inits CenteredSampler with specified distribution

Parameters:

space (StateSpace) – Sampler’s state space.
ranges (Tensor) – (2 * space.n_v,) multiplicative scale on noise distribution standard deviation.
unit_noise – Callback, returns coordinate-independent noise of nominal unit size.
x_0 (Tensor) – (space.n_x,) center of distribution, around which Lie algebra perturbation is applied by underlying WhiteNoiser.

class dair_pll.state_space.GaussianSampler(*args, **kwargs)[source]

Bases: CenteredSampler

Convenience CenteredSampler for Gaussian noise.

Inits CenteredSampler with specified distribution

Parameters:

space – Sampler’s state space.
ranges – (2 * space.n_v,) multiplicative scale on noise distribution standard deviation.
unit_noise – Callback, returns coordinate-independent noise of nominal unit size.
x_0 – (space.n_x,) center of distribution, around which Lie algebra perturbation is applied by underlying WhiteNoiser.