class_docstr: Update input callbacks to mutate batch with states entropy rewards.
class_docstr_tok: ["Update","input","callbacks","to","mutate","batch","with","states","entropy","rewards","."]
code: def on_learn_on_batch(
      self,
      *,
      policy: Policy,
      train_batch: SampleBatch,
      result: dict,
      **kwargs,
  ):
      super().on_learn_on_batch(
          policy=policy,
          train_batch=train_batch,
          result=result,
          **kwargs,
      )
      states_entropy = compute_states_entropy(
          train_batch[SampleBatch.OBS_EMBEDS],
          self.embeds_dim,
          self.k_nn,
      )
      states_entropy = update_beta(
          self.beta_schedule,
          self.beta,
          self.rho,
          RE3UpdateCallbacks._step,
      ) * np.reshape(
          self._rms(states_entropy),
          train_batch[SampleBatch.OBS_EMBEDS].shape[:-1],
      )
      train_batch[SampleBatch.REWARDS] = (
          train_batch[SampleBatch.REWARDS] + states_entropy
      )
      if Postprocessing.ADVANTAGES in train_batch:
          train_batch[Postprocessing.ADVANTAGES] = (
              train_batch[Postprocessing.ADVANTAGES] + states_entropy
          )
          train_batch[Postprocessing.VALUE_TARGETS] = (
              train_batch[Postprocessing.VALUE_TARGETS] + states_entropy
          )
  
imports: import gc
  import os
  import platform
  import tracemalloc
  from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, Union
  import gymnasium
  import numpy
  from ray.rllib.core.rl_module.rl_module import RLModule
  from ray.rllib.env.base_env import BaseEnv
  from ray.rllib.env.env_context import EnvContext
  from ray.rllib.evaluation.episode import Episode
  from ray.rllib.evaluation.episode_v2 import EpisodeV2
  from ray.rllib.evaluation.postprocessing import Postprocessing
  from ray.rllib.policy import Policy
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.utils.annotations import override, OverrideToImplementCustomLogic, PublicAPI
  from ray.rllib.utils.deprecation import deprecation_warning
  from ray.rllib.utils.exploration.random_encoder import _MovingMeanStd, compute_states_entropy, update_beta
  from ray.rllib.utils.typing import AgentID, EnvType, EpisodeType, PolicyID
  from ray.tune.callback import _CallbackMeta
  import psutil
label_desc: Use image node_id 2 for calling the RE3UpdateCallbacks obj's underlying member method code with example usage: obj.on_learn_on_batch() without return types