docstr: Tests whether checkpointing restores the state correctly.
code: def test_checkpointing(self):
      """Tests whether checkpointing restores the state correctly."""
      # Given an environment, and a model based on this environment.
      model = simulator.Simulator(catch.Catch())
      num_actions = model.action_spec().num_values
  
      model.reset()
  
      # Now, we save a checkpoint.
      model.save_checkpoint()
  
      ts = model.step(1)
  
      # Step the model once and load the checkpoint.
      timestep = model.step(np.random.randint(num_actions))
      model.load_checkpoint()
      self._check_equal(ts, model.step(1))
  
      while not timestep.last():
          timestep = model.step(np.random.randint(num_actions))
  
      # The model should require a reset.
      self.assertTrue(model.needs_reset)
  
      # Once we load checkpoint, the model should no longer require reset.
      model.load_checkpoint()
      self.assertFalse(model.needs_reset)
  
      # Further steps should agree with the original environment state.
      self._check_equal(ts, model.step(1))
  
file_path: acme/acme/agents/tf/mcts/models/simulator_test.py
imports: from acme.agents.tf.mcts.models import simulator
  from bsuite.environments import catch
  import dm_env
  import numpy
  from absl.testing import absltest
label_desc: Use image node_id 3 for calling the SimulatorTest obj's underlying member method code with example usage: obj.test_checkpointing() without return types