code: def test_train_step_updates_weights(self, agent_name="dr3"):
      create_agent_fn = self._create_agent_fn(agent_name)
      agent = create_agent_fn(
          4,
          "unused_string",
          replay_buffer_builder=self.create_replay_buffer,
      )
  
      # We skip sampling from the replay buffer with a mock, and set the
      # replay elements directly to the batch we want to learn from.
      agent._sample_from_replay_buffer = mock.create_autospec(
          agent._sample_from_replay_buffer, spec_set=True
      )
      agent.replay_elements = self.batch
  
      params_before = agent.online_params
      agent.train_step()
      params_after = agent.online_params
  
      for i, (param1, param2) in enumerate(
          zip(
              jax.tree_leaves(params_before),
              jax.tree_leaves(params_after),
          )
      ):
          with self.subTest("param_set_{}".format(i)):
              self.assertTrue((param1 != param2).any())
  
file_path: dopamine/tests/dopamine/labs/offline_rl/jax/offline_agent_test.py
imports: import functools
  from unittest import mock
  from absl.testing import absltest
  from absl.testing import parameterized
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.labs.offline_rl import fixed_replay
  from dopamine.labs.offline_rl.jax import networks
  from dopamine.labs.offline_rl.jax import offline_dr3_agent
  import gin
  import jax
  from jax import numpy
label_desc: Use image node_id 3 for calling the OfflineAgentTest obj's underlying member method code with example usage: obj.test_train_step_updates_weights(agent_name) without return types