mirror of
https://github.com/vale981/ray
synced 2025-03-09 21:06:39 -04:00
271 lines
9.7 KiB
Python
271 lines
9.7 KiB
Python
from collections import deque
|
||
import gym
|
||
import gym_minigrid
|
||
import numpy as np
|
||
import sys
|
||
import unittest
|
||
|
||
import ray
|
||
from ray import tune
|
||
from ray.rllib.agents.callbacks import DefaultCallbacks
|
||
import ray.rllib.agents.ppo as ppo
|
||
from ray.rllib.utils.test_utils import check_learning_achieved, \
|
||
framework_iterator
|
||
from ray.rllib.utils.numpy import one_hot
|
||
from ray.tune import register_env
|
||
|
||
|
||
class MyCallBack(DefaultCallbacks):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.deltas = []
|
||
|
||
def on_postprocess_trajectory(self, *, worker, episode, agent_id,
|
||
policy_id, policies, postprocessed_batch,
|
||
original_batches, **kwargs):
|
||
pos = np.argmax(postprocessed_batch["obs"], -1)
|
||
x, y = pos % 10, pos // 10
|
||
self.deltas.extend((x**2 + y**2)**0.5)
|
||
|
||
def on_sample_end(self, *, worker, samples, **kwargs):
|
||
print("mean. distance from origin={}".format(np.mean(self.deltas)))
|
||
self.deltas = []
|
||
|
||
|
||
class OneHotWrapper(gym.core.ObservationWrapper):
|
||
def __init__(self, env, vector_index, framestack):
|
||
super().__init__(env)
|
||
self.framestack = framestack
|
||
# 49=7x7 field of vision; 11=object types; 6=colors; 3=state types.
|
||
# +4: Direction.
|
||
self.single_frame_dim = 49 * (11 + 6 + 3) + 4
|
||
self.init_x = None
|
||
self.init_y = None
|
||
self.x_positions = []
|
||
self.y_positions = []
|
||
self.x_y_delta_buffer = deque(maxlen=100)
|
||
self.vector_index = vector_index
|
||
self.frame_buffer = deque(maxlen=self.framestack)
|
||
for _ in range(self.framestack):
|
||
self.frame_buffer.append(np.zeros((self.single_frame_dim, )))
|
||
|
||
self.observation_space = gym.spaces.Box(
|
||
0.0,
|
||
1.0,
|
||
shape=(self.single_frame_dim * self.framestack, ),
|
||
dtype=np.float32)
|
||
|
||
def observation(self, obs):
|
||
# Debug output: max-x/y positions to watch exploration progress.
|
||
if self.step_count == 0:
|
||
for _ in range(self.framestack):
|
||
self.frame_buffer.append(np.zeros((self.single_frame_dim, )))
|
||
if self.vector_index == 0:
|
||
if self.x_positions:
|
||
max_diff = max(
|
||
np.sqrt((np.array(self.x_positions) - self.init_x)**2 +
|
||
(np.array(self.y_positions) - self.init_y)**2))
|
||
self.x_y_delta_buffer.append(max_diff)
|
||
print("100-average dist travelled={}".format(
|
||
np.mean(self.x_y_delta_buffer)))
|
||
self.x_positions = []
|
||
self.y_positions = []
|
||
self.init_x = self.agent_pos[0]
|
||
self.init_y = self.agent_pos[1]
|
||
|
||
# Are we carrying the key?
|
||
# if self.carrying is not None:
|
||
# print("Carrying KEY!!")
|
||
|
||
self.x_positions.append(self.agent_pos[0])
|
||
self.y_positions.append(self.agent_pos[1])
|
||
|
||
# One-hot the last dim into 11, 6, 3 one-hot vectors, then flatten.
|
||
objects = one_hot(obs[:, :, 0], depth=11)
|
||
colors = one_hot(obs[:, :, 1], depth=6)
|
||
states = one_hot(obs[:, :, 2], depth=3)
|
||
# Is the door we see open?
|
||
# for x in range(7):
|
||
# for y in range(7):
|
||
# if objects[x, y, 4] == 1.0 and states[x, y, 0] == 1.0:
|
||
# print("Door OPEN!!")
|
||
|
||
all_ = np.concatenate([objects, colors, states], -1)
|
||
all_flat = np.reshape(all_, (-1, ))
|
||
direction = one_hot(
|
||
np.array(self.agent_dir), depth=4).astype(np.float32)
|
||
single_frame = np.concatenate([all_flat, direction])
|
||
self.frame_buffer.append(single_frame)
|
||
return np.concatenate(self.frame_buffer)
|
||
|
||
|
||
def env_maker(config):
|
||
name = config.get("name", "MiniGrid-Empty-5x5-v0")
|
||
framestack = config.get("framestack", 4)
|
||
env = gym.make(name)
|
||
# Only use image portion of observation (discard goal and direction).
|
||
env = gym_minigrid.wrappers.ImgObsWrapper(env)
|
||
env = OneHotWrapper(
|
||
env,
|
||
config.vector_index if hasattr(config, "vector_index") else 0,
|
||
framestack=framestack)
|
||
return env
|
||
|
||
|
||
register_env("mini-grid", env_maker)
|
||
CONV_FILTERS = [[16, [11, 11], 3], [32, [9, 9], 3], [64, [5, 5], 3]]
|
||
|
||
|
||
class TestCuriosity(unittest.TestCase):
|
||
@classmethod
|
||
def setUpClass(cls):
|
||
ray.init(num_cpus=3)
|
||
|
||
@classmethod
|
||
def tearDownClass(cls):
|
||
ray.shutdown()
|
||
|
||
def test_curiosity_on_large_frozen_lake(self):
|
||
config = ppo.DEFAULT_CONFIG.copy()
|
||
# A very large frozen-lake that's hard for a random policy to solve
|
||
# due to 0.0 feedback.
|
||
config["env"] = "FrozenLake-v0"
|
||
config["env_config"] = {
|
||
"desc": [
|
||
"SFFFFFFFFF",
|
||
"FFFFFFFFFF",
|
||
"FFFFFFFFFF",
|
||
"FFFFFFFFFF",
|
||
"FFFFFFFFFF",
|
||
"FFFFFFFFFF",
|
||
"FFFFFFFFFF",
|
||
"FFFFFFFFFF",
|
||
"FFFFFFFFFF",
|
||
"FFFFFFFFFG",
|
||
],
|
||
"is_slippery": False
|
||
}
|
||
# Print out observations to see how far we already get inside the Env.
|
||
config["callbacks"] = MyCallBack
|
||
# Limit horizon to make it really hard for non-curious agent to reach
|
||
# the goal state.
|
||
config["horizon"] = 23
|
||
# Local only.
|
||
config["num_workers"] = 0
|
||
config["lr"] = 0.001
|
||
|
||
num_iterations = 10
|
||
for _ in framework_iterator(config, frameworks="torch"):
|
||
# W/ Curiosity. Expect to learn something.
|
||
config["exploration_config"] = {
|
||
"type": "Curiosity",
|
||
"eta": 0.2,
|
||
"lr": 0.001,
|
||
"feature_dim": 128,
|
||
"feature_net_config": {
|
||
"fcnet_hiddens": [],
|
||
"fcnet_activation": "relu",
|
||
},
|
||
"sub_exploration": {
|
||
"type": "StochasticSampling",
|
||
}
|
||
}
|
||
trainer = ppo.PPOTrainer(config=config)
|
||
learnt = False
|
||
for i in range(num_iterations):
|
||
result = trainer.train()
|
||
print(result)
|
||
if result["episode_reward_max"] > 0.0:
|
||
print("Reached goal after {} iters!".format(i))
|
||
learnt = True
|
||
break
|
||
trainer.stop()
|
||
self.assertTrue(learnt)
|
||
|
||
# W/o Curiosity. Expect to learn nothing.
|
||
config["exploration_config"] = {
|
||
"type": "StochasticSampling",
|
||
}
|
||
trainer = ppo.PPOTrainer(config=config)
|
||
rewards_wo = 0.0
|
||
for _ in range(num_iterations):
|
||
result = trainer.train()
|
||
rewards_wo += result["episode_reward_mean"]
|
||
print(result)
|
||
trainer.stop()
|
||
self.assertTrue(rewards_wo == 0.0)
|
||
|
||
def test_curiosity_on_partially_observable_domain(self):
|
||
config = ppo.DEFAULT_CONFIG.copy()
|
||
config["env"] = "mini-grid"
|
||
config["env_config"] = {
|
||
# Also works with:
|
||
# - MiniGrid-MultiRoom-N4-S5-v0
|
||
# - MiniGrid-MultiRoom-N2-S4-v0
|
||
"name": "MiniGrid-Empty-8x8-v0",
|
||
"framestack": 1, # seems to work even w/o framestacking
|
||
}
|
||
config["horizon"] = 15 # Make it impossible to reach goal by chance.
|
||
config["num_envs_per_worker"] = 4
|
||
config["model"]["fcnet_hiddens"] = [256, 256]
|
||
config["model"]["fcnet_activation"] = "relu"
|
||
config["num_sgd_iter"] = 8
|
||
config["num_workers"] = 0
|
||
|
||
config["exploration_config"] = {
|
||
"type": "Curiosity",
|
||
# For the feature NN, use a non-LSTM fcnet (same as the one
|
||
# in the policy model).
|
||
"eta": 0.1,
|
||
"lr": 0.0003, # 0.0003 or 0.0005 seem to work fine as well.
|
||
"feature_dim": 64,
|
||
# No actual feature net: map directly from observations to feature
|
||
# vector (linearly).
|
||
"feature_net_config": {
|
||
"fcnet_hiddens": [],
|
||
"fcnet_activation": "relu",
|
||
},
|
||
"sub_exploration": {
|
||
"type": "StochasticSampling",
|
||
}
|
||
}
|
||
|
||
min_reward = 0.001
|
||
stop = {
|
||
"training_iteration": 25,
|
||
"episode_reward_mean": min_reward,
|
||
}
|
||
for _ in framework_iterator(config, frameworks="torch"):
|
||
# To replay:
|
||
# trainer = ppo.PPOTrainer(config=config)
|
||
# trainer.restore("[checkpoint file]")
|
||
# env = env_maker(config["env_config"])
|
||
# s = env.reset()
|
||
# for _ in range(10000):
|
||
# s, r, d, _ = env.step(trainer.compute_action(s))
|
||
# if d:
|
||
# s = env.reset()
|
||
# env.render()
|
||
|
||
results = tune.run("PPO", config=config, stop=stop, verbose=1)
|
||
check_learning_achieved(results, min_reward)
|
||
iters = results.trials[0].last_result["training_iteration"]
|
||
print("Reached in {} iterations.".format(iters))
|
||
|
||
# config_wo = config.copy()
|
||
# config_wo["exploration_config"] = {"type": "StochasticSampling"}
|
||
# stop_wo = stop.copy()
|
||
# stop_wo["training_iteration"] = iters
|
||
# results = tune.run(
|
||
# "PPO", config=config_wo, stop=stop_wo, verbose=1)
|
||
# try:
|
||
# check_learning_achieved(results, min_reward)
|
||
# except ValueError:
|
||
# print("Did not learn w/o curiosity (expected).")
|
||
# else:
|
||
# raise ValueError("Learnt w/o curiosity (not expected)!")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import pytest
|
||
sys.exit(pytest.main(["-v", __file__]))
|