docstr: Sample one-step dynamics based on the provided world model
docstr_tok: ["Sample","one-step","dynamics","based","on","the","provided","world","model"]
code: def sample_reward_next_state_terminal(
      self,
      state: rlt.FeatureData,
      action: rlt.FeatureData,
      mem_net: MemoryNetwork,
  ):
      """Sample one-step dynamics based on the provided world model"""
      wm_output = mem_net(state, action)
      num_mixtures = wm_output.logpi.shape[2]
      mixture_idx = (
          Categorical(torch.exp(wm_output.logpi.view(num_mixtures)))
          .sample()
          .long()
          .item()
      )
      next_state = Normal(
          wm_output.mus[0, 0, mixture_idx],
          wm_output.sigmas[0, 0, mixture_idx],
      ).sample()
      reward = wm_output.reward[0, 0]
      if self.terminal_effective:
          not_terminal_prob = torch.sigmoid(
              wm_output.not_terminal[0, 0]
          )
          not_terminal = (
              Bernoulli(not_terminal_prob).sample().long().item()
          )
      else:
          not_terminal_prob = 1.0
          not_terminal = 1
      return reward, next_state, not_terminal, not_terminal_prob
  
imports: import itertools
  import logging
  import random
  from typing import List, Optional, Tuple
  import numpy
  import scipy.stats
  import torch
  import torch.nn
  from reagent.core import types
  from reagent.core.parameters import CONTINUOUS_TRAINING_ACTION_RANGE
  from reagent.models.base import ModelBase
  from reagent.models.world_model import MemoryNetwork
  from reagent.training.utils import rescale_actions
  from torch.distributions.bernoulli import Bernoulli
  from torch.distributions.categorical import Categorical
  from torch.distributions.normal import Normal
label_desc: Use image node_id 5 for calling the CEMPlannerNetwork obj's underlying member method code with example usage: obj.sample_reward_next_state_terminal(state, action, mem_net) and returns: reward,  next_state,  not_terminal,  not_terminal_prob