class_docstr: A MuZero/Muesli-style loss on the rollouts of the dynamics model.
class_docstr_tok: ["A","MuZero\/Muesli-style","loss","on","the","rollouts","of","the","dynamics","model","."]
docstr: A convenient tree-mapped and configured call to rolling window.
  
  Stacks R = T - K + 1 action slices of length K = model_rollout_length from
  tensor x: [..., 0:K; ...; T-K:T, ...].
  
  Args:
    x: The tensor to select rolling slices from (along specified axis), with
      shape [..., T, ...] such that T = x.shape[axis].
    axis: The axis to slice from (defaults to 0).
  
  Returns:
    A tensor containing the stacked slices [0:K, ... T-K:T] from an axis of x
    with shape [..., K, R, ...] for input shape [..., T, ...].
docstr_tok: ["A","convenient","tree-mapped","and","configured","call","to","rolling","window",".","Stacks","R","=","T","-","K","+","1","action","slices","of","length","K","=","model_rollout_length","from","tensor","x",":","[","...",",","0",":","K",";","...",";","T-K",":","T",",","...","]",".","Args",":","x",":","The","tensor","to","select","rolling","slices","from","(","along","specified","axis",")",",","with","shape","[","...",",","T",",","...","]","such","that","T","=","x.shape","[","axis","]",".","axis",":","The","axis","to","slice","from","(","defaults","to","0",")",".","Returns",":","A","tensor","containing","the","stacked","slices","[","0",":","K",",","...","T-K",":","T","]","from","an","axis","of","x","with","shape","[","...",",","K",",","R",",","...","]","for","input","shape","[","...",",","T",",","...","]","."]
code: def _rolling_window(self, x: chex.Array, axis: int = 0) -> chex.Array:
      """A convenient tree-mapped and configured call to rolling window.
  
      Stacks R = T - K + 1 action slices of length K = model_rollout_length from
      tensor x: [..., 0:K; ...; T-K:T, ...].
  
      Args:
        x: The tensor to select rolling slices from (along specified axis), with
          shape [..., T, ...] such that T = x.shape[axis].
        axis: The axis to slice from (defaults to 0).
  
      Returns:
        A tensor containing the stacked slices [0:K, ... T-K:T] from an axis of x
        with shape [..., K, R, ...] for input shape [..., T, ...].
      """
  
      def rw(y):
          return mpo_utils.rolling_window(
              y,
              window=self._model_rollout_length,
              axis=axis,
              time_major=True,
          )
  
      return mpo_utils.tree_map_distribution(rw, x)
  
file_path: acme/acme/agents/jax/mpo/rollout_loss.py
imports: from typing import Tuple
  from acme import types
  from acme.adders import reverb
  from acme.agents.jax.mpo import categorical_mpo
  from acme.agents.jax.mpo import networks
  from acme.agents.jax.mpo import types
  from acme.agents.jax.mpo import utils
  from acme.jax import networks
  import chex
  import jax
  import jax.numpy
  import rlax
label_desc: Use image node_id 2 for calling the RolloutLoss obj's underlying member method code with example usage: obj._rolling_window(x, axis) and returns: mpo_utils, mpo_utils