[Dream to Fly: Model-Based Reinforcement Learning for Vision-Based Drone Flight] ([2501.14377] Dream to Fly: Model-Based Reinforcement Learning for Vision-Based Drone Flight) (JAngel Romero, Ashwin Shenai, Ismail Geles, Elie Aljalbout, Davide Scaramuzza / arxiv, 2025)

Summary

• 온보드 카메라의 raw 픽셀 입력만을 사용하여 인간 조종사처럼 드론을 레이스 트랙에서 자율 비행
• Intermediate Representation이나 Explicit State Estimation 없이, dreamerv3를 이용해 pixel-to-command MBRL의 실현 가능성을 입증

---

Takeaways

• 기존의 Model-free 방법론을 이용한 raw pixel → command 학습은 거의 불가능에 가까움
  ⇒ 이전 논문에서 다룬 여러가지 trick이 필요
  

• Model-based RL인 DreamerV3를 활용하여 훨씬 sample efficient 하게 학습 가능
  → world model을 학습해서 imaginary rollout으로 actor / critic 학습 가능
  → 여전히 real world에서 여러 task에 대해 world model을 잘 학습시킬 수 있을지는 의문
  

• Limitation
  • 실제로 학습하고 시간이 많이 든다고 함.
    → 수렴하는데 20 million steps ~ 240 시간 학습
    → RSSM world model을 학습하는데 가장 많은 시간 소요
    
  • computation이 많아 real-time operation이 어려움
    → policy만 돌리는게 아니고 world model을 같이 돌려줘야함
    → offboard가 필수적
    

Preliminary

• MBRL (Model-Based Reinforcement Learning): Transition Dynamics을 학습하는 World Model을 사용하여 샘플 효율성을 높이는 RL 방법론

• DreamerV3: MBRL의 최신 아키텍처로, 세계 모델 학습과 정책 훈련을 분리하여 효율성을 높인 알고리즘

• Model-based reinforcement learning = MBRL

  • world model을 활용해서 환경의 동역학을 근사하는 함수를 학습하는 방법. (~transition dynamics modelling)
    • 강화학습에서 환경은 어떤 action을 취했을 때, 어떻게 state가 변하는지를 알려줌.
    • model-based는 이 환경 자체를 학습
      → 실제 env rollout으로 world model 학습
      → policy 학습은 world model의 imagination rollout에서 수행
      
  • 표본 효율 향상 (sample efficiency)
    • model-free는 실제 환경을 매번 돌려야 하지만, MBRL은 world model 안에서 시뮬레이션처럼 상상(rollout) 가능.
    • 즉, 한 번의 경험으로 world model을 갱신해두면, 그걸로 수천 번의 가상 경험을 재활용할 수 있음.
  • 일반화 능력 향상 (generalization)
    • world model이 드론의 물리 dynamics를 잘 배우면, 비슷한 구조의 다른 task에도 쉽게 적용 가능.
  • 하지만 모델 복잡도가 증가하고, 학습시간이 증가한다는 단점이 있음.
    • world model을 배우면서 control policy를 동시에 학습해야하기 때문

Introduction

• 명시적 상태 추정 의존: 기존 자율비행 drone racing system은 explicit state estimation 사용. (IMU와 onboard 센서 이용 → 인간의 시각 기반 비행과 다름.

• 이전 연구에서 RL을 활용하여 state estimation 없이 곧바로 pixel → action 출력하려는 시도.

  • Demonstrating agile flight from pixels without state estimation, 2024
    • 픽셀을 이진 마스크로 단순화하여 사용 (게이트 seg) → Observation space의 complexity 해소
      
    • But, 게이트 정보 외의 유용한 시각 정보(배경, 중력 방향 등)를 손실
  • Bootstrapping reinforcement learning with imitation for vision-based agile flight, CoRL 2024
    • Imitation learning의 장점을 활용해서 sample efficiency 향상
    • But, extensive bootstrapping 필요

• 정리하자면

  • Model-Free RL (PPO 등): 고차원 픽셀 입력과 연속적인 action space에서 학습이 매우 비효율적이며 수렴에 실패
  • 기존 MBRL: 대부분 시뮬레이션에 국한되거나 상태 추정을 관찰로 사용
  • DreamerV3를 이용항 MBRL로 해결

Methodology

Task = Drone Racing

Goal : raw visual observation을 직접적으로 control command로 mapping하는 정책 학습.

Observation and Action Spaces, Reward Function

• Observation Space

  • raw RGB image
  • explicit state estimation 없고 직접 visual information 주입

• Action Space

  • CTBR format
  • action space $\mathcal A = [-1,1]^4$
    → tanh 함수로 squeeze
    → 실제 collective thrust + body rate limit 값으로 scaling
    

• Reward Function
  • Racing track은 각 게이트 center $g_k$를 선형적으로 연결한 waypoints로 정의
  • reward는 아래와 같이 정의
    $r(k) = \begin{cases} -4.0, \quad \text{if collision;} \\ +10.0, \quad \text{if gate passed;} \\ b_1(||g_k-p_{k-1}||-||g_k-p_k||)-b_2||w_k||, \quad \text{otherwise;} \\ \end{cases}$
  • $b_1(||g_k-p_{k-1}||-||g_k-p_k||)$ : drone이 target gate에 가깝게 이동하도록
    ☑️ 이전 timestep보다 지금 timestep에서 다음 게이트까지 거리가 더 가깝도록
    
  • $b_2||w_k||$ : 너무 큰 body rate를 패널티
  ☑️ 여기서 하나 알고가야하는 것은 reward 정의에서 실제 어떤 trajectory를 따라가라고 강요하지 않음
  → 단순히 우리는 목적만 부여하고 agent가 가장 효율적인 trajectory를 찾도록함
  

Model-based Reinforcement Learning : DreamerV3

• DreamerV3 : off-policy model-based reinforcement learning algorithm
  ☑️ 이 알고리즘의 핵심은 미래를 상상하는 능력을 가진 world model을 학습하는 것.

• 크게 world model과 actor-critic policy로 구성됨

  → agent가 환경과 상호작용하는 동안 replay buffer를 활용해 두 모듈이 번갈아가면서 학습

  → world model 학습할때는 미리 만들어둔 replay buffer로 학습, actor-critic 학습할때는 world model freeze

  

World model training

• World model은 compact latent space에서 state transition dynamics를 기술

• latent state는 $s_k=(h_k,z_k)$로 사용하고, state transition probability를 $P(s_{k+1}|s_k,a_k)$로 표현

• World model은 기본적으로 sensory input을 low-dim representation으로 encoding하며, agent의 action을 기반으로 미래의 latent state와 reward를 예측

• world model은 Recurrent State Space Model로 구현

• Encoder

  • Encoder network는 raw obs $x_k$를 stochastic latent representation $z_k$로 mapping

    $z_k \sim q_\phi(z_k | h_k,x_k)$

    ☑️ stochastic representation → 네트워크가 평균, 분산 출력하고 그 분포에서 sampling

• Recurrent Sequence Model
  • recurrent state $h_k$를 갖는 recurrent sequence model이 latent representation의 sequence를 예측
  • 이전 latent state와 이전 action을 통해 현재 recurrent state 구함
    $h_k=f_\phi(h_{k-1},z_{k-1},a_{k-1})$
    ☑️ $h_t$는 과거 전체 정보를 요약해주는 state 라고 생각 ☑️ 이후 $h_t$를 통해서 representation $\hat{z}$를 예측할 수 있음 → imagination rollout을 위해
    

• Dynamics, Reward, and Continue Prediction

  • dynamic predictor는 recurrent state $h_k$를 통해 stochastic state $\hat{z_k}$를 예측

  • reward, continue predictor는 어떤 latent state $s_k=(h_k,z_k)$가 주어졌을 때, reward $r_k$와 $c_k \in \{0,1\}$를 각각 예측

    $\text{Dynamic predictor: } \hat{z_k}\sim p_\phi(\hat{z_k}|h_k) \\ \text{Reward predictor: } \hat{r_k}\sim p_\phi(\hat{r_k}|h_k,z_k) \\ \text{Continue predictor: } \hat{c_k}\sim p_\phi(\hat{c_k}|h_k,z_k)$

• Decoder

  • latent representation으로부터 원래의 observation을 reconstruct

  • latent 변수 $z_k$가 환경으로부터 필수적인 정보를 포함하도록

    $\hat{x}_k\sim p_\phi(\hat{x_k}|h_k,z_k)$

• 이 모델들은 각각 다른 loss를 통해 학습됨.

• input $x_{1:T}$, action $a_{1:T}$, reward $r_{1:T}$, continuous flags $c_{1:T}$의 sequence가 주어진 상황에서 아래의 loss를 최소화하도록 end-to-end로 학습된다.

• pred loss는 decoder와 reward predictor, continue predictor를 학습.
  $\mathcal L_{pred}(\phi)=-lnp_\phi(x_t|z_t,h_t)-lnp_\phi(r_t|z_t,h_t)-lnp_\phi(c_t|z_t,h_t) \\$
  ☑️ state가 주어졌을 때 실제로 관측된 값 $x, ...$에 높은 확률을 주도록 모델을 학습하겠다는 의미
  

• dyn loss는 sequence model이 다음 representation을 잘 예측할 수 있도록 학습.
  → $p_\phi(z_t|h_t)$ 와 $q_\phi(z_t | h_t,x_t)$ 사이의 KL divergence를 줄여서 학습
• rep loss는 representation이 dynamics로부터 더 잘 예측될 수 있도록 학습.
  $\mathcal L_{dyn}(\phi)=\text{max }(1,KL[sg(q_\phi(z_t|h_t,x_t))|| p_\phi(z_t,h_t)]) \\ \mathcal L_{rep}(\phi)=\text{max }(1,KL[q_\phi(z_t|h_t,x_t)|| sg(p_\phi(z_t,h_t))])$

Actor-Critic Training

• Actor와 Critic은 학습된 world model로부터 생성된 상상의 trajectory를 활용하여 학습
  → 실제 환경과 상호작용없이 학습할 수 있음
  

• 실제 학습은 아래와 같이 이루어짐.

  1. 어떤 starting state representation $(h_{k-1},z_{k-1})$이 주어졌을 때, actor가 action $a_{k-1}$출력

  2. world model은 actor가 출력한 action $a_{k-1}$과 state $(h_{k-1},z_{k-1})$을 통해 현재의 state $(h_{k},z_{k})$를 계산하고, 이를 기반으로 reward $r_k$와 continuation flag $c_k$ 계산

  3. 반복적으로 수행해서 trajectory 생성

     → model states $(h_{1:T},z_{1:T})$, actions $a_{1:T}$, rewards $r_{1:T}$, flags $c_{1:T}$

     ☑️ 실제 환경과 interaction한 trajectory는 아니고, 상상으로 얻은 trajectory임

  4. Critic은 bootstrapped $\lambda$-return 분포을 예측하여 이 trajectory들을 평가하는 방법을 학습

  5. Actor는 이 $\lambda$-return을 최대화하도록 학습.

(참고) bootstrapped $\lambda$-return의 핵심 아이디어

• 보통 TD 학습이라고 하면 아래와 같음
  $V(s_t) \leftarrow r_t + \gamma V(s_{t+1})$
  → 한 step 뒤만 bootstrap → bias 큼
  
• n-step return은 TD(0)와 MC의 중간
  $G_t^{(n)}=r_t+\lambda r_{t+1}+...+\lambda^{n-1}r_{t+n-1}+\lambda^n V(s_{t+n})$
  → n이 커질수록 bias 낮아지지만, variance 커짐 → 여전히 bias-variance를 자유롭게 조절하긴 어려움
  

• $\lambda$-return : 그러면 모든 n-step return을 가중 평균

  $R_t=(1-\lambda)G_t^{(1)}+\lambda(1-\lambda)G_t^{(2)}+ \lambda^2(1-\lambda)G_t^{(3)}+ ...$

  → $\lambda=0$이면 TD(0), $\lambda=1$이면 MC

  → bias-variance tradeoff를 연속적으로 부드럽게 조절 가능

  ⚠️ 하지만 우리의 상황에서는 $n$을 충분히 늘리기 어려움

  • world model은 짧은 horizon (T=16)까지만 안정적
  • horizon 이후의 reward 값을 얻을 수 없음.

• bootstrapped $\lambda$-return

  $R_t^\lambda = r_t+\gamma c_t((1-\lambda)v_t+ \lambda R^\lambda_{t+1})$

  → horizon 내부에서는 각 time의 reward 값을 사용하다가, horizon 마지막에서 critic이 주는 value로 이후 값들을 bootstrapping

  → world model의 짧은 horizon에도 불구하고 long-term return을 예측 가능하게 함.
  

Quadrotor Dynamics

• 쿼트로터의 state를 정의
  $\vec{x}=[\vec{p}_I, \vec{q}_{IB},\vec{v}_I]^T$

• dynamics 식은 아래와 같이 정의

Experiments

• 구현 : dreamerv3-torch (Large Network Config) + stable-baselines3 (baseline PPO)

• 시뮬레이션 : Flightmare + Agilcious + Habitat Simulator

• 장비 : Quadro RTX 8000 한대

• 이미지는 64x64 RGB

Results

• 기본적으로 PPO로 학습한 모델은 학습자체가 거의 불가능함.
  ➡️ 예전에 다룬 여러 논문에서도 raw pixel → control 학습은 불가능하기 때문에 imitation learning을 한다던가, intermediate representation을 활용한다던가 했음.

• reconstruction된 그림을 보면 학습이 진행되면서, 실제로 제어에 필요한 정보들을 담게 된다.
• 그런데 reconstruction 자체가 목적은 아님. 제약의 관점에서 봐야함.
  → latent만 보아도 어느정도 observation를 설명할 수 있어야 한다 정도.
  

Perception-aware emergent behaviour

• 학습 이후에 agent는 이동할때 카메라가 gate를 향하는 방향으로 이동하는 경향을 보임
  → reward function에 perception 관련 제약을 주지 않았는데도 자연스레 학습
  → 이것은 바로 위 그림에서 reconstruction했을 때, 뒷 배경보다는 gate에 집중하고 있음

• 이를 검증하기 위해 racing track에 두 개의 추가적인 gate를 설치했음 (실제로는 지나지 않아도 되는 게이트. reward에 포함 안됨)
• 검은색 화살표가 드론 카메라의 방향을 의미하는데, 카메라 시야가 실제 가야하는 게이트 + 추가적인 가짜 게이트에 골고루 분산된다는 것을 확인할 수 있음
  → 드론 입장에서 강한 시각적 특징을 가지는 게이트를 인지하고 이를 비행 등에 활용한다고 해석할 수도 있음
  

Real-world Setup

• beta-flight활용하였고, offboard computer에서 계산한 collective thrust + bodyrate command를 직접 fc가 받아서 제어.

  • 드론에 RF receiver 달아서 조종 신호 받으며, 이 신호가 바로 FC로 전달되는 구조

    = 인간 레이서와 동일한 observation & control input modality

• real-world deploy에는 HIL setup 활용
  → observation인 이미지는 habitat simulator에서 얻어지고, 이를 통해 얻어진 command는 실제 드론에 전송.
  → real world dynamics에서 실험한 것
  ☑️ perception은 여전히 simulation이긴 함. end2end real-world x
  • real dynamics + real control만 검증하는 수준