jax-agents

# Acme JAX Agents Use this sub-skill when the task is to choose, configure, or adapt `acme.agents.jax` algorithms or `acme.jax.experiments` runners for continuous-control, discrete-control, offline RL, imitation, learning-from-demonstrations, model-based, or decentralized multiagent workflows. ## Quick Routing - Start with [`references/agent-catalog.md`](references/agent-catalog.md) to choose the JAX algorithm family from task signals such as action space, online vs offline data, demonstrations, model-based planning, and multiagent structure. - Use [`scripts/select_jax_agent.py`](scripts/select_jax_agent.py) for a safe local recommendation helper; it uses only Python standard-library imports and `--help` does not import Acme. - Use [`references/experiment-workflows.md`](references/experiment-workflows.md) to build `ExperimentConfig` or `OfflineExperimentConfig`, switch between local and Launchpad distributed execution, add evaluators, checkpointing, or offline datasets. - Use [`references/api-reference.md`](references/api-reference.md) for import paths, config classes, builder classes, network factories, experiment runner signatures, actor utilities, and saving helpers. - Use [`references/troubleshooting.md`](references/troubleshooting.md) before diagnosing dependency imports, JAX/JAXLIB wheels, Launchpad/Reverb/TF extras, PRNG handling, spec/network mismatches, and expensive examples. ## Boundaries - Route generic `acme.EnvironmentLoop`, environment wrappers, logging loops, and package-level setup to the Acme core workflows skill. - Route replay table, adders, dataset iterator, TFDS/RLDS, and Reverb data-pipeline internals to the replay/data skill unless the issue is selecting a JAX builder or offline experiment shape. - Route TensorFlow agent selection or `acme.agents.tf.*` code to the TF agents skill. - Do not require users to inspect Acme source docs, examples, or tests at runtime; the API names and workflow patterns needed for JAX agent adaptation are distilled in the bundled references.