code: def test_hidden_state_policy():
      search_space = CategoricalMultiple([1, 2], label="x")
  
      replay_buffer = ReplayBuffer(20)
      generator = TuningTrajectoryGenerator(
          search_space, HiddenStatePolicy
      )
      policy = generator.policy
  
      for _ in range(20):
          sample = generator.next_sample()
          assert not generator.sample_logits
          reward = sum(search_space.freeze(sample))
          trajectory = generator.send_reward(reward)
          assert np.all(
              trajectory.policy.hidden_state == np.array([0, 2])
          )
          assert np.all(trajectory.rew == np.array([0, 3]))
          assert np.all(trajectory.done == np.array([False, True]))
          replay_buffer.update(trajectory)
  
      policy.update(0, replay_buffer)
  
file_path: nni/test/algo/nas/strategy/test_rl_impl.py
imports: from functools import partial
  import numpy
  import pytest
  from tianshou.data import ReplayBuffer, Batch
  from tianshou.policy import BasePolicy
  from nni.mutable import Categorical, Numerical, CategoricalMultiple, MutableDict, MutableList, ExpressionConstraint
  from nni.nas.strategy._rl_impl import TuningTrajectoryGenerator, default_policy_fn
label_desc: Use image node_id 6 for calling a global function with example usage: test_hidden_state_policy() without return types