code: def embed_oar(
      features: Array, action: Array, reward: Array, num_actions: int
  ) -> Array:
      """Embed each of the (observation, action, reward) inputs & concatenate."""
      chex.assert_rank([features, action, reward], [2, 1, 1])
      action = jax.nn.one_hot(action, num_classes=num_actions)  # [B, A]
  
      reward = jnp.tanh(reward)
      while reward.ndim < action.ndim:
          reward = jnp.expand_dims(reward, axis=-1)
  
      embedding = jnp.concatenate(
          [features, action, reward], axis=-1
      )  # [B, D+A+1]
      return embedding
  
label_desc: Use image node_id 1 for calling a global function with example usage: embed_oar(features, action, reward, num_actions) and returns: embedding