wehub-resource-sync
68 lines
2.4 KiB
Python
68 lines
2.4 KiB
Python
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
|