class_docstr: ## Trainer
class_docstr_tok: ["#","#","Trainer"]
docstr: ### Sample data with current policy
docstr_tok: ["#","#","#","Sample","data","with","current","policy"]
code: def sample(self) -> Dict[str, torch.Tensor]:
      """
      ### Sample data with current policy
      """
  
      rewards = np.zeros(
          (self.n_workers, self.worker_steps), dtype=np.float32
      )
      actions = np.zeros(
          (self.n_workers, self.worker_steps), dtype=np.int32
      )
      done = np.zeros(
          (self.n_workers, self.worker_steps), dtype=np.bool
      )
      obs = np.zeros(
          (self.n_workers, self.worker_steps, 4, 84, 84), dtype=np.uint8
      )
      log_pis = np.zeros(
          (self.n_workers, self.worker_steps), dtype=np.float32
      )
      values = np.zeros(
          (self.n_workers, self.worker_steps + 1), dtype=np.float32
      )
  
      with torch.no_grad():
          # sample `worker_steps` from each worker
          for t in range(self.worker_steps):
              # `self.obs` keeps track of the last observation from each worker,
              #  which is the input for the model to sample the next action
              obs[:, t] = self.obs
              # sample actions from $\pi_{\theta_{OLD}}$ for each worker;
              #  this returns arrays of size `n_workers`
              pi, v = self.model(obs_to_torch(self.obs))
              values[:, t] = v.cpu().numpy()
              a = pi.sample()
              actions[:, t] = a.cpu().numpy()
              log_pis[:, t] = pi.log_prob(a).cpu().numpy()
  
              # run sampled actions on each worker
              for w, worker in enumerate(self.workers):
                  worker.child.send(("step", actions[w, t]))
  
              for w, worker in enumerate(self.workers):
                  # get results after executing the actions
                  (
                      self.obs[w],
                      rewards[w, t],
                      done[w, t],
                      info,
                  ) = worker.child.recv()
  
                  # collect episode info, which is available if an episode finished;
                  #  this includes total reward and length of the episode -
                  #  look at `Game` to see how it works.
                  if info:
                      tracker.add("reward", info["reward"])
                      tracker.add("length", info["length"])
  
          # Get value of after the final step
          _, v = self.model(obs_to_torch(self.obs))
          values[:, self.worker_steps] = v.cpu().numpy()
  
      # calculate advantages
      advantages = self.gae(done, rewards, values)
  
      #
      samples = {
          "obs": obs,
          "actions": actions,
          "values": values[:, :-1],
          "log_pis": log_pis,
          "advantages": advantages,
      }
  
      # samples are currently in `[workers, time_step]` table,
      # we should flatten it for training
      samples_flat = {}
      for k, v in samples.items():
          v = v.reshape(v.shape[0] * v.shape[1], *v.shape[2:])
          if k == "obs":
              samples_flat[k] = obs_to_torch(v)
          else:
              samples_flat[k] = torch.tensor(v, device=device)
  
      return samples_flat
  
imports: from typing import Dict
  import numpy
  import torch
  from torch import nn
  from torch import optim
  from torch.distributions import Categorical
  from labml import monit, tracker, logger, experiment
  from labml.configs import FloatDynamicHyperParam, IntDynamicHyperParam
  from labml_helpers.module import Module
  from labml_nn.rl.game import Worker
  from labml_nn.rl.ppo import ClippedPPOLoss, ClippedValueFunctionLoss
  from labml_nn.rl.ppo.gae import GAE
label_desc: Use image node_id 2 for calling the Trainer obj's underlying member method code with example usage: obj.sample() and returns: samples_flat