class_docstr: Custom model that adds an imitation loss on top of the policy loss.
class_docstr_tok: ["Custom","model","that","adds","an","imitation","loss","on","top","of","the","policy","loss","."]
code: def custom_loss(self, policy_loss, loss_inputs):
      # Create a new input reader per worker.
      reader = JsonReader(
          self.model_config["custom_model_config"]["input_files"]
      )
      input_ops = reader.tf_input_ops()
  
      # Define a secondary loss by building a graph copy with weight sharing.
      obs = restore_original_dimensions(
          tf.cast(input_ops["obs"], tf.float32), self.obs_space
      )
      logits, _ = self.forward({"obs": obs}, [], None)
  
      # Compute the IL loss.
      action_dist = Categorical(logits, self.model_config)
      self.policy_loss = policy_loss
      self.imitation_loss = tf.reduce_mean(
          -action_dist.logp(input_ops["actions"])
      )
      return policy_loss + 10 * self.imitation_loss
  
imports: import numpy
  from ray.rllib.models.modelv2 import ModelV2, restore_original_dimensions
  from ray.rllib.models.tf.tf_action_dist import Categorical
  from ray.rllib.models.tf.tf_modelv2 import TFModelV2
  from ray.rllib.models.tf.fcnet import FullyConnectedNetwork
  from ray.rllib.models.torch.torch_action_dist import TorchCategorical
  from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
  from ray.rllib.models.torch.fcnet import FullyConnectedNetwork
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.offline import JsonReader
label_desc: Use image node_id 4 for calling the CustomLossModel obj's underlying member method code with example usage: obj.custom_loss(policy_loss, loss_inputs) and returns: unknown