code: def test_reinforcement_learning(
      named_model_space, engine, reward_for_invalid, caplog
  ):
      # Failure rate for the base case is approximately 3/50 (6%).
      # Sometimes the first 20 samples are within an awkward range, making the policy hard to optimize.
  
      # import torch, numpy
      # torch.manual_seed(seed)
      # numpy.random.seed(seed)
  
      name, model_space = named_model_space
      strategy_kwargs = dict(
          samples_per_update=20,
          replay_buffer_size=200,
          policy_fn=naive_policy,
          update_kwargs={
              "batch_size": 50,
              "repeat": 50,
              "update_times": 5,
          },
          reward_for_invalid=reward_for_invalid,
      )
      strategy = PolicyBasedRL(**strategy_kwargs)
      assert (
          repr(strategy)
          == f"PolicyBasedRL(samples_per_update=20, replay_buffer_size=200, reward_for_invalid={reward_for_invalid})"
      )
      if name == "numerical":
          with pytest.raises(ValueError, match="Categorical"):
              strategy(model_space, engine)
          return
  
      strategy(model_space, engine)
  
      # Failure rate for constraint + reward_for_invalid is approximately 9/50 (18%).
      # The odd condition implies mutual dependency between variables, unfriendly to the simple policy.
      if name != "base" and reward_for_invalid is not None:
          assert next(strategy.list_models()).metric >= 0.8
      else:
          assert next(strategy.list_models()).metric == 1.0
  
      if name == "constraint" and reward_for_invalid == -1:
          return  # FIXME: fails too often
  
      prev_models = list(engine.list_models())
      state_dict = strategy.state_dict()
      strategy2 = PolicyBasedRL(**strategy_kwargs)
      strategy2.load_state_dict(state_dict)
      assert "not restored" in caplog.text
      engine.max_model_count += 10
  
      strategy2.initialize(model_space, engine)
      strategy2.load_state_dict(state_dict)
  
      strategy2.run()
      new_models = [
          m for m in engine.list_models() if m not in prev_models
      ]
      assert len(new_models) == 10
      if name == "base":
          assert max(m.metric for m in new_models) == 1.0
  
imports: from __future__ import annotations
  import pytest
  import nni
  from nni.mutable import *
  from nni.nas.evaluator import FunctionalEvaluator
  from nni.nas.execution import SequentialExecutionEngine
  from nni.nas.space import *
  from nni.nas.strategy import GridSearch, Random, RegularizedEvolution, PolicyBasedRL
  import torch
  from torch import nn
  from tianshou.policy import A2CPolicy
  from tianshou.data import to_torch
label_desc: Use image node_id 8 for calling a global function with example usage: test_reinforcement_learning(named_model_space, engine, reward_for_invalid, caplog) without return types