code: def main(_):
      key = jax.random.PRNGKey(FLAGS.seed)
      key_demonstrations, key_learner = jax.random.split(key, 2)
  
      # Create an environment and grab the spec.
      environment = gym_helpers.make_environment(task=FLAGS.env_name)
      environment_spec = specs.make_environment_spec(environment)
  
      # Get a demonstrations dataset.
      transitions_iterator = tfds.get_tfds_dataset(
          FLAGS.dataset_name, FLAGS.num_demonstrations
      )
      demonstrations = tfds.JaxInMemoryRandomSampleIterator(
          transitions_iterator,
          key=key_demonstrations,
          batch_size=FLAGS.batch_size,
      )
  
      # Create the networks to optimize.
      networks = cql.make_networks(environment_spec)
  
      # Create the learner.
      learner = cql.CQLLearner(
          batch_size=FLAGS.batch_size,
          networks=networks,
          random_key=key_learner,
          policy_optimizer=optax.adam(FLAGS.policy_learning_rate),
          critic_optimizer=optax.adam(FLAGS.critic_learning_rate),
          fixed_cql_coefficient=FLAGS.fixed_cql_coefficient,
          cql_lagrange_threshold=FLAGS.cql_lagrange_threshold,
          demonstrations=demonstrations,
          num_sgd_steps_per_step=1,
      )
  
      def evaluator_network(
          params: hk.Params, key: jax.Array, observation: jax.Array
      ) -> jax.Array:
          dist_params = networks.policy_network.apply(
              params, observation
          )
          return networks.sample_eval(dist_params, key)
  
      actor_core = actor_core_lib.batched_feed_forward_to_actor_core(
          evaluator_network
      )
      variable_client = variable_utils.VariableClient(
          learner, "policy", device="cpu"
      )
      evaluator = actors.GenericActor(
          actor_core, key, variable_client, backend="cpu"
      )
  
      eval_loop = acme.EnvironmentLoop(
          environment=environment,
          actor=evaluator,
          logger=loggers.TerminalLogger("evaluation", time_delta=0.0),
      )
  
      # Run the environment loop.
      while True:
          for _ in range(FLAGS.evaluate_every):
              learner.step()
          eval_loop.run(FLAGS.evaluation_episodes)
  
file_path: acme/examples/offline/run_cql_jax.py
imports: from absl import app
  from absl import flags
  import acme
  from acme import specs
  from acme.agents.jax import actor_core
  from acme.agents.jax import actors
  from acme.agents.jax import cql
  from acme.datasets import tfds
  from acme.examples.offline import helpers
  from acme.jax import variable_utils
  from acme.utils import loggers
  import haiku
  import jax
  import optax
label_desc: Use image node_id 1 for calling a global function with example usage: main(_) and returns: networks