class_docstr: Generate a sample with the given search space and policy (factory) that is expected to work on the space.
  When the metric is ready, it returns the whole trajectory.
  
  The generator is single-threaded, meaning that it can't generate multiple samples,
  and wait for their rewards simultaneously.
  For parallelization, you can create multiple generators and share one policy.
  
  Parameters
  ----------
  policy
      The policy to use.
      For the first generator, as the environment is never created,
      it's recommended to use a factory.
      For the following generators, you can reuse the same policy.
  search_space
      The search space to search.
  
  Attributes
  ----------
  policy
      The policy it uses. Use the attribute as parameter of another generator to reuse the policy.
  env
      The environment it uses.
  sample
      The sample it just generated.
  sample_logits
      The logits of the sample.
class_docstr_tok: ["Generate","a","sample","with","the","given","search","space","and","policy","(","factory",")","that","is","expected","to","work","on","the","space",".","When","the","metric","is","ready",",","it","returns","the","whole","trajectory",".","The","generator","is","single-threaded",",","meaning","that","it","ca","n't","generate","multiple","samples",",","and","wait","for","their","rewards","simultaneously",".","For","parallelization",",","you","can","create","multiple","generators","and","share","one","policy",".","Parameters","--","--","--","--","--","policy","The","policy","to","use",".","For","the","first","generator",",","as","the","environment","is","never","created",",","it","'s","recommended","to","use","a","factory",".","For","the","following","generators",",","you","can","reuse","the","same","policy",".","search_space","The","search","space","to","search",".","Attributes","--","--","--","--","--","policy","The","policy","it","uses",".","Use","the","attribute","as","parameter","of","another","generator","to","reuse","the","policy",".","env","The","environment","it","uses",".","sample","The","sample","it","just","generated",".","sample_logits","The","logits","of","the","sample","."]
code: def __init__(
      self,
      search_space: Mutable,
      policy: PolicyFactory | BasePolicy | None = None,
  ) -> None:
      self.simplified_space: dict[
          str, Categorical | CategoricalMultiple
      ] = {}
      self.search_labels: list[tuple[str, int | None]] = []
      self.num_choices: list[int] = []
      for label, mutable in search_space.simplify().items():
          if isinstance(mutable, MutableAnnotation):
              # Skip annotations like constraints.
              continue
          # Expand CategoricalMultiple to Categorical by default.
          if isinstance(mutable, CategoricalMultiple):
              if mutable.n_chosen is None:
                  for i in range(len(mutable.values)):
                      self.search_labels.append((label, i))
                      self.num_choices.append(2)  # true and false
              elif mutable.n_chosen == 1:
                  self.search_labels.append((label, None))
                  self.num_choices.append(len(mutable.values))
              else:
                  # TODO: support categorical multiple.
                  raise ValueError(
                      "CategoricalMultiple with n_chosen > 1 is not supported yet."
                  )
              self.simplified_space[label] = mutable
          elif isinstance(mutable, Categorical):
              self.simplified_space[label] = mutable
              self.search_labels.append((label, None))
              # TODO: sampling priors are neglected.
              self.num_choices.append(len(mutable))
          else:
              raise ValueError(
                  "RL algorithms only supports Categorical for now."
              )
      self.env = TuningEnvironment(self.num_choices)
  
      if policy is None:
          self.policy = default_policy_fn(self.env)
      elif isinstance(policy, BasePolicy):
          self.policy = policy
      else:
          self.policy = policy(self.env)
  
      self._last_action: int | None = None
      self._trajectory: list[Batch] | None = None
      self._transition: Batch | None = None
  
      self.sample: Sample | None = None
      self.sample_logits: Sample | None = None
  
imports: from __future__ import annotations
  from copy import deepcopy
  from typing import Tuple, Callable
  import gym
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from typing_extensions import TypedDict
  from gym import spaces
  from tianshou.data import ReplayBuffer, Batch, to_torch, to_numpy
  from tianshou.policy import BasePolicy, PPOPolicy
  from nni.mutable import Categorical, CategoricalMultiple, Mutable, MutableAnnotation, Sample
label_desc: Use image node_id 1 to create a new TuningTrajectoryGenerator object  with example: obj = TuningTrajectoryGenerator(search_space, policy)