class_docstr: Large-scale implementation of Evolution Strategies in Ray.
class_docstr_tok: ["Large-scale","implementation","of","Evolution","Strategies","in","Ray","."]
code: def step(self):
      config = self.config
  
      theta = self.policy.get_flat_weights()
      assert theta.dtype == np.float32
      assert len(theta.shape) == 1
  
      # Put the current policy weights in the object store.
      theta_id = ray.put(theta)
      # Use the actors to do rollouts. Note that we pass in the ID of the
      # policy weights as these are shared.
      results, num_episodes, num_timesteps = self._collect_results(
          theta_id,
          config["episodes_per_batch"],
          config["train_batch_size"],
      )
      # Update our sample steps counters.
      self._counters[NUM_AGENT_STEPS_SAMPLED] += num_timesteps
      self._counters[NUM_ENV_STEPS_SAMPLED] += num_timesteps
  
      all_noise_indices = []
      all_training_returns = []
      all_training_lengths = []
      all_eval_returns = []
      all_eval_lengths = []
  
      # Loop over the results.
      for result in results:
          all_eval_returns += result.eval_returns
          all_eval_lengths += result.eval_lengths
  
          all_noise_indices += result.noise_indices
          all_training_returns += result.noisy_returns
          all_training_lengths += result.noisy_lengths
  
      assert len(all_eval_returns) == len(all_eval_lengths)
      assert (
          len(all_noise_indices)
          == len(all_training_returns)
          == len(all_training_lengths)
      )
  
      self.episodes_so_far += num_episodes
  
      # Assemble the results.
      eval_returns = np.array(all_eval_returns)
      eval_lengths = np.array(all_eval_lengths)
      noise_indices = np.array(all_noise_indices)
      noisy_returns = np.array(all_training_returns)
      noisy_lengths = np.array(all_training_lengths)
  
      # Process the returns.
      proc_noisy_returns = utils.compute_centered_ranks(noisy_returns)
  
      # Compute and take a step.
      g, count = utils.batched_weighted_sum(
          proc_noisy_returns[:, 0] - proc_noisy_returns[:, 1],
          (
              self.noise.get(index, self.policy.num_params)
              for index in noise_indices
          ),
          batch_size=500,
      )
      g /= noisy_returns.size
      assert (
          g.shape == (self.policy.num_params,)
          and g.dtype == np.float32
          and count == len(noise_indices)
      )
      # Compute the new weights theta.
      theta, update_ratio = self.optimizer.update(
          -g + config["l2_coeff"] * theta
      )
  
      # Update our train steps counters.
      self._counters[NUM_AGENT_STEPS_TRAINED] += num_timesteps
      self._counters[NUM_ENV_STEPS_TRAINED] += num_timesteps
  
      # Set the new weights in the local copy of the policy.
      self.policy.set_flat_weights(theta)
      # Store the rewards
      if len(all_eval_returns) > 0:
          self.reward_list.append(np.mean(eval_returns))
  
      # Bring restored workers back if necessary.
      # We will sync filters right next.
      self.workers.probe_unhealthy_workers()
  
      # Now sync the filters
      FilterManager.synchronize(
          {DEFAULT_POLICY_ID: self.policy.observation_filter},
          self.workers,
      )
  
      info = {
          "weights_norm": np.square(theta).sum(),
          "grad_norm": np.square(g).sum(),
          "update_ratio": update_ratio,
          "episodes_this_iter": noisy_lengths.size,
          "episodes_so_far": self.episodes_so_far,
      }
  
      reward_mean = np.mean(self.reward_list[-self.report_length :])
      result = {
          "sampler_results": {
              "episode_reward_mean": reward_mean,
              "episode_len_mean": eval_lengths.mean(),
          },
          "timesteps_this_iter": noisy_lengths.sum(),
          "info": info,
      }
  
      return result
  
imports: import logging
  import random
  import time
  from collections import namedtuple
  from typing import Dict, List, Optional
  import numpy
  from rllib_es.es import optimizers, utils
  from rllib_es.es.es_tf_policy import ESTFPolicy, rollout
  import ray
  from ray.rllib.algorithms import Algorithm
  from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
  from ray.rllib.env.env_context import EnvContext
  from ray.rllib.evaluation.worker_set import WorkerSet
  from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
  from ray.rllib.utils import FilterManager
  from ray.rllib.utils.actor_manager import FaultAwareApply
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.deprecation import Deprecated
  from ray.rllib.utils.metrics import NUM_AGENT_STEPS_SAMPLED, NUM_AGENT_STEPS_TRAINED, NUM_ENV_STEPS_SAMPLED, NUM_ENV_STEPS_TRAINED
  from ray.rllib.utils.torch_utils import set_torch_seed
  from ray.rllib.utils.typing import PolicyID
label_desc: Use image node_id 4 for calling the ES obj's underlying member method code with example usage: obj.step() and returns: result