---

name: jax-agents description: "Select, configure, and adapt Acme JAX agents and JAX experiment workflows." disable-model-invocation: true metadata: disco-role: operating license: Apache 2.0

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 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 for a safe local recommendation helper; it uses only Python standard-library imports and --help does not import Acme.
• Use 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 for import paths, config classes, builder classes, network factories, experiment runner signatures, actor utilities, and saving helpers.
• Use 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.