Shared model with rotation-manifold projection layers.
Control modes and their action dimensions::
cmd_ctatt_euler → 4D [thrust, roll, pitch, yaw] cmd_ctatt_rotvec → 4D [thrust, rx, ry, rz] cmd_ctatt_quat → 5D [thrust, qx, qy, qz, qw] cmd_ctatt_rotmat → 10D [thrust, r00..r22]
Projection layers ensure the rotation component of the policy output lies on a valid SO(3) manifold. Available variants per rotation type:
============== ============================================================
Euler / rotvec EulerProjection (tanh clamp)
Quaternion QuatProjection (tanh + L2-normalize)
QuatPlusProjection (tanh + sigmoid w + L2-normalize)
QuatOffsetProjection (offset by identity quat)
QuatPlusOffsetProjection (offset + sigmoid w)
Rotation RotmatProjection (tanh + SVD-orthogonalize)
matrix RotmatOffsetProjection (offset by identity matrix)
============== ============================================================
Select a non-default projection by passing projection_cls when
constructing :class:SharedAttitude. The config module re-exports the
projection classes for compatibility with existing experiments.
Classes:
| Name | Description |
|---|---|
EulerProjection |
Bounded projection for euler-angle or rotation-vector actions. |
QuatProjection |
L2-normalize the 4-D quaternion at indices [1:5]. |
QuatPlusProjection |
Quaternion projection with double-cover elimination. |
QuatOffsetProjection |
Quaternion projection centred at the identity rotation. |
QuatPlusOffsetProjection |
Quaternion projection combining offset and double-cover elimination. |
RotmatProjection |
SVD-orthogonalize the 9-D rotation-matrix at indices [1:10]. |
RotmatOffsetProjection |
Rotation-matrix projection centred at the identity matrix. |
SharedAttitude |
Shared actor-critic with rotation-manifold policy projection. |
Bases: Module
Bounded projection for euler-angle or rotation-vector actions.
Applies tanh to clamp all components to [-1, 1] so that
the environment action scaling (e.g. multiplying by
attitude_range) stays within the expected range. No geometric
constraint is needed beyond clipping.
References
so3_primer.rotations.modules.ddpg.activations — plain
nn.Tanh() used for euler / tangent quadrotor actors.
Methods:
| Name | Description |
|---|---|
forward |
Clamp x to |
forward(x: Tensor) -> torch.Tensor
Clamp x to [-1, 1] via tanh.
Bases: Module
L2-normalize the 4-D quaternion at indices [1:5].
Applies tanh to bound all components, then normalizes the
quaternion slice to unit length. The thrust channel (index 0) is
passed through unchanged.
This is the simplest quaternion projection — the double-cover
ambiguity (q and -q represent the same rotation) is left
for the policy to resolve.
References
so3_primer.rotations.modules.ddpg.activations.QuatQuadrotor
Methods:
| Name | Description |
|---|---|
forward |
Normalize quaternion slice of x. |
forward(x: Tensor) -> torch.Tensor
Normalize quaternion slice of x.
Bases: Module
Quaternion projection with double-cover elimination.
The scalar component w is mapped through a sigmoid channel
(tanh → [0, 1]) so that the output quaternion always satisfies
w ≥ 0, breaking the q ≡ -q ambiguity. The vector
components (x, y, z) use standard tanh. This mirrors
so3_primer.rotations.modules.ddpg.activations.QuatPlusQuadrotor
and so3_primer.rotations.modules.ppo.std.QuatPlusQuadrotorActor.
Methods:
| Name | Description |
|---|---|
forward |
Match the primer quadrotor quat-plus head. |
forward(x: Tensor) -> torch.Tensor
Match the primer quadrotor quat-plus head.
QuatOffsetProjection()
Bases: Module
Quaternion projection centred at the identity rotation.
Adds the identity quaternion [0, 0, 0, 1] to the policy output,
clamps to [-1, 1], then L2-normalizes. This centres the action
distribution on zero-rotation so that a zero policy output commands
hover attitude.
References
so3_primer.rotations.modules.ddpg.activations.QuatQuadrotorOffset
Register the identity-quaternion offset buffer.
Methods:
| Name | Description |
|---|---|
forward |
Offset x by identity quaternion, then normalize. |
forward(x: Tensor) -> torch.Tensor
Offset x by identity quaternion, then normalize.
QuatPlusOffsetProjection()
Bases: Module
Quaternion projection combining offset and double-cover elimination.
Offsets the vector part by zero (adds [0, 0, 0]), maps w
through a sigmoid channel, then L2-normalizes.
References
Combines so3_primer QuatPlusQuadrotor (sigmoid w)
with QuatQuadrotorOffset (identity offset).
Register the zero-offset buffer for the vector part.
Methods:
| Name | Description |
|---|---|
forward |
Offset vector part by zero, force |
forward(x: Tensor) -> torch.Tensor
Offset vector part by zero, force w ≥ 0, then normalize.
Bases: Module
SVD-orthogonalize the 9-D rotation-matrix at indices [1:10].
Reshapes the matrix slice to (..., 3, 3), applies tanh to
bound elements, then computes the least-squares orthogonal
Procrustes projection via SVD. A determinant correction ensures
the result lies in SO(3) (det=+1) rather than O(3).
The thrust channel (index 0) is passed through unchanged.
References
so3_primer.rotations.modules.ddpg.activations.MatrixQuadrotor
Methods:
| Name | Description |
|---|---|
forward |
SVD-project the matrix slice of x onto SO(3). |
forward(x: Tensor) -> torch.Tensor
SVD-project the matrix slice of x onto SO(3).
RotmatOffsetProjection()
Bases: Module
Rotation-matrix projection centred at the identity matrix.
Adds the flattened 3×3 identity to the policy output, clamps to
[-1, 1], then SVD-projects onto SO(3). A zero policy output
therefore commands hover attitude.
References
so3_primer.rotations.modules.ddpg.activations.MatrixQuadrotorOffset
Register the identity-matrix offset buffer.
Methods:
| Name | Description |
|---|---|
forward |
Offset x by identity matrix, then SVD-project. |
forward(x: Tensor) -> torch.Tensor
Offset x by identity matrix, then SVD-project.
SharedAttitude(observation_space, state_space, action_space, device, clip_actions=False, clip_log_std=True, min_log_std=-20, max_log_std=2, reduction='sum', projection_cls: type[Module] | None = None)
Bases: GaussianMixin, DeterministicMixin, Model
Shared actor-critic with rotation-manifold policy projection.
The trunk is a 2-hidden-layer MLP (128 units, ELU activation)
shared between policy and value heads. The policy mean passes
through an optional :class:nn.Module projection layer that
enforces the rotation component to lie on the correct manifold
(e.g. unit quaternion, orthogonal matrix).
The projection is selected by action_dim via
:data:_PROJECTION_MAP unless an explicit projection_cls is
provided.
Initialize shared trunk, heads, and optional projection layer.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
|
Environment observation space. |
required | |
|
Environment state space. |
required | |
|
Environment action space. |
required | |
|
Torch device for parameters and computation. |
required | |
|
Whether to clip actions to the action space
bounds (passed to :class: |
False
|
|
|
Whether to clip the log-standard-deviation
(passed to :class: |
True
|
|
|
Minimum log-std value. |
-20
|
|
|
Maximum log-std value. |
2
|
|
|
Reduction mode for the log-probability. |
'sum'
|
|
|
type[Module] | None
|
Optional override for the rotation-manifold
projection layer. When |
None
|
Methods:
| Name | Description |
|---|---|
act |
Dispatch to the appropriate mixin based on role. |
compute |
Compute policy mean or state value. |
act(inputs: Any, role: str) -> Any
Dispatch to the appropriate mixin based on role.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
|
Any
|
Model inputs dict (must contain |
required |
|
str
|
|
required |
Returns:
| Type | Description |
|---|---|
Any
|
Action outputs for the given role, or |
compute(inputs: Any, role: str) -> Any
Compute policy mean or state value.
For the "policy" role the mean passes through the
rotation-manifold projection (if configured) and the trunk
features are cached for reuse by the value head.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
|
Any
|
Model inputs dict (must contain |
required |
|
str
|
|
required |
Returns:
| Type | Description |
|---|---|
Any
|
|
Any
|
or |