chore: import upstream snapshot with attribution

doc/source/rllib/doc_code/dreamerv3_inference.py

@@ -0,0 +1,67 @@
+import gymnasium as gym
+import numpy as np
+
+import tree  # pip install dm_tree
+
+from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config
+from ray.rllib.core.columns import Columns
+from ray.rllib.utils.framework import convert_to_tensor
+
+
+env_name = "CartPole-v1"
+# Use the vector env API.
+env = gym.make_vec(env_name, num_envs=1, vectorization_mode="sync")
+
+terminated = truncated = False
+# Reset the env.
+obs, _ = env.reset()
+# Every time, we start a new episode, we should set is_first to True for the upcoming
+# action inference.
+is_first = 1.0
+
+# Create the algorithm from a simple config.
+config = (
+    DreamerV3Config()
+    .environment("CartPole-v1")
+    .training(model_size="XS", training_ratio=1024)
+)
+algo = config.build()
+
+# Extract the actual RLModule from the local (Dreamer) EnvRunner.
+rl_module = algo.env_runner.module
+# Get initial states from RLModule (note that these are always B=1, so this matches
+# our num_envs=1; if you are using a vector env >1, you would have to repeat the
+# returned states `num_env` times to get the correct batch size):
+states = rl_module.get_initial_state()
+# Batch the states to B=1.
+states = tree.map_structure(lambda s: s.unsqueeze(0), states)
+
+while not terminated and not truncated:
+    # Use the RLModule for action computations directly.
+    # DreamerV3 expects this particular batch format:
+    # obs=[B, T, ...]
+    # prev. states=[B, ...]
+    # `is_first`=[B]
+    batch = {
+        # States is already batched (see above).
+        Columns.STATE_IN: states,
+        # `obs` is already batched (due to vector env), but needs time-rank.
+        Columns.OBS: convert_to_tensor(obs, framework="torch")[None],
+        # Set to True at beginning of episode.
+        "is_first": convert_to_tensor(is_first, "torch")[None],
+    }
+    outs = rl_module.forward_inference(batch)
+    # Alternatively, call `forward_exploration` in case you want stochastic, non-greedy
+    # actions.
+    # outs = rl_module.forward_exploration(batch)
+
+    # Extract actions (remove time-rank) from outs.
+    actions = outs[Columns.ACTIONS].numpy()[0]
+    # Extract states from out. States are returned as batched.
+    states = outs[Columns.STATE_OUT]
+
+    # Perform a step in the env. Note that actions are still batched, which
+    # is ok, because we have a vector env.
+    obs, reward, terminated, truncated, info = env.step(actions)
+    # Not at the beginning of the episode anymore.
+    is_first = 0.0