code: def run(*, trajectories_per_actor, num_actors, unroll_len):
      """Runs the example."""
  
      # Construct the agent network. We need a sample environment for its spec.
      env = catch.Catch()
      num_actions = env.action_spec().num_values
      net = hk.without_apply_rng(
          hk.transform(lambda ts: SimpleNet(num_actions)(ts))
      )  # pylint: disable=unnecessary-lambda
  
      # Construct the agent and learner.
      agent = Agent(net.apply)
      opt = optax.rmsprop(5e-3, decay=0.99, eps=1e-7)
      learner = Learner(agent, opt.update)
  
      # Initialize the optimizer state.
      sample_ts = env.reset()
      sample_ts = preprocess_step(sample_ts)
      ts_with_batch = jax.tree_util.tree_map(
          lambda t: np.expand_dims(t, 0), sample_ts
      )
      params = jax.jit(net.init)(jax.random.PRNGKey(428), ts_with_batch)
      opt_state = opt.init(params)
  
      # Create accessor and queueing functions.
      current_params = lambda: params
      batch_size = 2
      q = queue.Queue(maxsize=batch_size)
  
      def dequeue():
          batch = []
          for _ in range(batch_size):
              batch.append(q.get())
          batch = jax.tree_util.tree_map(
              lambda *ts: np.stack(ts, axis=1), *batch
          )
          return jax.device_put(batch)
  
      # Start the actors.
      for i in range(num_actors):
          key = jax.random.PRNGKey(i)
          args = (
              agent,
              key,
              current_params,
              q.put,
              unroll_len,
              trajectories_per_actor,
          )
          threading.Thread(target=run_actor, args=args).start()
  
      # Run the learner.
      num_steps = num_actors * trajectories_per_actor // batch_size
      for i in range(num_steps):
          traj = dequeue()
          params, opt_state = learner.update(params, opt_state, traj)
  
imports: import functools
  import queue
  import threading
  from typing import Any, Callable, NamedTuple
  from absl import app
  from absl import logging
  from bsuite.environments import catch
  import dm_env
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  import rlax
label_desc: Use image node_id 3 for calling a global function with example usage: run() and returns: jax