[Learning Quadrotor Control From Visual Features Using Differentiable Simulation] ([2410.15979] Learning Quadrotor Control From Visual Features Using Differentiable Simulation) (Johannes Heeg, Yunlong Song, Davide Scaramuzza / ICRA, 2025)

Summary

• Vision-based UAV control에서 sample inefficiency 문제를 해결하기 위한 방법으로 differentiable simulation을 제안
  • surrogate dynamics model을 활용해서 backpropagation 속도 향상
  • state representation learning을 활용해서 convergence 속도 향상
• 기존 model-free approach에 비해서 학습 효율성이 크게 향상

---

My Insights / Takeaways

• Differentiable simulation을 UAV RL control에 활용하는 방법에 대해 공부
  → forward / backward process가 어떻게 되는지 이해
  → backward process는 조금 더 단순화해도 학습에 지장 없고, 훨씬 더 빠르게 학습 가능
  ⇒ 기존 model-free 기반 방법론과 다르게, sampling 기반이 아니고 직접적으로 gradient를 구해서 업데이트하므로 sample efficiency 훨씬 높음
  

• 시뮬레이션 구현 난이도가 높을 것 같음 → 한번 살펴보면 좋을 것 같음.

• 이론적으로는 깔끔한거 같은데 Real world deployment는 실제로 의미있는 것인지 의문

  • 해당 연구에서는 visual feature point coord를 기반으로 stabilize했는데,
    → 실제로는 이 point를 quadrotor에 facing하고 있다는 보장도 없고
    → pixel coord를 추출하는 model의 정확도에 대한 보장도 없어서 ..
    

• State representation pretraining은 추후에도 쓸만한 방법인 것 같음.
  • implicit한 state estimation을 먼저 조금 배워두는 느낌으로

---

Introduction

Problem definition

• Vision-based RL Control의 가장 큰 문제 → Sample inefficiency

  • Input dimension이 높아서 더 많은 sample 필요
  • NN이 implicit한 state representation과 control law를 동시에 학습해야함
  • 실제로 이전 pure-vision-based 연구들은 한 모델 학습하는데 몇 일이 걸릴 수 있다고 함.

• Differentiable simulation을 활용하면 sample efficiency를 크게 향상시킬 수 있음

  • Dynamics의 derivative를 활용하여 policy gradient를 sampling하지않고 직접 계산
    → parameter update의 variance를 크게 줄여서 sample efficiency를 높임
    

Contribution

Vision based quadrotor control 학습을 위해 differentiable simulation을 활용하는 방법 제안

• 이 논문에서는 quadrotor를 던졌을 때, camera로부터 얻어지는 visual feature만을 활용해서 stabilize하는 문제를 풀어보려고 함.

1. state based control task와 vision based control task 모두 실험

   → model-free RL (PPO)와 비교하여 월등한 성능

2. Backward pass에서 단순화된 surrogate model을 활용하면, gradient 계산을 훨씬 빠르게 수행할 수 있음

   → sample efficiency 향상

3. Control policy를 보조적인 representation task로 pretraining하면 전체적인 학습을 가속화 / 안정화 가능

Method

• Differentiable simulation의 핵심은 dyanmics를 통해 objective function으로부터 policy parameter까지 직접 gradient를 back-propagation 할 수 있다는 것

---

참고) model-free & model-based policy gradient

• Model-free policy gradient에서는 사실상 world의 dynamics transition을 모른다고 가정하기 때문에, reward를 미분하는게 아니고, policy objective를 잘 정의해서, reward로 weighted된 log-probability gradient로 policy update
  = policy gradient theorem
  → reward를 sampling해서 log-probability를 좋은 행동을 하는 방향으로 조금씩 조작해주는 것
  $\nabla_\theta J(\theta)=\mathbb E [\nabla_\theta log \pi_\theta (a|s) \cdot \hat{A}]$

• Model-based policy gradient에서는 실제로 reward를 policy 파라미터로 미분해서 직접적으로 파라미터 업데이트

  → dynamics가 파라미터로부터 reward까지를 연결해주기 때문에 가능한 일 (아래와 같이)

  $\frac{\partial r }{\partial \theta} = \frac{\partial r }{\partial x} \frac{\partial x }{\partial u} \frac{\partial u }{\partial \theta}$

⚠️ 그리고 여기서 model-based는 DreamerV3와 같은 실제 model-based RL의 개념은 아님.

• 이 논문에서 model-based라 하는 것은 world physics를 알고 있어서, 물리적 state를 가지고 기술된 dynamics를 미분해서 학습 또는 최적화를 수행한다는 것
• 보통 RL에서 model-based라 하는 것은 실제 world physics를 모르기 때문에, 이를 근사하는 미분 가능한 world 모델을 데이터로부터 학습하겠다는 것

---

이제 differentiable simulation이 policy optimization에 어떻게 활용되는지 살펴보자.

Policy Optimization via Differentiable Simulation

• 먼저 quadrotor는 discrete-time dynamic system으로 정의

• system state : $x_t \in \mathcal X$

• control input : $u_t \in \mathcal U$

  → system dynamics는 아래와 같이

  $f : \mathcal X \times \mathcal U \to \mathcal X, \; \text{s.t.} \; x_{t+1} = f(x_t,u_t)$

• observation $o_t \in \mathcal O$

  → observation은 각 time step마다 observation model을 통해 state로 만들어진다.

  $h: \mathcal X \to \mathcal O,\; \text{s.t. } o_t=h(x_t)$

• 또한 각 time step마다 reward를 받는다. → 여기선 reward도 $x,u$의 함수
  $r_t= r(x_t,u_t)$

• control policy는 deterministic differentiable function (NN)
  $u_t=\pi_\theta(o_t)$

• Objective function $R(\theta)$는 policy parameter가 주어졌을 때 누적 reward값으로, optimal parameter를 찾기 위해 gradient ascent로 maximize
  $\theta^*=\underset{\theta}{\text{argmax}} \; R(\theta) \\ R(\theta)=\sum^{N-1}_{t=0}r(x_t,u_t)=\sum^{N-1}_{t=0}r(x_t,\pi_\theta(o_t))=\sum^{N-1}_{t=0}r(x_t,\pi_\theta(h(x_t))) \\ \theta_{k+1} \leftarrow \theta_k + \alpha \nabla_\theta R(\theta_k)$
  ✅ model-free RL과 핵심 차이는 이 objective function이 sampling기반이 아니라 (expection이 아님) 실제 결정론적인 함수라는 것
  

• Back-propagation Through Time (BPTT)

  • 그래서 이 $R(\theta)$의 구조를 통해 chain rule을 적용하여 policy gradient를 구하면

    $\nabla_\theta R(\theta)= \frac{1}{N}\sum^{N-1}_{t=0} (\sum^t_{i=1}\frac{\partial r_t}{\partial x_t} \prod_{j=i}^t (\frac{dx_j}{dx_{j-1}})\frac{\partial x_i}{\partial \theta}+ \frac{\partial r_t}{\partial u_t}\frac{\partial u_t}{\partial \theta})$

  • 여기서 $dx_j/dx_{j-1}$은 dynamic system의 jacobian

    → 과거 state의 변화가 현재 state에 얼마나 영향을 주는지를 의미

  ✅ BPTT의 policy gradient를 두 항으로 나눠 해석하자면 아래와 같다고 생각

  • $\sum^t_{i=1}\frac{\partial r_t}{\partial x_t} \prod_{j=i}^t (\frac{dx_j}{dx_{j-1}})\frac{\partial x_i}{\partial \theta}$ ⇒ state를 통한 효과
    : 파라미터를 바꾸면 action이 바뀌고, 그게 dynamics를 통해 state를 바꾸고, 그 state가 reward에 영향을 준다
  • $\frac{\partial r_t}{\partial u_t}\frac{\partial u_t}{\partial \theta}$ ⇒ action을 통한 효과
    : 파라미터를 바꾸면 action이 바뀌고, 그 action이 reward에 직접 반영되는 경우 (e.g., control effort penalty)

Quadrotor dynamics

• Simple quadrotor dynamics를 아래와 같이 표현

• $p,R,v$는 world frame position, orientation, velocity

• $w$는 body frame angular velocity

• $c$는 collective thrust

• full dynamics는 motor, drag, delay, 등을 모두 포함한 수식

Jax-based Differentiable Simulator

• Differentiable simulator는 Flightmare 시뮬레이터와 비슷하게 air-drag, low-level control 구조, motor speed, transmission delay를 simulate
  → 최대한 현실적으로 만들어서 real-world zero-shot 가능하도록
  

• 모든 simulation framework는 JAX를 활용하여 구현

  → 시뮬레이션 + BPTT를 포함한 함수 자체를 실행 시점에 컴파일 (JIT)해서 GPU에서 바로 실행되도록 함.

  ⚠️ 실제 구현이 어떻게 되어있는지는 코드를 봐야 알거같은데, 아마 JAX로 컴파일한 어떤 함수에서 시뮬레이션 rollout → back propagation까지 해서 gradient구해주고, 우리는 외부에서 모델만 계속 업데이트해주는 구조일 것 같음

     # 이런식이려나
     @jax.jit
     def loss_and_grad(theta, x0):
         R = rollout(theta, x0)
         return R, jax.grad(R)(theta)

→ 그래서 매우 빠른 실행이 가능함 (1초에 몇백만 steps까지)

→ 또한 위 표에서 gradient 계산을 빠르게 해주는 또 다른 요소가 아래 나오는 fast surrogate gradients

Fast Surrogate Gradients

• Differentiable simulator는 위에서 설명한 것처럼 fully differentiable.
• 하지만, full low-level control stack까지 모두 미분하는것은 바람직하지 않음 → computational graph가 너무 expensive해짐
  

➡️ 그래서 forward path에서는 full model $f$를 활용하고, backpropagation에서는 simpler model $\hat{f}$를 활용

• forward path → 실제 시뮬레이션에서 활용되는 모델 (flightmare 등)
  $x_{t+1}=f(x_t,u_t)$

• backward path → 최소한의 물리 관계만 유지 = 위에 Quadrotor dynamics에 표현된 식
  $\frac{\partial x_{t+1}}{\partial x_t} := \left.\frac{\partial \hat{f}} {\partial x}\right|_{(x_t,u_t)}, \frac{\partial x_{t+1}}{\partial u_t} := \left.\frac{\partial \hat{f}} {\partial u}\right|_{(x_t,u_t)}$

Pretraining on State Representation Learning

• 지금까지는 differentiable simulation과 그 gradient 계산 방법 등에 대해서 알아보았음.
• 근데 실제로 vision-based control을 생각해보면, policy training 전에, visual feature representation을 학습하는 것 자체가 또 어려운 task임.

  1. visual feature → command mapping 학습이 훨씬 어려움.

     • state-based : state -> control
     • vision-based : image -> (implicit state representation) -> control

     ⇒ NN은 이미지에서 암묵적으로 state를 추정해야함.

  2. Observation model 떄문에 gradient가 우회해서 흐르게 됨

     • state-based : $x \to f \to u \to r$
     • vision-based : $x \to o \to f \to u \to r$

➡️ Neural network를 state representation에 대해 Pretraining하자

1. 데이터 수집

   • control 자체는 중요하지 않고, random 초기화된 policy로 데이터 $(x_t,o_t)$ 수집

2. State representation learning

   • NN은 visual observation으로부터 quadrotor state를 예측하도록 학습한다.
   $\frac{1}{N} \sum_{(x_i,o_i)\in \mathcal D} || x_i-\psi_\theta(o_i)||^2$

3. 학습된 weight를 policy network $\pi_\theta$의 초기값으로 사용한다.

---

→ policy network 내부적으로 state space를 latent로 가지게 되어서 vision → state mapping이 어느정도 가능

→ 이 상태에서 control만 추가로 잘 가르치면 되는 식.

→ 이러한 방법은 RL의 convergence에 큰 도움을 준다고 함.

“Integrating State Representation Learning Into Deep Reinforcement Learning” (RAL 2018) 참고

---

Training Details

• Policy

  • two hidden layer MLP
  • size → 512 for state-based, 1024 for vision-based

• Action space

  • body rate 와 collective thrust
  • on-board controller는 이를 받아서 execute (~50Hz)
    → FC에서 state estimation이 필요 없음
    → RL이 stabilization까지 모두 수행하는 구조

• Observation space

  • state-based에서는 $o=[p^T,vec(R)^T,v^T]$

  • vision-based에서는 아래 그림과 같이 바닥에 놓인 7개의 feature point pixel coord
    

  • 한 frame만 가지고는 vel, acc를 추정할 수 없으므로, 5개의 이전 frame과 3개의 이전 action을 추가로 observation으로 활용

  • Camera model도 differentiable model 사용

• Reward function

  • Quadrotor stabilization이라는 것은 원하는 위치 $p_{des}$에서 드론을 hovering 시키는 것
    $r_t=r_t^{pos}+r_t^{vel} + r_t^{act}$
    • $r_t^{pos}=-0.2 \cdot L_H(5 \cdot (p_t-p_{des}))$
    • $r_t^{vel}=-0.1 \cdot L_H(v_t)-0.1 \cdot L_H(w_t)$
    • $r_t^{act}=-0.5 \cdot L_H(u_t-u_{hover})-0.01 \cdot L_J(u_t-u_{t-1})$
  • $L_H$는 huber loss
  • $u_{hover}=[9.81,0,0,0]^T$

Experiments

• 여기서 task는 random condition으로부터 stabilize하는 것

• state-based RL과 vision-based RL로 각각 실험

• PPO와 BPTT를 비교

• 모든 학습은 Nvidia Titan RTX (24GB VRAM)에서 진행

Sample efficiency and Training time

Learning State-based Control

• BPTT는 훨씬 적은 parallel environment로도 PPO와 비슷한 성능 도출
• 같은 reward를 도달하기 위한 sample 수, time에서도 우위

Feature-based Control without state estimation

• BPTT는 state-based control에 유사한 성능을 낼 수 있음.

Real world transfer

• 실제 드론으로 HITL 실험
• 실제 드론의 pose는 MOCAP으로 측정
• actor는 가상의 observation을 받음 (아래 그림 오른쪽 아래 feature point observation 처럼) → feature detection model 없이도 control 성능 측정 가능
  

• observation model은 sim2real gap 없지만, MOCAP 측정 오류는 존재 가능한 상황

Simple model gradient computation

• Surrogate function을 사용했을때, training time을 줄이면서, sample efficiency나 성능에 손해를 보지 않는다

Pretraining on state representation

• vision-based control task에서 state representation learning을 해줬을 때 성능이 훨씬 좋다.
  → 네트워크가 vision → state mapping에 대한 이해를 가지게 됨
  → gradient quality를 높임