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","."]
docstr: Create a new trajectory, and return the sample when it's done.
  
  The sample isn't yet validated and thus can be possibly invalid.
  The caller needs to freeze the mutable with the sample by itself.
  
  The class will be in a state pending for reward after a call of :meth:`next_sample`.
  It will either receive the reward via :meth:`send_reward` or be reset via another :meth:`next_sample`.
docstr_tok: ["Create","a","new","trajectory",",","and","return","the","sample","when","it","'s","done",".","The","sample","is","n't","yet","validated","and","thus","can","be","possibly","invalid",".","The","caller","needs","to","freeze","the","mutable","with","the","sample","by","itself",".","The","class","will","be","in","a","state","pending","for","reward","after","a","call","of",":","meth",":","`","next_sample","`",".","It","will","either","receive","the","reward","via",":","meth",":","`","send_reward","`","or","be","reset","via","another",":","meth",":","`","next_sample","`","."]
code: def next_sample(self) -> Sample:
      """Create a new trajectory, and return the sample when it's done.
  
      The sample isn't yet validated and thus can be possibly invalid.
      The caller needs to freeze the mutable with the sample by itself.
  
      The class will be in a state pending for reward after a call of :meth:`next_sample`.
      It will either receive the reward via :meth:`send_reward` or be reset via another :meth:`next_sample`.
      """
      obs, info = self.env.reset()
      last_state = None  # hidden state
  
      self._trajectory = []
      self._transition = Batch(
          obs=obs,  # only obs is set.
          act={},  # the rest are empty.
          rew={},
          terminated={},
          truncated={},
          done={},
          obs_next={},
          info=info,
          policy={},
      )
  
      self.sample = {}
      self.sample_logits = {}
  
      step_count = 0
  
      while True:
          obs_batch = Batch(
              [self._transition]
          )  # the first dimension is batch-size
          policy_result = self.policy(obs_batch, last_state)
          # get bounded and remapped actions first (not saved into buffer)
          self._last_action = self.policy.map_action(
              to_numpy(policy_result.act)
          )[0]
          # TODO: exploration noise
  
          # Update the sample and the logits.
          last_label, label_count = self.search_labels[step_count]
          mutable = self.simplified_space[last_label]
          if "logits" in policy_result:
              logits = to_numpy(policy_result.logits)[0][
                  : self.num_choices[step_count]
              ].tolist()
          else:
              logits = None
  
          if label_count is None:
              # Categorical, or CategoryMultiple with n_chosen == 1.
              if isinstance(mutable, CategoricalMultiple):
                  self.sample[last_label] = [
                      mutable.values[self._last_action]
                  ]
              else:
                  self.sample[last_label] = mutable.values[
                      self._last_action
                  ]
              if logits is not None:
                  self.sample_logits[last_label] = logits
          else:
              # CategoryMultiple with n_chosen is None.
              if last_label not in self.sample:
                  self.sample[last_label] = []
              if self._last_action == 0:
                  self.sample[last_label].append(
                      mutable.values[label_count]
                  )
              if logits is not None:
                  if last_label not in self.sample_logits:
                      self.sample_logits[last_label] = []
                  self.sample_logits[last_label].append(logits[0])
  
          self._transition.update(
              policy=self._canonicalize_policy_data(policy_result),
              act=to_numpy(policy_result.act)[0],  # the unmapped action
          )
  
          step_count += 1
          if step_count == len(self.num_choices):
              return self.sample
  
          obs_next, rew, terminated, truncated, info = self.env.step(
              self._last_action
          )
          assert (
              not terminated
          ), "The environment should not be done yet."
  
          self._transition.update(
              obs_next=obs_next,
              rew=rew,
              terminated=terminated,
              truncated=truncated,
              done=terminated,
              info=info,
          )
  
          self._trajectory.append(deepcopy(self._transition))
  
          last_state = policy_result.get("state", None)
          self._transition.obs = self._transition.obs_next
  
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 3 for calling the TuningTrajectoryGenerator obj's underlying member method code with example usage: obj.next_sample() and returns: self