code: def test_sampled_policy_distillation_loss(self):
      targets = policy_targets.PolicyTarget(
          actions=jnp.array([0, 1], dtype=jnp.int32),
          weights=jnp.array([0.5, 0.0], dtype=jnp.float32),
      )
      distribution = distrax.Categorical(
          probs=jnp.array([[0.7, 0.3], [0.9, 0.1]], dtype=jnp.float32)
      )
      loss = policy_targets.sampled_policy_distillation_loss(
          distribution=distribution, policy_targets=targets
      )
      expected_loss = 0.089169
      np.testing.assert_allclose(expected_loss, loss, atol=1e-4)
  
file_path: rlax/rlax/_src/policy_targets_test.py
imports: from absl.testing import absltest
  import distrax
  import jax
  import jax.numpy
  import numpy
  from rlax._src import policy_targets
label_desc: Use image node_id 1 for calling the PolicyTargetsTest obj's underlying member method code with example usage: obj.test_sampled_policy_distillation_loss() without return types