class_docstr: Importance weighted actor/learner architecture (IMPALA) Algorithm
  
  == Overview of data flow in IMPALA ==
  1. Policy evaluation in parallel across `num_workers` actors produces
     batches of size `rollout_fragment_length * num_envs_per_worker`.
  2. If enabled, the replay buffer stores and produces batches of size
     `rollout_fragment_length * num_envs_per_worker`.
  3. If enabled, the minibatch ring buffer stores and replays batches of
     size `train_batch_size` up to `num_sgd_iter` times per batch.
  4. The learner thread executes data parallel SGD across `num_gpus` GPUs
     on batches of size `train_batch_size`.
class_docstr_tok: ["Importance","weighted","actor\/learner","architecture","(","IMPALA",")","Algorithm","==","Overview","of","data","flow","in","IMPALA","==","1",".","Policy","evaluation","in","parallel","across","`","num_workers","`","actors","produces","batches","of","size","`","rollout_fragment_length","*","num_envs_per_worker","`",".","2",".","If","enabled",",","the","replay","buffer","stores","and","produces","batches","of","size","`","rollout_fragment_length","*","num_envs_per_worker","`",".","3",".","If","enabled",",","the","minibatch","ring","buffer","stores","and","replays","batches","of","size","`","train_batch_size","`","up","to","`","num_sgd_iter","`","times","per","batch",".","4",".","The","learner","thread","executes","data","parallel","SGD","across","`","num_gpus","`","GPUs","on","batches","of","size","`","train_batch_size","`","."]
code: def training_step(self) -> ResultDict:
      # First, check, whether our learner thread is still healthy.
      if (
          not self.config._enable_new_api_stack
          and not self._learner_thread.is_alive()
      ):
          raise RuntimeError("The learner thread died while training!")
  
      use_tree_aggregation = (
          self._aggregator_actor_manager
          and self._aggregator_actor_manager.num_healthy_actors() > 0
      )
  
      # Get sampled SampleBatches from our workers (by ray references if we use
      # tree-aggregation).
      unprocessed_sample_batches = self.get_samples_from_workers(
          return_object_refs=use_tree_aggregation,
      )
      # Tag workers that actually produced ready sample batches this iteration.
      # Those workers will have to get updated at the end of the iteration.
      workers_that_need_updates = {
          worker_id for worker_id, _ in unprocessed_sample_batches
      }
  
      # Send the collected batches (still object refs) to our aggregation workers.
      if use_tree_aggregation:
          batches = self.process_experiences_tree_aggregation(
              unprocessed_sample_batches
          )
      # Resolve collected batches here on local process (using the mixin buffer).
      else:
          batches = self.process_experiences_directly(
              unprocessed_sample_batches
          )
  
      # Increase sampling counters now that we have the actual SampleBatches on
      # the local process (and can measure their sizes).
      for batch in batches:
          self._counters[NUM_ENV_STEPS_SAMPLED] += batch.count
          self._counters[NUM_AGENT_STEPS_SAMPLED] += batch.agent_steps()
      # Concatenate single batches into batches of size `train_batch_size`.
      self.concatenate_batches_and_pre_queue(batches)
      # Using the Learner API. Call `update()` on our LearnerGroup object with
      # all collected batches.
      if self.config._enable_new_api_stack:
          train_results = self.learn_on_processed_samples()
          module_ids_to_update = set(train_results.keys()) - {
              ALL_MODULES
          }
          additional_results = self.learner_group.additional_update(
              module_ids_to_update=module_ids_to_update,
              timestep=self._counters[
                  NUM_ENV_STEPS_TRAINED
                  if self.config.count_steps_by == "env_steps"
                  else NUM_AGENT_STEPS_TRAINED
              ],
              # TODO (sven): Feels hacked, but solves the problem of algos inheriting
              #  from IMPALA (like APPO). In the old stack, we didn't have this
              #  problem b/c IMPALA didn't need to call any additional update methods
              #  as the entropy- and lr-schedules were handled by
              #  `Policy.on_global_var_update()`.
              **self._get_additional_update_kwargs(train_results),
          )
          for key, res in additional_results.items():
              if key in train_results:
                  train_results[key].update(res)
      else:
          # Move train batches (of size `train_batch_size`) onto learner queue.
          self.place_processed_samples_on_learner_thread_queue()
          # Extract most recent train results from learner thread.
          train_results = self.process_trained_results()
  
      # Sync worker weights (only those policies that were actually updated).
      with self._timers[SYNCH_WORKER_WEIGHTS_TIMER]:
          if self.config._enable_new_api_stack:
              if train_results:
                  pids = list(set(train_results.keys()) - {ALL_MODULES})
                  self.update_workers_from_learner_group(
                      workers_that_need_updates=workers_that_need_updates,
                      policy_ids=pids,
                  )
          else:
              pids = list(train_results.keys())
              self.update_workers_if_necessary(
                  workers_that_need_updates=workers_that_need_updates,
                  policy_ids=pids,
              )
  
      # With a training step done, try to bring any aggregators back to life
      # if necessary.
      # Aggregation workers are stateless, so we do not need to restore any
      # state here.
      if self._aggregator_actor_manager:
          self._aggregator_actor_manager.probe_unhealthy_actors(
              timeout_seconds=self.config.worker_health_probe_timeout_s,
              mark_healthy=True,
          )
  
      if self.config._enable_new_api_stack:
          if train_results:
              # Store the most recent result and return it if no new result is
              # available. This keeps backwards compatibility with the old
              # training stack / results reporting stack. This is necessary
              # any time we develop an asynchronous algorithm.
              self._results = train_results
          return self._results
      else:
          return train_results
  
imports: import copy
  from functools import partial
  import logging
  import platform
  import queue
  import random
  from typing import Callable, List, Optional, Set, Tuple, Type, Union
  import numpy
  import tree
  import ray
  from ray import ObjectRef
  from ray.rllib import SampleBatch
  from ray.rllib.algorithms.algorithm import Algorithm
  from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
  from ray.rllib.core.rl_module.rl_module import SingleAgentRLModuleSpec
  from ray.rllib.evaluation.worker_set import handle_remote_call_result_errors
  from ray.rllib.execution.buffers.mixin_replay_buffer import MixInMultiAgentReplayBuffer
  from ray.rllib.execution.learner_thread import LearnerThread
  from ray.rllib.execution.multi_gpu_learner_thread import MultiGPULearnerThread
  from ray.rllib.policy.policy import Policy
  from ray.rllib.policy.sample_batch import concat_samples
  from ray.rllib.utils.actor_manager import FaultAwareApply, FaultTolerantActorManager, RemoteCallResults
  from ray.rllib.utils.actors import create_colocated_actors
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.metrics import ALL_MODULES
  from ray.rllib.utils.deprecation import DEPRECATED_VALUE, deprecation_warning
  from ray.rllib.utils.metrics import NUM_AGENT_STEPS_SAMPLED, NUM_AGENT_STEPS_TRAINED, NUM_ENV_STEPS_SAMPLED, NUM_ENV_STEPS_TRAINED, NUM_SYNCH_WORKER_WEIGHTS, NUM_TRAINING_STEP_CALLS_SINCE_LAST_SYNCH_WORKER_WEIGHTS, SYNCH_WORKER_WEIGHTS_TIMER, SAMPLE_TIMER
  from ray.rllib.utils.metrics.learner_info import LearnerInfoBuilder
  from ray.rllib.utils.replay_buffers.multi_agent_replay_buffer import ReplayMode
  from ray.rllib.utils.replay_buffers.replay_buffer import _ALL_POLICIES
  from ray.rllib.utils.schedules.scheduler import Scheduler
  from ray.rllib.utils.typing import PartialAlgorithmConfigDict, PolicyID, ResultDict, SampleBatchType
  from ray.tune.execution.placement_groups import PlacementGroupFactory
label_desc: Use image node_id 4 for calling the Impala obj's underlying member method code with example usage: obj.training_step() and returns: self, train_results