WMP(World Model-bassed Perception for Visual Legged Locomotion) 리뷰 및 테스트

World Model-based Perception for Visual Legged Locomotion 논문 리뷰

Abstract

다양한 지형에서 Legged Robot이 이동하는 것은 어려운 문제이며, 이때 proprioception과 vision를 통한 정확한 환경 인지가 필수적이다. 그러나 고차원 시각 정보에서 직접 학습하는 것은 데이터 효율성이 낮고 복잡한 과정이 요구된다.
이를 해결하기 위해, 기존 연구에서는 privileged information를 활용한 teacher policy를 먼저 학습한 후, student policy가 시각 정보를 기반으로 teacher의 행동을 모방하는 방식을 사용. 하지만 이러한 imitation framework은 입력 간 정보 차이로 인해 student policy가 최적의 성능을 달성하지 못하는 문제가 있음
실제 동물들은 priviledged information을 알지 못한 상태에서도, 세계를 이해하는 능력을 바탕으로 다양한 지형을 직관적으로 이동하는 방법을 배움. 이러한 자연적 학습 방식에서 영감을 받아, World Model-based Perception 방법을 제안. WMP는 환경의 월드 모델을 구축하고, 이를 바탕으로 policy를 학습.

1. Introduction

강화 학습은 최근 다양한 지형에서 Legged Robot의 이동을 학습하는 데 있어 큰 서오가를 거두었으며, 물리적 시뮬레이션에서 policy를 학습한 후 실제 환경으로 전이(sim to real transfer)하는 방식을 널리 사용한다. 이때 proprioception과 visual 정보를 입력으로 받아 각 관절에 필요한 Position이나 Effort 제어 값을 출력한다.
그러나 고해상도 픽셀 입력을 직접 사용하여, Reward Signal만으로 정책을 학습하는 것은 데이터 효율성이 매우 낮다. 또한 로봇이 전방 카메라를 사용할 경우, 앞으로 이동할 지형을 예측하기 위해 과거의 시각 정보를 기억해야 하는 추가적인 학습 난제가 발생한다. 이를 해결하기 위해, Priviledged Learning이 제안되었으며, Teacher Policy와 Student Policy로 분리된다.

Privileged Learning의 한계

하지만, Privileged Learning은 Student Policy의 Generalization Error로 인해 Teacher Policy의 행동을 완벽하게 복제하지 못하며, Privileged Information과 Visual Information의 gap이 클수록 성능 차이가 심화된다. 또한, Extra Information 필요하다. Teacher Policy는 학습 괒어에서 다양한 추가 정보를 필요로 하므로, 데이터 설계가 어렵고 노동집약적이다.

자연스러운 학습 방식에서 영감을 얻다.

반면, 동물은 비정형 지형을 직관적으로 학습하며, 제한된 감각 정보만으로도 익숙하지 않은 환경에서 적절한 결정을내릴 수 있다. 인간을 포함한 동물들이 세계를 이해하는 Mental Model을 구축하여 환경을 예측한 능력을 갖추고 있기 때문
동물들은 행동을 수행할 때, 과거 정보를 활용하여 현재를 인지하고, 미ㅐ 감각 데이터를 예측하는 능력을 가지고 있다. 이러한 개념에서 영감을 받아, Model-based Reinforcement Learning이 제안되었다. MBRL 기법 기반의 Legged Locomotion은 거의 이뤄지지 않음

WMP : World Model based Perception

이 논문에서는 World Model Learning이 vision based legged robot locomotion에 도움이 될 수 있을지를 탐구한다. WMP라는 새로운 end-to-end 프레임워크를 제안한다.

1. World Model 학습

• 시뮬레이션에서 환경의 월드 모델을 학습하여 과거 시각 정보로부터 미래의 시각 정보를 예측

2. Policy 학습

• World Model에서 추출된 추상적 표현을 기반으로 정책을 학습
• 단순히 시각 정보만을 사용하는 기존의 방법보다 의미 있는 정보 표현을 생성하고, 보다 효과저인 의사 결정을 가능하게 함

실험 결과

WMP를 기존 최신 기법과 비교하여 다양한 지형에서 평가함. Teacher Policy는 유사한 최적 보상을 달성하였으며, Student Pilicy는 높은 성능을 기록하였다. Unitree A1 로봇 기반으로 진행. WMP는 World Model을 활용한 최초의 vision based legged robot locomotion으로 기존의 privileged learning이 가진 한계를 극복할 수 있는 패러다임을 제시

2. Related Work

1. 강화학습 기반 legged locomotion 방법론
   이전 연구에서는 Proprioception만을 입력으로 받는 정책, Biologically Inspired Rewards design, Domain Randomiation, Curriculum Learning 등의 기법을 활용하여 다양한 실제 환경에서 이동할 수 있음을 보여줌.
   그러나 Visual perception이 없을 경우는 매우 어려움. 이를 해결하기 위해 Privileged Learning 프레임워크가 개발. 또한 최근 연구에서는, 과거 관측 데이터를 활용하여 Scandost를 추정하는 방법이 제안. 이에 반해 WMP은 이러한 한계를 극복하며 더 일반적이고 효과적인 시각 기반 학습 프레임워크를 제안
2. 모델 기반 강화학습 방법론 (Model-based Reinforcement Learning, MBRL)
   모델 기반 강화 핛브은 환경의 동역학 모델을 학습하여 의사 결정을 돕는 방식. 로봇 제어 정책을 학습하는 데 유망
   MBRL에서는 Partially Observable Environment이나 Pixel-Input Environments에서 정확한 동역학 모델을 학습하는 연구가 활발히 진행

3. Preliminries

Legged Locomotion을 Partially observable markov Decision Process, POMDP)로 정의하며, 이과정을 다음과 같은 튜플로 표현
S: 상태 공간(State Space)
O: 관측 공간(Observation Space)
A: 행동 공간(Action Space)
T(st+1 | st, at): 상태 전이 확률(Transition Density)
시간 t에서 상태 st와 행동 at가 주어졌을 때, 다음 상태 st+1로 전이될 확률 분포
r(st, at): 보상 함수(Reward Function)
γ ∈ (0, 1): 할인 계수(Discount Factor)
RL의 목표는 Expected Return을 극대화하는 Optimal Policy n*을 찾는 것. 할인된 보상의 합을 최대로 만드는 정책을 학습하는 것을 목적

강화 학습에서는 일반적으로 전태 상태 st를 직접 관측할 수 있다고 가정하지만, 실제 환경에서는 센서의 한계로 인해 Partial Observation만 획득이 가능하다. 즉 로봇은 특정 시점 t에서 완전한 상태 st를 알 수 없으며, 오직 부분 관측값만을 활용할 수 있다.
자기 수용 감각 & Depth Image만 사용
반면, 시뮬레이션은 추가적인 Priviledged Information를 포함한 완전한 상태 st를 얻을 수 있다.

4. Method

해당 섹션에서는 WMP(World Model-based Perception)을 소개한다. WMP는 고차원 센서 입력에서 유용한 정보를 추출하기 위해 World Model을 활용하는 End-to-End 프레임워크이다. 기존의 priviledged learning과 다르게, 단일 학습 단계에서 world model과 policy를 동시에 학습한다.

A. World Model Learning

기존 연구에 따라, 이 논문에서는 Recurrent State-Space Model,RSSM 변형을 월드 모델 구조로 채택하였다. 하지만 시뮬레이션에서 depth image를 얻는 데 걸리는 계산 비용과, 실제 로봇에서 RSSM을 실행하는 데 걸리는 Time Cost를 고려하여 기존 RSSM 모델을 수정하여 정책보다 낮은 빈도로 월드 모델을 업데이트하도록 설계하였다.
RSSM은 매 k스텝마다 순환 상태 ht를 업데이트한다. 4가지 주요 컴포넌트로 구성된다.

1. 순환 모델 : 저차원 잠재 공간에서 작동하며, 과거의 순환상태,과거 행동 시퀀스,이전 확률 상태를 이용해 새로운 결정론적 순환 상태 ht를 예측
2. 인코더 : 현재 관측값을 받아, 후방 확률 상태를 계산
3. 동적 예측기 : 실제 관측값 없이도, 미래 상태를 예측할 수 있도록 설계
4. 디코더 : 고차원 관측값을 재구성할 수 있도록 학습

순환 모델은 GRU 네트워크를 사용하여 구현되었으며, 깊이 이미지는 CNN, 자기수용 감각은 MLP을 활용하여 인코딩 및 디코딩을 수행
RSSM의 손실함수는 다음과 같으며, 재구성 손실항과 KL 정규화 항으로 구성된다.
재구성 손실항 : 후방 상태가 현재 관측값의 충분한 정보를 포함하도록 유도, 원래의 고차원 데이터를 잘 복원할 수 있도록 함
KL 정규화 : 후방 상태와 사전 상태 간의 KL Divergence를 최소화하여, 관측값이 없어도 미래 상태를 예측할 수 있도록 모델을 규제

B.Policy Learning

Vision Based Locomotion에서 정책 학습은 Partial Observability 문제로 인해 어려움이 존재, 그러나 잘 학습된 월드 모델의 순환 상태는 충분한 정보를 포함하고 있어 미래 상태를 예측할 수 있다. 이는 실제 마르코프 상태에 가까운 정보를 제공한다는 점에서 중요한 의미를 가진다.

• Asymmetric Actor-Critic Framework : 훈련 중에 추가 정보를 활용하여 가치 예측을 더 정확하게 하면서, 실제 실행 시에는 액터가 시각 정보만을 사용하여 일반화할 수 있도록 유도. 순환 상태가 크리틱 학습에서 중요한 역할을 하는 것을 발견. RSSM에서 얻은 순환상태는 보다 일반적인 지형 정보를 포함할 수 있음
• Proxiaml Policy Optimization(PPO)알고리즘을 사용하여 액터-크리틱 모델을 학습
  그러나, 기존 모델 기반 강화학습 기법과 달리, WMP에서는 월드 모델을 사용하여 정책 학습을 위한 롤아웃 데이터를 생성하지 않음. 시뮬레이터가 가장 정확한 데이터이기 때문에, 시뮬레이터에서 직접 샘플링된 데이터를 사용하여 정책을 학습

C. Training Details

Environment

Legged Gym 코드 베이스 기반으로 실험 수행, IsaacGYM 시뮬레이터를 사용. 4096대의 Unitree A1 로봇 사용
6가지 지형에서 실험을 수행 (Slope, Stair, Gap, Climb, Crawl, Tilt)
난이도(Curriculum Learning)은 가장 쉬운 난이도에서 시작하며, 로봇이 지형을 통과하면 더 어려운 지형으로 이동, 목표 속도의 절반 이하로 이동하면 더 쉬운 지형으로 이동,
로봇의 행동 주기는 50hz / Depth image는 k 타임스텝마다 계산되며, 100ms 지연을 추가하여 sim-to-real 전이를 용이하게 함
물리적 파라미터 랜덤화를 통해 policy의 강건성을 높임

State and Action Space

로봇의 관측값은 Proprioception과 depth image로 구성된다.
Proprioception은 총 45차원 (기본 속성 - 기본 각속도, 중력 벡터, 제어 명령, 관절 정보, 이전 행동)
Priviliged Information (시뮬레이터에서만) - Scandots, 발 접촉력, 랜덤화된 물리적 파라미터
Depth image - 로봇자기 중심 시야에서 촬영된 깊이 영상, 시야각 58 x 58

Action

액션에서는 각 관절의 목표 위치를 지정한다.

Reward Function

로봇은 3차원의 속도 명령을 추적하도록 학습된다. 기존 연구에서는 경로를 따라 수동으로 웨이포인트를 설정해야하는 반면, 이 논문은 단순화된 속도 추적 보상을 사용하여 웨이포인트 설정을 최소화하였다.

AMP(Adversarial Motion Priors) 스타일 보상

AMP을 활용하여 로봇이 보다 자연스러운 움직임을 학습하도록 유도

Experimental Result

이 논문에서는 다음 세가지 질문을 검증하는 것을 목표로 한다.

1. WMP는 SOTA와 비교했을 떄 얼마나 좋은 성능을 보이는가?
2. 시뮬레이터에서 학습된 월드 모델이 실제 환경에서도 궤적을 정확히 예측할 수 있는가?
3. WMP의 시뮬레이션 성과가 실제 로봇에서도 그대로 유지될 수 있는 가?

A. Simulation Comparision

WMP의 성능을 검증하기 위해, RL 반환값과 속도 추적 오차(MSE)를 다양한 지형에서 측정

B. Empirical Study

WMP의 월드 모델이 정책 학습에 기여하는 방식을 분석하기 위한 추가 실험 진행

1. 순환 상태 시각화 : 지형별 차이를 효과적으로 인코딩하며 이를 통해 정책이 특정 환경에 맞는 적절한 행동을 학습할 수 있음
2. 모델 업데이트 간격
3. 학습 데이터 길이
4. 실제 환경에서의 예측 성능
   • 실제 환경에서 로봇의 이동 궤적(Trajectroy) 수집
   • 초기 관측값과 행동 시퀀스를 입력으로주고, 월드 모델이 미래 깊이 이미지를 예측하도록 함
   • intermediate depth image 없이 예측이 가능한지 확인

C. Real-World Evaluation

WMP 및 기존 방법을 실제 A1 로봇에 적용하여 실험 수행

• 모든 방법은 Jetson NX 하드웨어에서 직접 실행
• Depth Images는 Intel D435i 전방 카메라에서 60Hz 속도로 읽어오며, 해상도는 424 x 240
• 시뮬레이터와 실제 환경 간의 시각적 차이를 줄이기 위해, 공간 및 시간 필터를 적용하여 노이즐를 보정
• 보정된 이미지를 64x64로 크롭 및 다운 샘플링 후, 100ms 지연과 함께 월드 모델로 입력
• PD 제어기 설정

Conclusion

WMP는 Student Policy보다 훨씬 더 어려운 지형을 성공적으로 이동하며, Sim-to-Real 차이를 최소화하는 강력한 성능을 보임

World Model-based Perception for Visual Legged Locomotion Test

conda create --name wmp python=3.8 -y
conda activate wmp
https://github.com/bytedance/WMP
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117
conda create --name wmp python=3.8 -y
apt-get install build-essential --fix-missing
apt-get install ninja-build
pip install setuptools==59.5.0
pip install ruamel_yaml==0.17.4
apt install libgl1-mesa-glx -y
pip install opencv-contrib-python
git clone https://github.com/bytedance/WMP.git
pip install -r requirements.txt

# Train 
python legged_gym/scripts/train.py --task=a1_amp --headless --sim_device=cuda:0
# Test
python legged_gym/scripts/play.py --task=a1_amp --sim_device=cuda:0 --terrain=climb

Reference

• https://wmp-loco.github.io/