code: def build_experiment_config() -> experiments.ExperimentConfig:
      """Returns a configuration for GAIL/DAC experiments."""
  
      # Create an environment, grab the spec, and use it to create networks.
      environment = helpers.make_environment(task=FLAGS.env_name)
      environment_spec = specs.make_environment_spec(environment)
  
      # Create the direct RL agent.
      td3_config = td3.TD3Config(
          min_replay_size=1, samples_per_insert_tolerance_rate=2.0
      )
      td3_networks = td3.make_networks(environment_spec)
  
      # Create the discriminator.
      def discriminator(*args, **kwargs) -> networks_lib.Logits:
          return ail.DiscriminatorModule(
              environment_spec=environment_spec,
              use_action=True,
              use_next_obs=False,
              network_core=ail.DiscriminatorMLP(
                  hidden_layer_sizes=[
                      64,
                  ],
                  spectral_normalization_lipschitz_coeff=1.0,
              ),
          )(*args, **kwargs)
  
      discriminator_transformed = hk.without_apply_rng(
          hk.transform_with_state(discriminator)
      )
  
      def network_factory(
          environment_spec: specs.EnvironmentSpec,
      ) -> ail.AILNetworks:
          return ail.AILNetworks(
              ail.make_discriminator(
                  environment_spec, discriminator_transformed
              ),
              # reward balance = 0 corresponds to the GAIL reward: -ln(1-D)
              imitation_reward_fn=ail.rewards.gail_reward(
                  reward_balance=0.0
              ),
              direct_rl_networks=td3_networks,
          )
  
      # Create demonstrations function.
      dataset_name = helpers.get_dataset_name(FLAGS.env_name)
      num_demonstrations = FLAGS.num_demonstrations
  
      def make_demonstrations(batch_size, seed: int = 0):
          transitions_iterator = tfds.get_tfds_dataset(
              dataset_name,
              num_demonstrations,
              env_spec=environment_spec,
          )
          return tfds.JaxInMemoryRandomSampleIterator(
              transitions_iterator, jax.random.PRNGKey(seed), batch_size
          )
  
      # Create DAC agent.
      ail_config = ail.AILConfig(
          direct_rl_batch_size=td3_config.batch_size
          * td3_config.num_sgd_steps_per_step
      )
  
      env_name = FLAGS.env_name
  
      def environment_factory(seed: int) -> dm_env.Environment:
          del seed
          return helpers.make_environment(task=env_name)
  
      td3_builder = td3.TD3Builder(td3_config)
  
      dac_loss = ail.losses.add_gradient_penalty(
          ail.losses.gail_loss(entropy_coefficient=1e-3),
          gradient_penalty_coefficient=10.0,
          gradient_penalty_target=1.0,
      )
  
      ail_builder = ail.AILBuilder(
          rl_agent=td3_builder,
          config=ail_config,
          discriminator_loss=dac_loss,
          make_demonstrations=make_demonstrations,
      )
  
      return experiments.ExperimentConfig(
          builder=ail_builder,
          environment_factory=environment_factory,
          network_factory=network_factory,
          seed=FLAGS.seed,
          max_num_actor_steps=FLAGS.num_steps,
      )
  
file_path: acme/examples/baselines/imitation/run_gail.py
imports: from absl import flags
  from acme import specs
  from acme.agents.jax import ail
  from acme.agents.jax import td3
  from acme.datasets import tfds
  import helpers
  from absl import app
  from acme.jax import experiments
  from acme.jax import networks
  from acme.utils import lp_utils
  import dm_env
  import haiku
  import jax
  import launchpad
label_desc: Use image node_id 1 for calling a global function with example usage: build_experiment_config() and returns: experiments, ail, ail, tfds, helpers