Equivariant Reinforcement Learning for Quadrotor UAV

Overview

The paper exploits the rotational symmetry of a quadrotor about the vertical gravity axis to improve the sample efficiency of reinforcement learning for low-level control of the quadrotor. The policy maps the quadrotor state directly to its four motor thrusts for position stabilization. A standard policy treats horizontally rotated versions of the same flight condition as different inputs, even though they require the same motor commands when actions are represented in the body-fixed frame.

The authors map every state to a canonical representative of its rotational equivalence class before passing it to the actor and critic. This removes one continuous input dimension and makes the learned policy and value functions invariant to arbitrary rotations about the gravity direction. The approach is applied to TD3 and SAC and improves their convergence in simulation.

$S^1$ Symmetry of Quadrotor

The quadrotor state and action are

$$s=(x,v,R,\Omega), \qquad a=(T_1,T_2,T_3,T_4),$$

where $x,v\in\mathbb{R}^3$ are expressed in the inertial frame, $R\in SO(3)$ maps body coordinates to inertial coordinates, and $\Omega\in\mathbb{R}^3$ is expressed in the body frame.

The paper shows that the rotations around the vertical gravity axis result in equivalent states:

In the image, gravity points into the screen and $\vec{e}_1$ and $\vec{e}_2$ form the world horizontal plane (parallel to the ground).

The body-frame quantities $\Omega$ and $a$ remain numerically unchanged because their reference frame rotates with the vehicle. Their corresponding inertial-frame quantities are still rotated through $R$. In other words, in the image, the rotor that is on the inner side of the circle remains on the inner side after rotation; so, the action vector does not change.

Writing the quadrotor dynamics as $\dot{s}=F(s,a)$, they satisfy

$$F\circ g_\theta=g_\theta\circ F.$$

where $g_{\theta}$ is the rotation of states and actions by an angle $\theta$ starting from $\vec{e}_1$.
Thus, rotating a valid state-action trajectory about $\vec{e}_3$ produces another valid trajectory. Two trajectories are equivalent when one is a vertical-axis rotation of the other, and the equivalence class

$$[s,a]=\{g_\theta(s,a):\theta\in S^1\}$$

contains all horizontally rotated versions of the same control problem.

Equivariant RL

For a group action $g$, assume that the dynamics are equivariant and the reward is invariant:

$$F\circ g=g\circ F, \qquad r\circ g=r.$$

Given a policy $\pi$, its transformed policy is

$$\widetilde{\pi}(\widetilde{s}) =g\pi(g^{-1}\widetilde{s}).$$

The corresponding value functions satisfy

$$V_\pi=V_{\widetilde{\pi}}\circ g,$$

and the optimal policy is equivariant:

$$\pi^*\circ g=g\circ\pi^*.$$

For the selected quadrotor symmetry, the group action on motor thrusts is the identity. Therefore, the optimal policy, value function, and action-value function are invariant across each state orbit:

$$\pi^*(g_\theta s)=\pi^*(s),$$ $$V^*(g_\theta s)=V^*(s), \qquad Q^*(g_\theta s,a)=Q^*(s,a).$$

Rather than constructing specialized equivariant layers, the authors choose one representative from each state orbit. For a state with horizontal position $(x_1,x_2)$, they use

$$[\theta]=-\operatorname{atan2}(x_2,x_1), \qquad [s]=g_{[\theta]}s.$$

This rotates the state so that

$$[x]_2=0.$$

The same rotation is applied to $x$, $v$, and $R$, while $\Omega$ remains unchanged. Since $[x]_2$ is always zero, it is removed from the neural network input. The actor and critic receive

$$([x]_1,[x]_3,[v],[R],[\Omega])\in\mathbb{R}^{17}$$

instead of the original 18-component matrix representation of the state.

The actor maps the canonical state directly to body-frame motor commands,

$$a=\pi([s]),$$

and the critic estimates

$$Q([s],a).$$

No action transformation is needed because the motor commands are invariant under the selected group action. Consequently, all states differing only by a vertical-axis rotation generate the same network input, so experience is shared across the entire $S^1$ equivalence class.