code: def update(self, env_outputs, action, stats):
      self.prev_action = action
      self.running_reward += env_outputs["reward"]
      self.weighted_returns *= self.discounting
      self.weighted_returns += env_outputs["reward"]
      self.weighted_returns_rms.update(self.weighted_returns)
  
      self.scaled_reward = env_outputs["reward"] / torch.sqrt(
          self.weighted_returns_rms.var + 1e-8
      )
  
      self.step_count += 1
  
      done = env_outputs["done"]
  
      episode_return = self.running_reward * done
      episode_step = self.step_count * done
  
      episodes_done = done.sum().item()
      if episodes_done > 0:
          stats["mean_episode_return"] += (
              episode_return.sum().item() / episodes_done
          )
          stats["mean_episode_step"] += (
              episode_step.sum().item() / episodes_done
          )
      stats["steps_done"] += done.numel()
      stats["episodes_done"] += episodes_done
  
      stats["running_reward"] += self.running_reward.mean().item()
      stats["running_step"] += self.step_count.mean().item()
  
      not_done = ~done
  
      self.running_reward *= not_done
      self.weighted_returns *= not_done
      self.step_count *= not_done
  
imports: import copy
  import dataclasses
  import logging
  import torch
  import moolib
label_desc: Use image node_id 2 for calling the EnvBatchState obj's underlying member method code with example usage: obj.update(env_outputs, action, stats) without return types