mirror of
https://github.com/saymrwulf/stable-baselines3.git
synced 2026-05-18 21:30:19 +00:00
5d47296b8d
* Fix return computation + add test for GAE * Rename `last_dones` to `episode_starts` for clarification * Revert advantage * Cleanup test * Rename variable * Clarify return computation * Clarify docs * Add multi-episode rollout test * Reformat Co-authored-by: Anssi "Miffyli" Kanervisto <kaneran21@hotmail.com>
114 lines
3.9 KiB
Python
114 lines
3.9 KiB
Python
import gym
|
|
import numpy as np
|
|
import pytest
|
|
import torch as th
|
|
|
|
from stable_baselines3 import A2C, PPO
|
|
from stable_baselines3.common.callbacks import BaseCallback
|
|
from stable_baselines3.common.policies import ActorCriticPolicy
|
|
|
|
|
|
class CustomEnv(gym.Env):
|
|
def __init__(self, max_steps=8):
|
|
super(CustomEnv, self).__init__()
|
|
self.observation_space = gym.spaces.Box(low=-1, high=1, shape=(2,), dtype=np.float32)
|
|
self.action_space = gym.spaces.Box(low=-1, high=1, shape=(2,), dtype=np.float32)
|
|
self.max_steps = max_steps
|
|
self.n_steps = 0
|
|
|
|
def seed(self, seed):
|
|
self.observation_space.seed(seed)
|
|
|
|
def reset(self):
|
|
self.n_steps = 0
|
|
return self.observation_space.sample()
|
|
|
|
def step(self, action):
|
|
self.n_steps += 1
|
|
|
|
done = False
|
|
reward = 0.0
|
|
if self.n_steps >= self.max_steps:
|
|
reward = 1.0
|
|
done = True
|
|
|
|
return self.observation_space.sample(), reward, done, {}
|
|
|
|
|
|
class CheckGAECallback(BaseCallback):
|
|
def __init__(self):
|
|
super(CheckGAECallback, self).__init__(verbose=0)
|
|
|
|
def _on_rollout_end(self):
|
|
buffer = self.model.rollout_buffer
|
|
rollout_size = buffer.size()
|
|
|
|
max_steps = self.training_env.envs[0].max_steps
|
|
gamma = self.model.gamma
|
|
gae_lambda = self.model.gae_lambda
|
|
value = self.model.policy.constant_value
|
|
# We know in advance that the agent will get a single
|
|
# reward at the very last timestep of the episode,
|
|
# so we can pre-compute the lambda-return and advantage
|
|
deltas = np.zeros((rollout_size,))
|
|
advantages = np.zeros((rollout_size,))
|
|
# Reward should be 1.0 on final timestep of episode
|
|
rewards = np.zeros((rollout_size,))
|
|
rewards[max_steps - 1 :: max_steps] = 1.0
|
|
# Note that these are episode starts (+1 timestep from done)
|
|
episode_starts = np.zeros((rollout_size,))
|
|
episode_starts[::max_steps] = 1.0
|
|
|
|
# Final step is always terminal (next would episode_start = 1)
|
|
deltas[-1] = rewards[-1] - value
|
|
advantages[-1] = deltas[-1]
|
|
for n in reversed(range(rollout_size - 1)):
|
|
# Values are constants
|
|
episode_start_mask = 1.0 - episode_starts[n + 1]
|
|
deltas[n] = rewards[n] + gamma * value * episode_start_mask - value
|
|
advantages[n] = deltas[n] + gamma * gae_lambda * advantages[n + 1] * episode_start_mask
|
|
|
|
# TD(lambda) estimate, see Github PR #375
|
|
lambda_returns = advantages + value
|
|
|
|
assert np.allclose(buffer.advantages.flatten(), advantages)
|
|
assert np.allclose(buffer.returns.flatten(), lambda_returns)
|
|
|
|
def _on_step(self):
|
|
return True
|
|
|
|
|
|
class CustomPolicy(ActorCriticPolicy):
|
|
"""Custom Policy with a constant value function"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super(CustomPolicy, self).__init__(*args, **kwargs)
|
|
self.constant_value = 0.0
|
|
|
|
def forward(self, obs, deterministic=False):
|
|
actions, values, log_prob = super().forward(obs, deterministic)
|
|
# Overwrite values with ones
|
|
values = th.ones_like(values) * self.constant_value
|
|
return actions, values, log_prob
|
|
|
|
|
|
@pytest.mark.parametrize("model_class", [A2C, PPO])
|
|
@pytest.mark.parametrize("gae_lambda", [1.0, 0.9])
|
|
@pytest.mark.parametrize("gamma", [1.0, 0.99])
|
|
@pytest.mark.parametrize("num_episodes", [1, 3])
|
|
def test_gae_computation(model_class, gae_lambda, gamma, num_episodes):
|
|
env = CustomEnv(max_steps=64)
|
|
rollout_size = 64 * num_episodes
|
|
model = model_class(
|
|
CustomPolicy,
|
|
env,
|
|
seed=1,
|
|
gamma=gamma,
|
|
n_steps=rollout_size,
|
|
gae_lambda=gae_lambda,
|
|
)
|
|
model.learn(rollout_size, callback=CheckGAECallback())
|
|
|
|
# Change constant value so advantage != returns
|
|
model.policy.constant_value = 1.0
|
|
model.learn(rollout_size, callback=CheckGAECallback())
|