class_docstr: Discriminator module that concatenates its inputs.
class_docstr_tok: ["Discriminator","module","that","concatenates","its","inputs","."]
code: def __call__(
      self,
      observations: networks_lib.Observation,
      actions: networks_lib.Action,
      next_observations: networks_lib.Observation,
      is_training: bool,
      rng: networks_lib.PRNGKey,
  ) -> networks_lib.Logits:
      observations = self._observation_embedding(observations)
      if self._use_next_obs:
          next_observations = self._observation_embedding(
              next_observations
          )
          data = jnp.concatenate(
              [observations, next_observations], axis=-1
          )
      else:
          data = observations
      if self._use_action:
          action_spec = self._environment_spec.actions
          if isinstance(action_spec, specs.DiscreteArray):
              actions = jax.nn.one_hot(actions, action_spec.num_values)
          data = jnp.concatenate([data, actions], axis=-1)
      output = self._network_core(data, is_training, rng)
      output = jnp.squeeze(output, axis=-1)
      return output
  
imports: import dataclasses
  import functools
  from typing import Any, Callable, Generic, Iterable, Optional
  from acme import specs
  from acme import types
  from acme.jax import networks
  from acme.jax import utils
  from acme.jax.imitation_learning_types import DirectRLNetworks
  import haiku
  import jax
  from jax import numpy
  import numpy
label_desc: Use image node_id 2 for calling the DiscriminatorModule obj's underlying member method code with example usage: obj.__call__(observations, actions, next_observations, is_training, rng) and returns: output