mbrl/models/model.py [193:231]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    def sample(
        self,
        act: torch.Tensor,
        model_state: Dict[str, torch.Tensor],
        deterministic: bool = False,
        rng: Optional[torch.Generator] = None,
    ) -> Tuple[
        torch.Tensor,
        Optional[torch.Tensor],
        Optional[torch.Tensor],
        Optional[Dict[str, torch.Tensor]],
    ]:
        """Samples a simulated transition from the dynamics model.

        This method will be used by :class:`ModelEnv` to simulate a transition of the form.
            o_t+1, r_t+1, d_t+1, st = sample(at, s_t), where

            - a_t: action taken at time t.
            - s_t: model state at time t (as returned by :meth:`reset()` or :meth:`sample()`.
            - r_t: reward at time t.
            - d_t: terminal indicator at time t.

        If the model doesn't simulate rewards and/or terminal indicators, it can return
        ``None`` for those.

        Args:
            act (tensor): the action at.
            model_state (tensor): the model state st.
            deterministic (bool): if ``True``, the model returns a deterministic
                "sample" (e.g., the mean prediction). Defaults to ``False``.
            rng (`torch.Generator`, optional): an optional random number generator
                to use.

        Returns:
            (tuple): predicted observation, rewards, terminal indicator and model
                state dictionary. Everything but the observation is optional, and can
                be returned with value ``None``.
        """
        raise NotImplementedError(
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



mbrl/models/model.py [409:421]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    def sample(
        self,
        act: torch.Tensor,
        model_state: Dict[str, torch.Tensor],
        deterministic: bool = False,
        rng: Optional[torch.Generator] = None,
    ) -> Tuple[
        torch.Tensor,
        Optional[torch.Tensor],
        Optional[torch.Tensor],
        Optional[Dict[str, torch.Tensor]],
    ]:
        raise NotImplementedError(
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -