Assistax: A Multi-Agent Hardware-Accelerated Reinforcement Learning Benchmark for Assistive Robotics

기술 요약: Assistax

문제 정의

보조 로봇 공학 분야는 청소부터 이동성이 저하된 개인을 돌보는 일까지, 인간의 일상 활동을 도울 수 있는 자율 시스템을 개발하는 것을 목표로 한다. 이 영역의 핵심적인 과제는 인간-로봇 상호작나이 본질적으로 다중 에이전트(multi-agent)라는 점이다. 즉, 인간은 수동적인 수혜자가 아니라 다양한 선호도, 능력, 그리고 잠재적 장애를 가진 능동적인 협력자이다. 현재의 강화 학습(RL) 시뮬레이션 환경은 이러한 복잡성을 충분히 다루지 못하는 경우가 많다. 많은 환경이 너무 단순하여(예: 그리드 월드) 복잡한 3D 연속 제어에 대한 통찰을 제공하지 못하거나, 고충실도 로보틱스 시뮬레이터는 처리량이 낮아 강화 학습 훈련이 지나치게 느려지는 문제를 겪는다. 또한, 기존의 벤치마크 중에는 로봇과 반응형 인간 파트너 사이의 다중 에이전트 상호작용을 능동적으로 모델링하기보다는 인간을 정적이거나 스크립트된 개체로 취급하는 경우가 많다. 하드웨어 가속, 3D 연속 제어, 그리고 로봇이 사전 적응 없이 보이지 않는 파트너와 협력해야 하는 애드혹 팀워크(Ad-Hoc Teamwork, AHT) 패러다임을 결합한 벤치마크는 구체적으로 부족한 실정이다.

방법론

저자들은 JAX와 MuJoCo MJX를 기반으로 구축된 다중 에이전트 강화 학습(MARL) 벤치마크인 Assistax를 소개한다. 이 아키텍처는 엔드 투 엔드 하드웨어 가속을 가능하게 하여 단일 GPU에서 대규모 병렬화를 지원한다.

환경 제품군 (Environment Suite)

Assistax는 실제 시나리오에서 영감을 얻은 다섯 가지의 뚜렷한 보조 작업을 구현하며, 각 작업은 토크 제어 로봇 에이전트와 단순화된 휴머노이드 에이전트(인간 파트너를 시뮬레이션)를 특징으로 한다:

1. 긁기 (Scratch): 로봇은 휴머노이드의 팔에 있는 무작위 타겟으로 이동하여 특정 접촉력을 가하면서 긁는 동작을 모방해야 한다.
2. 양치질 (Tooth Brushing): 로봇은 칫솔을 휴머노이드의 입으로 안내하며, 이를 위해 정렬, 접촉력, 그리고 접선 속도가 필요하다.
3. 식사 (Feeding): 로봇은 숟가락을 휴머노이드의 입으로 안내하며, 흘리는 것을 방지하기 위해 방향과 접근 속도를 관리해야 한다.
4. 목욕 (Bed Bathing): 로봇은 분산된 타겟 지점들을 방문하며 휴머노이드의 팔을 닦는다.
5. 팔 보조 (Arm Assist): 로봇은 갈고리를 사용하여 휴머노이드의 팔을 들어 올려 타겟 방향으로 재배치한다.

이 환경들은 시뮬레이션된 떨림, 관절 약화, 제한된 가동 범위와 같은 현실적인 도전 과제들을 포함한다.

선호도 모델링 (Preference Modeling)

인간의 다양성을 포착하기 위해, Assistax는 **선호도 보상(Preference Rewards)**을 도입한다. 이는 다음 세 가지 축을 기준으로 작업 보상을 증강한다:

• 힘 (Force): 선호되는 접촉력 범위.
• 속도 (Speed): 선호되는 말단 장치(end-effector) 속도 범위.
• 접촉 횟수 (Contact Count): 지속적인 상호작용을 장려하기 위해 과도한 재접촉(찌르기)에 대한 페널티 부여.
  이러한 선호도는 두 에이전트 모두를 위한 관측값으로 인코딩되어, 로봇이 특정 인간의 제약 사항을 추론하고 적응할 수 있도록 한다.

훈련 및 평가 파이프라인

본 논문은 2단계 파이프라인을 제안한다:

1. MARL 공동 훈련 (Co-training): 로봇과 휴머노이드 에이전트는 협력적으로 작업을 해결하기 위해 표준 MARL 알고리즘(IPPO, MAPPO, MASAC 및 그 순환 변형들)을 사용하여 공동 훈련된다.

2. 애드혹 팀워크 (AHT) 벤치마킹: 제로샷 협업(zero-shot coordination)을 평가하기 위해, 저자들은 다양한 630개의 휴머노이드 파트너 집단을 사전 훈련한다. 이 다양성은 다음을 통해 유도된다:

   • 행동 관습 (공동 훈련 역학을 통해)
   • 선호도 파라미터 (힘, 속도, 접촉)
   • 장애 설정 (이동성 제한)

   그 후 벤치마크는 특정 하위 집합의 파트너(예: 5개 에이전트)로 훈련된 로봇 정책을 보유된 테스트 세트(예: 625개의 미보지 에이전트)에 대해 평가한다. 이 설정은 로봇이 추가적인 적응 없이 새로운 선호도 조합과 장애 프로필에 일반화할 수 있는 능력을 테스트한다.

주요 기여

1. 하드웨어 가속 벤치마크: Assistax는 보조 케어 맥락에서 두 개의 뚜렷한 체화된(embodied) 에이전트를 가진 3D 연속 환경을 위한 최초의 AHT 벤치마크이다. JAX와 MuJoCo MJX를 활용함으로써, 단일 GPU에서 유사한 CPU 기반 환경(예: Assistive Gym) 대비 최대 412배 빠른 오픈 루프 시뮬레이션과 25배 빠른 벽시계 시간(wall-clock) 훈련을 달성한다.
2. MARL 베이스라인 및 구현: 저자들은 JAX로 구현된 PPO 및 SAC 변형(순환 구조 포함)에 대한 튜닝된 베이스라인을 제공하여, 로봇-휴머노이드 팀의 확장 가능한 공동 훈련을 가능하게 한다.
3. AHT 파이프라인 및 데이터 공개: 본 논문은 엄격한 AHT 평가 파이프라인을 도입하고, Hugging Face를 통해 수천 개의 사전 훈련된 휴머노이드 파트너 정책을 공개한다. 이를 통해 연구자들은 제로샷 협업 결과를 재현하고, 다양하게 정의된 파트너 집단에 대해 새로운 AHT 알고리즘을 테스트할 수 있다.

결과

• MARL 성능: 공동 훈련 시나리오에서 IPPO와 MAPPO의 피드포워드 구현이 순환 변형보다 일관되게 우수한 성능을 보였다. MASAC 또한 강력한 성능을 입증했다. 선호도 보상의 도입은 베이스라인의 상대적 성능 순위를 크게 바꾸지는 않았으나, 보상 지형의 복잡성을 더했다.
• AHT 협업 격차: 새로운 선호도 조합을 가진 미보지 파트너를 대상으로 기존 RL 알고리즘을 평가했을 때(특히 훈련 및 테스트 분포 간의 중첩이 최소화된 "하드" 설정에서), 상당한 **협업 격차(coordination gap)**가 관찰되었다. 로봇 정책은 분포 외(out-of-distribution) 선호도에 적응하는 데 어려움을 겪었다. 그러나 훈련과 테스트 분포가 더 유사하거나 훈련 세트가 더 클 경우 이 격차가 유의미하게 줄어든다는 점을 논문은 명시하고 있다.
• 효율성: 1024개의 병렬 환경에서 4,000만 타임스텝을 수행하는 전형적인 훈련 과정이 약 20분 만에 완료되었으며, 이는 복잡한 MARL 연구를 위한 하드웨어 가속의 실질적인 생존 가능성을 입증한다.

의의

본 논문은 Assistax를 보조 케어 분야의 RL 발전을 위한 실질적이고 필수적인 단계로 자리매김한다. 고립된 에이전트 훈련에서 다중 에이전트 및 하드웨어 가속 환경으로 전환함으로써, Assistax는 복잡한 로보틱스 영역 연구를 제한해 온 처리량 병목 현상을 해결한다. 결정적으로, 상호작용을 애드혹 팀워크 문제로 프레임화하고 다양한 인간 파트너와의 제로샷 협업을 평가하기 위한 파이프라인을 제공함으로써, Assistax는 보조 로봇 공학의 핵심 목표인 "모든 사용자에 대한 일반화"를 위한 현실적인 테스트베드를 제공한다. 사전 훈련된 파트너 정책의 공개는 AHT 평가를 표준화하여, 커뮤니티가 정적인 벤치마크를 넘어 역동적이고 인간 중심적인 시나리오에서 로봇 정책의 적응력을 엄격하게 테스트할 수 있도록 돕는다.