Running Assistax without a cluster

Assistax is a hardware-accelerated benchmark for assistive robotics, a set of care tasks where a robot has to help a simulated human who is an active agent with their own preferences, rather than a mannequin that holds still. It runs on MuJoCo MJX under JAX, and the headline claim is a very large speedup from vectorising thousands of environments onto one GPU.

I build soft robotic rehabilitation hardware for a living, so a simulator that treats the patient as an agent is directly interesting to me. I wanted to run it rather than read about it. This is what that cost.

The Mac was a dead end

My main machine is a MacBook Pro M2 Pro. JAX’s Metal backend is not a viable path for this: jax-metal support has been effectively stalled, and MJX wants CUDA. There is no clever workaround here, if you are on Apple silicon and you want to run this, you need another machine or a rented GPU. I spent longer than I should have confirming that.

The gaming laptop got further than expected

The machine that worked first was a Lenovo Legion with a GTX 1660 Ti, 6 GB of VRAM, CUDA compute capability 7.5. I dual-booted Ubuntu, installed the Nvidia drivers, and got the thing JAX expects:

>>> import jax; jax.devices()
[CudaDevice(id=0)]

Training IPPO on the scratch-itch task with 1024 parallel environments sat at roughly full GPU utilisation and about 4.4 GB of VRAM. A 40M-timestep run took a little under four hours. For a five-year-old laptop GPU that is a genuinely surprising result, and it is the paper’s central point made concrete: the vectorised JAX implementation moves the bottleneck somewhere you can actually afford.

Getting a usable run out of it took three attempts.

Three runs, two of them wasted

Run one died with my SSH session. I launched training over SSH from the Mac, the connection dropped, and the process took SIGHUP with it. Entirely my fault, and the fix is the oldest advice in the book: start long jobs inside tmux, always. I also added keepalives on the client side:

# ~/.ssh/config
Host legion
  ServerAliveInterval 30
  ServerAliveCountMax 6

Run two completed and threw the weights away. This one took a while to understand. Training finished, the metrics looked reasonable, and there were no saved parameters anywhere on disk. The reason is that the save path writes into a tempfile.TemporaryDirectory() and then uploads the contents as a Weights & Biases artifact. I had been running with WANDB_MODE=disabled because I did not want to think about accounts yet, so the upload was a no-op, the temp directory was cleaned up, and four hours of compute evaporated.

If you are reproducing this, decide up front whether you want wandb on. If you do not, patch the save function to write somewhere persistent before you start a long run, not after.

Evaluation then hit an out-of-memory error. The default evaluation setting runs a large batch of episodes at once, which is fine on a datacentre card and not fine on 6 GB. Dropping NUM_EVAL_EPISODES from 32 to 8 got evaluation through. Run three, launched in tmux with wandb online, produced artifacts that still exist.

Where the laptop stopped being enough

The 1660 Ti is fine for one task and one algorithm. It is not fine for the thing the benchmark is actually built to measure, zero-shot coordination against a large held-out population of simulated partners, because that means many runs, not one, and the memory ceiling starts dictating configuration choices rather than the experiment doing so.

So the next runs went to an L4 on GCP. Two practical notes for anyone following: a Google Cloud free trial account cannot attach GPUs at all, and even on a paid account your GPU quota starts at zero and needs a manual increase request. Budget a day for that paperwork before you budget any hours for training.

TODO before publishing: fill in the L4 numbers, wall-clock for a comparable 40M-timestep run, environments per step, peak memory, and cost per run. Compare against the laptop figures above.

What I would tell someone starting today

Rent the GPU. I learned a lot from making a laptop work and I do not regret it, but if the goal is to understand the benchmark rather than to understand your own driver stack, an L4 for a few hours costs less than an evening.

Everything else is unglamorous discipline: run inside tmux, verify that your checkpoints land on a real filesystem before committing hours to a run, and check memory-dependent defaults against the card you actually have rather than the one the authors had.

The code is at assistive-autonomy/assistax, and it ran on everything I pointed it at once I stopped fighting my own environment.