Technical Summary: Assistax

Problem Statement

The field of assistive robotics aims to develop autonomous systems capable of aiding humans in daily activities, from cleaning to caring for individuals with mobility impairments. A critical challenge in this domain is that human-robot interaction is inherently multi-agent; the human is an active collaborator with varying preferences, capabilities, and potential disabilities, rather than a passive recipient. Current reinforcement learning (RL) simulation environments often fail to address this complexity adequately. Many are too simple (e.g., grid-worlds) to provide insights into complex 3D continuous control, while high-fidelity robotics simulators often suffer from low throughput, making RL training prohibitively slow. Furthermore, few existing benchmarks actively model the multi-agent interaction between a robot and a reactive human partner, often treating the human as a static or scripted entity. There is a specific lack of benchmarks that combine hardware acceleration, 3D continuous control, and the Ad-Hoc Teamwork (AHT) paradigm, where a robot must coordinate with unseen partners without prior adaptation.

Methodology

The authors introduce Assistax, a multi-agent reinforcement learning (MARL) benchmark for assistive robotics built entirely on JAX and MuJoCo MJX. This architecture enables end-to-end hardware acceleration, allowing for massive parallelization on a single GPU.

Environment Suite

Assistax implements five distinct assistive tasks inspired by real-world scenarios, all featuring a torque-controlled robot agent and a simplified humanoid agent (simulating a human partner):

1. Scratch: The robot must navigate to a random target on the humanoid's arm and apply a specific contact force while mimicking a scratching motion.
2. Tooth Brushing: The robot guides a toothbrush to the humanoid's mouth, requiring alignment, contact force, and tangential speed.
3. Feeding: The robot guides a spoon to the humanoid's mouth, managing orientation and approach speed to prevent spilling.
4. Bed Bathing: The robot wipes the humanoid's arm, visiting a set of distributed target points.
5. Arm Assist: The robot uses a hook to lift and reposition the humanoid's arm toward a target.

The environments incorporate realistic challenges such as simulated tremors, joint weakness, and limited range of motion.

Preference Modeling

To capture the diversity of human needs, Assistax introduces Preference Rewards. These augment the task reward based on three axes:

• Force: Preferred range of contact force.
• Speed: Preferred range of end-effector speed.
• Contact Count: Penalties for excessive re-contact (poking) to encourage sustained interaction.
  These preferences are encoded as observations for both agents, allowing the robot to infer and adapt to specific human constraints.

Training and Evaluation Pipeline

The paper proposes a two-stage pipeline:

1. MARL Co-training: The robot and humanoid agents are co-trained using standard MARL algorithms (IPPO, MAPPO, MASAC, and their recurrent variants) to solve the tasks collaboratively.

2. Ad-Hoc Teamwork (AHT) Benchmarking: To evaluate zero-shot coordination, the authors pre-train a diverse population of 630 humanoid partners per task. This diversity is induced by varying:

   • Behavioral conventions (via co-training dynamics).
   • Preference parameters (force, speed, contact).
   • Disability settings (mobility limitations).

   The benchmark then evaluates a robot policy trained on a subset of these partners (e.g., 5 agents) against a withheld test set (e.g., 625 unseen agents). This setup tests the robot's ability to generalize to novel preference combinations and disability profiles without further adaptation.

Key Contributions

1. Hardware-Accelerated Benchmark: Assistax is the first AHT benchmark for 3D continuous environments with two distinct embodied agents in an assistive context. By leveraging JAX and MuJoCo MJX, it achieves up to 412× faster open-loop simulation and 25× faster wall-clock training compared to comparable CPU-based environments (e.g., Assistive Gym) on a single GPU.
2. MARL Baselines and Implementation: The authors provide tuned baselines for PPO and SAC variants (including recurrent architectures) implemented in JAX, enabling scalable co-training of robot-humanoid teams.
3. AHT Pipeline and Data Release: The paper introduces a rigorous AHT evaluation pipeline and releases thousands of pre-trained humanoid partner policies via Hugging Face. This allows researchers to reproduce zero-shot coordination results and test new AHT algorithms against a diverse, pre-defined population of partners.

Results

• MARL Performance: In co-training scenarios, feed-forward implementations of IPPO and MAPPO consistently outperformed recurrent variants. MASAC also demonstrated strong performance. The introduction of preference rewards did not significantly alter the relative performance ranking of the baselines, though it added complexity to the reward landscape.
• AHT Coordination Gap: When evaluating existing RL algorithms against unseen partners with novel preference combinations (specifically in a "hard" setting where training and test distributions have minimal overlap), a significant coordination gap was observed. The robot policies struggled to adapt to out-of-distribution preferences. However, the paper notes that this gap closes significantly when the training and test distributions are more similar or when the training set is larger.
• Efficiency: A typical training run of 40 million timesteps with 1024 parallel environments completed in approximately 20 minutes, demonstrating the practical viability of hardware acceleration for complex MARL research.

Significance

The paper positions Assistax as a practical and necessary step forward for advancing RL in assistive care. By shifting from isolated agent training to multi-agent, hardware-accelerated environments, Assistax addresses the throughput bottlenecks that have historically limited research in complex robotics domains. Crucially, by framing the interaction as an Ad-Hoc Teamwork problem and providing a pipeline to evaluate zero-shot coordination with diverse human partners, Assistax offers a more realistic testbed for the "generalization to any user" goal central to assistive robotics. The release of pre-trained partner policies aims to standardize AHT evaluation, allowing the community to move beyond static benchmarks and rigorously test the adaptability of robotic policies in dynamic, human-centric scenarios.