class_docstr: Large-scale implementation of Augmented Random Search in Ray.
class_docstr_tok: ["Large-scale","implementation","of","Augmented","Random","Search","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.
      results, num_episodes, num_timesteps = self._collect_results(
          theta_id, config["num_rollouts"]
      )
      # 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)
  
      # keep only the best returns
      # select top performing directions if rollouts_used < num_rollouts
      max_rewards = np.max(noisy_returns, axis=1)
      if self.rollouts_used > self.num_rollouts:
          self.rollouts_used = self.num_rollouts
  
      percentile = 100 * (1 - (self.rollouts_used / self.num_rollouts))
      idx = np.arange(max_rewards.size)[
          max_rewards >= np.percentile(max_rewards, percentile)
      ]
      noise_idx = noise_indices[idx]
      noisy_returns = noisy_returns[idx, :]
  
      # Compute and take a step.
      g, count = utils.batched_weighted_sum(
          noisy_returns[:, 0] - noisy_returns[:, 1],
          (
              self.noise.get(index, self.policy.num_params)
              for index in noise_idx
          ),
          batch_size=min(500, noisy_returns[:, 0].size),
      )
      g /= noise_idx.size
      # scale the returns by their standard deviation
      if not np.isclose(np.std(noisy_returns), 0.0):
          g /= np.std(noisy_returns)
      assert (
          g.shape == (self.policy.num_params,) and g.dtype == np.float32
      )
      # Compute the new weights theta.
      theta, update_ratio = self.optimizer.update(-g)
  
      # 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)
      # update the reward list
      if len(all_eval_returns) > 0:
          self.reward_list.append(eval_returns.mean())
  
      # Bring restored workers back if necessary.
      self.restore_workers(self.workers)
  
      # Now sync the filters
      FilterManager.synchronize(
          {DEFAULT_POLICY_ID: self.policy.observation_filter},
          self.workers,
      )
  
      info = {
          "weights_norm": np.square(theta).sum(),
          "weights_std": np.std(theta),
          "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 Optional
  import numpy
  from rllib_ars.ars.ars_tf_policy import ARSTFPolicy
  import ray
  from ray.rllib.algorithms import Algorithm
  from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
  from ray.rllib.algorithms.es import optimizers, utils
  from ray.rllib.algorithms.es.es_tf_policy import rollout
  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.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
label_desc: Use image node_id 4 for calling the ARS obj's underlying member method code with example usage: obj.step() and returns: result