Steering Your Diffusion Policy with Latent Space Reinforcement Learning

저자: Andrew Wagenmaker, Mitsuhiko Nakamoto, Yunchu Zhang, Seohong Park, Waleed Yagoub, Anusha Nagabandi, Abhishek Gupta, Sergey Levine · CoRL 2025

라이선스

• 오픈 액세스
• 논문: CC BY 4.0

아래 내용은 모두 논문의 라이선스(CC BY 4.0)에 따라 논문 전문을 바탕으로 하며, 그림에는 저자를 표기했습니다.

요약

DSRL은 입력 노이즈를 선택하는 작은 RL 정책을 학습하여 고정된 확산 또는 플로우 정책을 개선하며, 수십에서 수백 에피소드 만에 실제 로봇과 π0를 적응시킨다.

중요한 이유

팀은 작동 중인 행동 복제 정책을 가중치를 건드리지 않고 그 자리에서, 수십 개의 실제 에피소드와 희소한 성공 보상만으로 개선할 수 있으며, 이 덕분에 블랙박스 모델이나 API로만 제공되는 모델도 적응시킬 수 있다. π0 실험에서 DSRL은 8 GB RTX 3070에서 학습되었는데, π0 저장소는 LoRA 미세조정에 최소 22.5 GB가 필요하다고 밝힌다.

문제

오늘날 흔히 확산 정책인 행동 복제 정책은 자신의 배포 경험으로부터 학습할 수 없으므로, 부진한 성능을 고치려면 대개 인간 시연을 더 수집해야 한다. RL은 이를 자율적으로 개선할 수 있지만, 확산 정책을 RL로 미세조정하려면 다단계 디노이징 체인을 통한 역전파가 필요하여 계산량이 많고 수치적으로 불안정하며, 표준 RL은 실제 로봇에 쓰기에는 너무 많은 샘플을 필요로 한다.

제안 방법

DSRL은 사전학습된 확산 정책(결정론적 DDIM 또는 플로우 샘플링 사용)을 환경의 일부로 취급하고, 초기 노이즈를 N(0, I)에서 뽑는 대신 선택하는 별도의 잠재 노이즈 정책을 학습한다. 이로써 적응은 '잠재 행동 MDP'에 대한 일반적인 RL이 되며, 기반 정책의 가중치가 아니라 순전파만 필요하다. 노이즈 에일리어싱 변형인 DSRL-NA는 실제 행동에 대해 크리틱을 학습하고 이를 노이즈 공간 크리틱으로 증류하므로, 오프라인 데이터를 사용할 수 있고 같은 행동으로 매핑되는 많은 노이즈 값을 활용할 수 있다.

핵심 통찰

• 결정론적 DDIM 또는 플로우 샘플링에서는 입력 노이즈가 행동을 완전히 결정하므로, 노이즈를 선택하면 정책을 바꾸지 않고도 정책을 조향할 수 있다.
• 기반 정책은 분포 내 행동만 출력하므로, 그 노이즈에 대해 최적화하면 명시적인 페널티 항 없이도 오프라인 RL의 보수성을 얻는다.
• 많은 노이즈 값이 거의 같은 행동으로 매핑된다. 행동 공간 크리틱을 노이즈 공간 크리틱으로 증류하면 이 에일리어싱을 활용하며, 노이즈 공간에서 일반 SAC 대비 약 절반의 샘플만 필요했다.
• π0의 1,600차원 청크 행동 공간에서는 32차원 노이즈 벡터 하나를 청크 전체에 반복하는 것만으로도 정책을 조향하기에 충분한 표현력을 갖는다.
• 기반 정책이 어느 정도 작업 해결 동작을 보이는 한, 크기, 학습 기간, 데이터 품질이 서로 다른 기반 정책 전반에서 조향이 효과를 보였다.

주요 결과

• Robomimic과 OpenAI Gym에서의 온라인 적응은 DPPO, IDQL, DQL, DIPO, QSM보다 최소 5~10배 적은 샘플로 거의 최적의 성능에 도달했다.
• 10개의 OGBench 오프라인 작업에서 DSRL은 약 절반에서 최고 수준(state-of-the-art)이었으며, 조향하지 않은 플로우 정책보다 크게 개선되었다.
• 실제 Franka 단일 작업 큐브를 그릇에 넣기: 3,500 온라인 스텝(약 40 에피소드) 만에 성공률이 20%에서 거의 100%로 올랐다. RLPD는 인간 개입 유무와 관계없이 작업을 해결하지 못했다.
• Bridge V2로 학습한 다중 작업 정책을 사용하는 실제 WidowX 250: 100~150 에피소드 내에 집어 놓기, 서랍 닫기, 블록 쌓기에서 상당한 향상을 보였으며, RLPD는 학습하지 못했다.
• LIBERO에서의 π0: 약 10,000개의 온라인 샘플 후 성공률이 약 20%에서 약 100%로 올랐다. ResiP와 V-GPS는 이를 유의미하게 개선하지 못했다.
• 공개 π0-DROID 체크포인트를 사용한 실제 Franka: 토스터 작업과 접시에 숟가락 놓기 작업을 각각 80 및 65 에피소드(약 11,000 및 10,000 스텝) 만에 개선했으며, 저자는 이를 π0의 최초 실제 환경 RL 미세조정이라고 부른다.

한계

• 탐색은 기반 정책에 의해 제한된다. 모든 확산 정책이 조향 가능하다는 보장은 없으며, 매우 좁은 행동 분포는 선택지를 너무 적게 남길 수 있다(저자).
• 주어진 기반 정책을 얼마나 개선할 수 있는지 미리 정량화할 방법은 아직 없다(저자).
• 다른 RL 방법과 마찬가지로 모방 학습이 요구하는 것 이상으로 보상 신호, 온라인 변형의 경우 온라인 롤아웃, 그리고 실제 환경에서의 리셋이 필요하다(저자).
• 실제 환경 작업은 수작업으로 정의한 0-1 성공 보상을 사용하며, 평가 규모가 작다(WidowX에서 작업당 10~20회 시행).

분야에 미친 영향

CoRL 2025 Best Paper / Best Student Paper 상 최종 후보이며, 현재까지 피인용 242회이다. 논문은 프로젝트 웹사이트(diffusion-steering.github.io)를 링크하고 공개 π0 체크포인트를 기반으로 하지만, 라이선스가 부여된 PMLR 사본에는 코드 저장소가 명시되어 있지 않다.

초록

인간 시연으로부터 학습한 로봇 제어 정책은 많은 실제 응용에서 인상적인 결과를 달성해 왔다. 그러나 새로운 개방형 환경에서 흔히 그렇듯 초기 성능이 만족스럽지 않은 상황에서는, 이러한 행동 복제(BC)로 학습한 정책의 동작을 더 개선하려면 일반적으로 추가 인간 시연을 수집해야 하며, 이는 비용과 시간이 많이 드는 과정이다. 반면 강화학습(RL)은 자율적인 온라인 정책 개선을 가능하게 할 잠재력이 있지만, 일반적으로 요구되는 샘플 수가 많아 이를 달성하지 못하는 경우가 많다. 본 연구에서 우리는 효율적인 실제 환경 RL을 통해 BC로 학습한 정책의 빠른 자율 적응을 가능하게 하는 방향으로 나아간다. 특히 최신 BC 방법론인 확산 정책(diffusion policy)에 초점을 맞추어, 우리는 강화학습을 통한 확산 조향(diffusion steering via reinforcement learning, DSRL)을 제안한다. 이는 BC 정책의 잠재 노이즈 공간에서 RL을 실행하여 BC 정책을 적응시키는 방법이다. 우리는 DSRL이 샘플 효율이 매우 높고, BC 정책에 대한 블랙박스 접근만을 필요로 하며, 효과적인 실제 환경 자율 정책 개선을 가능하게 함을 보인다. 또한 DSRL은 확산 정책 미세조정과 관련된 많은 어려움을 피하며, 기반 정책의 가중치를 수정할 필요를 아예 없앤다. 우리는 시뮬레이션 벤치마크, 실제 로봇 작업, 그리고 사전학습된 범용 정책의 적응에서 DSRL을 시연하여, 실제 환경 정책 개선에서의 샘플 효율과 효과적인 성능을 보여 준다.

링크

• arXiv 2506.15799: https://arxiv.org/abs/2506.15799
• DOI 10.48550/arXiv.2506.15799: https://doi.org/10.48550/arXiv.2506.15799
• PDF: https://raw.githubusercontent.com/mlresearch/v305/main/assets/wagenmaker25a/wagenmaker25a.pdf
• 프로젝트 페이지: https://diffusion-steering.github.io/