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 held 98% GPU utilisation for three hours and forty minutes and got through 40M timesteps on two seeds. 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.

It sat at about 4.4 GB of the card’s 6 GB the whole time, which I read as comfortable headroom and which turned out to mean nothing at all. More on that below.

I never got a usable set of weights off it. Three attempts, three different ways to lose one.

Three runs, three ways to lose one

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.

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 past it.

Run three trained for three hours and forty minutes and then lost the weights to a DNS failure. This is the one that still annoys me. Training completed, the checkpoints were written, the artifact directory was assembled, and then the upload could not resolve wandb:

  Saved all_params.safetensors
wandb: Adding directory to artifact (/tmp/tmp0f5e9fvv)... Done. 1.0s
socket.gaierror: [Errno -3] Temporary failure in name resolution

wandb raised an AuthenticationError, Hydra tore the job down, and /tmp/tmp0f5e9fvv went with it. Same trap as run two, opposite cause: the first time I had switched the upload off myself, the second time my home connection switched it off for me.

So the lesson is not “decide whether you want wandb on”, which is what I thought after run two. It is that a temporary directory plus a network call is not a save path. Write the checkpoint somewhere real, then upload a copy of it. Anything else means an unrelated thirty-second network blip can charge you four hours.

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 run 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.

The same configuration on the L4, 1024 environments and two seeds and 40M timesteps off the same commit, trained in two hours flat against the laptop’s three hours forty. That is 1.8×, less than I expected from a datacentre card, and the L4 averaged 87% GPU utilisation against the laptop’s 98%. A thousand environments does not actually keep it busy.

The number that matters is underneath that one. The L4 run reported a peak working set of 7.7 GB, which is more memory than the 1660 Ti has in total. Both cards look like they were sitting at three quarters full, but that figure is XLA preallocating 75% of whatever card you hand it, and it stays flat for the whole run regardless of what the workload is doing. So the 6 GB card was never going to run the default 32-episode evaluation. Dropping to 8 was not tuning, it was the only setting that fit, which is what I mean about the ceiling dictating the experiment rather than the other way round.

	GTX 1660 Ti	L4
Parallel environments	1024	1024
Training, 40M timesteps	3 h 39 m	2 h 00 m
Throughput, per seed	3,040 timesteps/s	5,550 timesteps/s
Preallocated by XLA	4.4 of 6 GB	16.8 of 22.5 GB
Peak working set	not reported	7.7 GB
Evaluation episodes	8	32

The laptop never reported a peak because it died before printing one. The L4 went on to run the full 32-episode evaluation in another forty-five minutes and upload the model, the evaluation data and the trajectories, which is where every figure in that table comes from.

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. The whole GCP bill for this came to about £3, and that includes leaving the VM running longer than I needed because I was new to it. The laptop’s last attempt cost three hours forty and produced nothing.

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.