code: def _hk_mock_policy_params(s_tm1):
      """Returns mock policy params."""
      # Outputs of the network are mu and sigma. Both shaped [B, ACTION_DIM].
      pi_out = hk.nets.MLP(
          output_sizes=[2 * ACTION_DIM],
          w_init=hk.initializers.VarianceScaling(1e-3),
          activation=jnp.tanh,
          activate_final=False,
          name="online_policy",
      )(s_tm1)
      pi_mean, pi_cov = jnp.split(pi_out, 2, axis=-1)
      pi_cov = jax.nn.softplus(pi_cov)
      pi_mean = pi_mean + _MEAN_OFFSET
      return {"mean": pi_mean, "stddev": pi_cov}
  
imports: import functools
  import math
  from typing import Tuple
  from absl.testing import absltest
  from absl.testing import parameterized
  import chex
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  from rlax._src import distributions
  from rlax._src import mpo_ops
label_desc: Use image node_id 1 for calling a global function with example usage: _hk_mock_policy_params(s_tm1) and returns: dict