class_docstr: Module for embedding (observation, action, reward) inputs together.
docstr: Embed each of the (observation, action, reward) inputs & concatenate.
docstr_tok: ["Embed","each","of","the","(","observation",",","action",",","reward",")","inputs","&","concatenate","."]
code: def __call__(
      self, inputs: observation_action_reward.OAR
  ) -> jnp.ndarray:
      """Embed each of the (observation, action, reward) inputs & concatenate."""
  
      # Add dummy batch dimension to observation if necessary.
      # This is needed because Conv2D assumes a leading batch dimension, i.e.
      # that inputs are in [B, H, W, C] format.
      expand_obs = len(inputs.observation.shape) == 3
      if expand_obs:
          inputs = inputs._replace(
              observation=jnp.expand_dims(inputs.observation, axis=0)
          )
      features = self.torso(inputs.observation)  # [T?, B, D]
      if expand_obs:
          features = jnp.squeeze(features, axis=0)
  
      # Do a one-hot embedding of the actions.
      action = jax.nn.one_hot(
          inputs.action, num_classes=self.num_actions
      )  # [T?, B, A]
  
      # Map rewards -> [-1, 1].
      reward = jnp.tanh(inputs.reward)
  
      # Add dummy trailing dimensions to rewards if necessary.
      while reward.ndim < action.ndim:
          reward = jnp.expand_dims(reward, axis=-1)
  
      # Concatenate on final dimension.
      embedding = jnp.concatenate(
          [features, action, reward], axis=-1
      )  # [T?, B, D+A+1]
  
      return embedding
  
file_path: acme/acme/jax/networks/embedding.py
imports: import dataclasses
  from acme.wrappers import observation_action_reward
  import haiku
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling the OAREmbedding obj's underlying member method code with example usage: obj.__call__(inputs) and returns: embedding