code: def _init_params(key):
      init_fn, _ = hk.transform(_hk_mock_policy_params)
      key_seq = hk.PRNGSequence(key)
      s_tm1 = jax.random.normal(
          next(key_seq), (TIME_DIM, BATCH_DIM, ACTION_DIM), jnp.float32
      )
      online_params = init_fn(next(key_seq), s_tm1)
      return dict(
          online=online_params,
          target=online_params,
          mpo=dict(
              temperature=jnp.array(_INIT_TEMPERATURE),
              alpha_mean=_INIT_ALPHA_MEAN,
              alpha_covariance=_INIT_ALPHA_COVARIANCE,
          ),
      )
  
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 2 for calling a global function with example usage: _init_params(key) and returns: dict