code: def get_per_learner_step_batch_size(config: AILConfig) -> int:
      """Returns how many transitions should be sampled per direct learner step."""
      # If the iterators are tied, the discriminator learning batch size has to
      # match the direct RL one.
      if config.share_iterator:
          assert (
              config.direct_rl_batch_size
              % config.discriminator_batch_size
          ) == 0
          return config.direct_rl_batch_size
      # Otherwise each iteration of the discriminator will sample a batch which will
      # be split in num_sgd_steps_per_step batches, each of size
      # discriminator_batch_size.
      return (
          config.discriminator_batch_size
          * config.num_sgd_steps_per_step
      )
  
imports: import dataclasses
  from typing import Optional
  import optax
label_desc: Use image node_id 1 for calling a global function with example usage: get_per_learner_step_batch_size(config) and returns: unknown, config