world model : 빠른 논문 리뷰 : TD-MPC2: Scalable, Robust World Models for Continuous Control

용어 설명

• TD-MPC (Temporal Difference Learning for Model Predictive Control): 이 논문의 바탕이 되는 model-based RL 알고리즘. 환경의 원본 이미지나 상태를 복원(reconstruction)하는 대신, 행동에 따른 미래의 잠재 상태(latent state)와 보상(reward)만을 예측하는 모델을 학습하고, 그 잠재 공간 안에서 최적의 행동 궤적을 계획(planning)합니다.
• Implicit (decoder-free) world model: 미래의 시각적 이미지나 원본 관측값을 다시 그려내는(decoding) 과정 없이, 오직 제어(control)에 필요한 결과값(reward, value)만을 예측하도록 학습된 world model을 뜻합니다.
• SimNorm (Simplicial Normalization): 이 논문에서 도입한 새로운 정규화(normalization) 기법. 잠재 상태(latent representation)를 softmax 연산을 통해 여러 개의 고정 차원 심플렉스(simplices)로 투영합니다. 이를 통해 표현의 희소성(sparsity)을 유도하고 값이 폭주하는 것을 막아 학습 안정성을 크게 높입니다.
• Discrete regression (in log-transformed space): 보상이나 가치(value)를 연속적인 실수 값으로 예측하는 대신, 값을 여러 구간(bin)으로 나누어 분류(classification) 문제처럼 접근하는 방식입니다. 태스크마다 보상의 크기(scale)가 천차만별인 multi-task 환경에서 알고리즘이 흔들리지 않게 잡아주는 핵심 역할을 합니다.
• Generalist embodied agent: 하나의 단일 모델로 강아지 로봇, 사람형 로봇, 로봇 팔 등 다양한 물리적 형태(embodiment)와 행동 공간(action space)을 넘나들며 여러 태스크를 수행할 수 있는 범용 AI 에이전트입니다.

---

Purpose of the Paper

• 기존 연구의 한계: 기존의 RL 알고리즘들은 태스크가 바뀔 때마다 하이퍼파라미터(hyperparameters)를 일일이 수정해야 하는 극심한 민감도를 가졌습니다. 또한, 다양한 로봇 형태와 행동 공간이 섞인 multi-task 데이터를 학습시키려면 도메인 지식이 필수적이거나, 전문가 수준의 데모 데이터(behavior cloning)에 크게 의존해야 했습니다. 특히 기존 TD-MPC는 모델의 크기(parameters)를 키우면 오히려 성능이 하락하는 확장성(scalability)의 한계가 있었습니다.
• 새로운 문제 정의 및 접근: 이 연구는 **"단 하나의 하이퍼파라미터 세트로, 전문가의 데모 없이(uncurated data), 모델 크기를 키울수록 성능이 우상향하는 범용 world model을 만들 수 있는가?"**라는 질문에 답하고자 합니다. 이를 위해 기존 TD-MPC 알고리즘의 뼈대를 뜯어고쳐 극도의 강건성(robustness)과 확장성을 지닌 TD-MPC2를 제시합니다.

---

Key Contributions

• 알고리즘 강건성(Robustness)의 극대화:
  • 기여: Reward와 Value 예측을 log-transformed space 기반의 discrete regression으로 변경하고, 잠재 상태에 SimNorm을 적용했습니다. 또한 LayerNorm, Mish 활성화 함수, 5개의 Q-functions 앙상블을 도입했습니다.
  • 참신성: 태스크마다 보상 스케일이 달라도 학습이 붕괴되지 않으며, 기울기 폭발(gradient explosion)을 원천 차단하여 104개의 전혀 다른 태스크를 단일 하이퍼파라미터로 정복하는 전례 없는 안정성을 달성했습니다.
• 도메인 지식 없는 Massively Multi-task Architecture 도입:
  • 기여: 학습 가능한 정규화된 task embeddings를 도입하고, 크기가 다른 observation/action space를 하나로 통합하기 위해 zero-padding과 action masking 기법을 사용했습니다.
  • 참신성: 로봇의 관절 수나 센서 데이터 형태가 완전히 달라도, 인간의 사전 지식 개입 없이 하나의 네트워크가 모든 데이터를 동시에 소화할 수 있는 범용 구조를 제안했습니다.
• RL World Model의 Scaling Law 증명:
  • 기여: 1M 파라미터부터 317M 파라미터까지 모델 크기를 키우며 80개의 multi-task 데이터셋을 학습시키는 데 성공했습니다.
  • 참신성: 모델을 키우면 오히려 성능이 떨어지던 기존 RL의 고질적 문제를 해결하고, 파라미터 수의 로그에 비례하여 에이전트의 능력이 선형적으로 증가한다는 것을 실험적으로 입증했습니다.

---

Experimental Highlights

• Single-task SOTA 달성: DMControl, Meta-World, ManiSkill2, MyoSuite 등 4개 도메인의 104개 연속 제어 태스크에서 SAC, DreamerV3 등 기존 state-of-the-art 모델들을 압도했습니다. 특히, 다른 baseline들은 튜닝 없이 풀지 못하는 고난이도 로봇 보행(Dog)이나 다중 물체 조작(Pick YCB) 태스크에서 압도적인 데이터 효율성과 최종 성능을 보여주었습니다.
• Massively Multi-task 학습 성공: 80개의 태스크를 동시에 학습하는 실험에서, 모델 파라미터가 1M일 때는 Normalized score가 약 16점에 불과했으나, 317M 파라미터로 확장했을 때 70.6점까지 급상승하는 뚜렷한 스케일링 효과를 그래프로 입증했습니다.
• Task Embeddings의 의미적 군집화 (T-SNE 시각화): 학습된 task embeddings를 시각화한 결과, 단순히 목표(목적)가 비슷한 태스크보다 물리적 구조(embodiment)나 물체 상호작용 방식(dynamics)이 유사한 태스크끼리 임베딩 공간에서 가깝게 군집화되는 흥미로운 현상을 발견했습니다 (예: Door Open과 Door Close가 인접).
• Few-shot Learning 가능성: 70개 태스크로 pre-trained된 19M 모델을 완전히 새로운 10개의 태스크에 fine-tuning 한 결과, 처음부터 학습하는 것보다 데이터가 적은 초기 구간(20k steps)에서 성능이 2배 이상 빠르게 향상되는 것을 확인했습니다.

---

Limitations and Future Work

• Limitation 1: Discrete Action Spaces로의 확장 한계
  • 설명: TD-MPC2는 근본적으로 연속적 행동 공간에 맞춰진 MPC(Model Predictive Control) 계획 알고리즘에 의존하므로, Atari나 Minecraft 같은 이산적(discrete) 행동 공간에는 당장 적용하기 어렵습니다.
  • Future Work: MPC를 MCTS(Monte Carlo Tree Search)와 같은 이산 공간용 계획 알고리즘으로 교체하거나, 샘플링을 통해 MPC를 이산 공간에 적용하는 방법을 발견한다면 진정한 만능 에이전트가 될 수 있습니다.
• Limitation 2: Reward 의존성 문제
  • 설명: 대규모 pre-training을 하려면 태스크의 성공/실패를 명확히 알려주는 보상(reward) 함수가 필요한데, 실제 현실 환경(in the wild)에서는 이런 보상 데이터를 대량으로 얻기 매우 힘듭니다.
  • Future Work: 사람의 선호도(human preferences), 단순 성공 라벨, 또는 관측값과 목표 지점 간의 임베딩 거리 등 '일반화된 형태의 보상'을 활용한 unsupervised pre-training 연구가 필수적입니다.
• Limitation 3: 자율성 부여에 따른 안전성(Safety) 위험
  • 설명: 강력한 generalist 모델에게 실제 물리적 로봇의 통제권을 무제한으로 넘기면, 잘못 설계된 보상을 악용하거나 예측 불가능한 치명적 사고를 낼 위험이 큽니다.
  • Future Work: 모델이 계획(planning)을 세울 때 불확실성이 높은 행동을 스스로 페널티 처리하는 test-time regularization이나, 추가적인 안전망(safety checks)을 알고리즘 내부에 통합하는 연구가 병행되어야 합니다.

---

Overall Summary

이 논문은 기존 RL 알고리즘들의 하이퍼파라미터 민감도와 확장성 부족 문제를 근본적으로 해결한 강력한 model-based RL 알고리즘인 TD-MPC2를 제안합니다. 구조적 정규화(SimNorm)와 스케일에 구애받지 않는 손실 함수(Discrete regression)를 도입함으로써, 단 하나의 세팅으로 104개의 다양한 로봇 제어 태스크에서 SOTA를 달성했습니다. 나아가 도메인 지식 없이도 317M 파라미터 규모의 단일 모델이 80개의 서로 다른 로봇 형태와 환경을 동시에 제어할 수 있음을 증명했습니다. 이는 거대 언어 모델처럼 물리적 로봇 제어 분야에서도 대규모 데이터와 컴퓨팅을 활용한 "Foundation Model (Generalist Agent)" 구축이 가능함을 보여준 매우 중요한 이정표입니다.

---

쉬운 설명

기존의 로봇 인공지능들은 '강아지 로봇으로 걷기'를 배울 때와 '로봇 팔로 문 열기'를 배울 때마다 연구자가 AI의 뇌 구조(hyperparameter)를 일일이 그 상황에 맞게 튜닝해줘야 했습니다. 게다가 AI의 뇌(parameter)를 무작정 크게 만들면 오히려 똑똑해지기는커녕 학습이 무너져버렸습니다.

TD-MPC2는 이 문제를 **"평가 방식의 변경"**과 **"상상력의 제한"**으로 해결했습니다. 점수(Reward)의 스케일이 10점이든 1000점이든 흔들리지 않게 객관식(Discrete regression)으로 채점 방식을 바꾸고, AI가 미래를 상상할 때 값이 너무 커져서 폭주하지 않도록 뇌 신경망에 안전장치(SimNorm)를 달았습니다.

그 결과, 연구자가 세팅을 단 한 번도 바꾸지 않고도 무려 104가지의 완전히 다른 로봇 조종법을 마스터했습니다. 심지어 뇌 크기를 3억 개 이상으로 엄청나게 키웠더니, 80가지의 각기 다른 로봇과 임무를 하나의 AI가 동시에 다 해내는 **"다재다능한 만능 로봇 두뇌"**가 탄생한 것입니다. 이는 텍스트나 이미지를 넘어, 물리적으로 움직이는 로봇계에도 챗GPT 같은 만능 모델이 나올 수 있다는 것을 보여줍니다.

 

 

 

 

 

 

 

 

관측값 -> 임베딩
미래 관측값 -> 임베딩 

현재 z와 a를 신경망에 넣고 다음 z를 예측 후 정규화

다음 예측값과  다음 액션을 q 펑션에 넣는다. 이런게 5개가 있음. 
그중 2개 골라서 더 작은값을 정답으로 사용

현재 z와 a를 q펑션에 넣고 위에서 구한 미래의 정답을 loss로 사용


현재 z와 a를 리워드 신경망에 넣고 보상을 예측. 버퍼에 있던 실제 보상과 비교해 loss로 사용

아까 역학 모델로 예측해둔 '다음 z'와 처음에 만들어둔 '미래 관측값 임베딩(정답지)'을 비교해서 loss로 사용

현재 z를 정책(Policy) 신경망에 넣고 최적의 액션을 예측

방금 뱉어낸 이 새로운 액션이 진짜 좋은 액션인지 평가를 받아야 함.

그래서 현재 z와 이 새로운 액션을 아까 위에서 학습시키고 있던 메인 Q 펑션에 한 번 넣어봄.

그러면 메인 Q 펑션이 "어, 그 상태에서 그 액션을 하면 미래 가치(점수)가 이 정도 될 거 같은데?" 하고 평가 점수를 알려줌 (이 점수가 바로 Q값임).

이 액션의 Q값을 가장 높이면서, 너무 뻔한 행동만 하지 않도록(엔트로피 유지) 정책 loss로 사용

지금까지 구한 Q loss, 리워드 loss, 다음 z loss를 다 더해서 메인 신경망(인코더, 역학, 리워드, 메인 Q)들을 한 번에 업데이트

정책 신경망은 방금 구한 정책 loss 하나만 써서 따로 업데이트

방금 똑똑해진 메인 Q 펑션 5개의 가중치를 타겟 Q 펑션 5개(정답지 만들 때 쓴 애들)에 아주 조금씩(EMA 방식) 덮어씌워서 타겟 최신화