mirror of
https://github.com/vale981/ray
synced 2025-03-06 10:31:39 -05:00
131 lines
4.4 KiB
Python
131 lines
4.4 KiB
Python
from gym.spaces import Box, Dict, Discrete, MultiDiscrete, Tuple
|
|
import unittest
|
|
|
|
import ray
|
|
from ray import tune
|
|
from ray.rllib.examples.env.random_env import RandomEnv
|
|
from ray.rllib.examples.env.stateless_cartpole import StatelessCartPole
|
|
from ray.rllib.models.catalog import ModelCatalog
|
|
from ray.rllib.models.tf.attention_net import GTrXLNet
|
|
from ray.rllib.utils.test_utils import framework_iterator
|
|
|
|
|
|
class TestAttentionNets(unittest.TestCase):
|
|
|
|
config = {
|
|
"env": StatelessCartPole,
|
|
"gamma": 0.99,
|
|
"num_envs_per_worker": 20,
|
|
"framework": "tf",
|
|
}
|
|
|
|
stop = {
|
|
"episode_reward_mean": 150.0,
|
|
"timesteps_total": 5000000,
|
|
}
|
|
|
|
@classmethod
|
|
def setUpClass(cls) -> None:
|
|
ray.init(num_cpus=5)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls) -> None:
|
|
ray.shutdown()
|
|
|
|
def test_attention_nets_w_prev_actions_and_prev_rewards(self):
|
|
"""Tests attention prev-a/r input insertions using complex actions."""
|
|
config = {
|
|
"env": RandomEnv,
|
|
"env_config": {
|
|
"config": {
|
|
"action_space": Dict({
|
|
"a": Box(-1.0, 1.0, ()),
|
|
"b": Box(-1.0, 1.0, (2, )),
|
|
"c": Tuple([
|
|
Discrete(2),
|
|
MultiDiscrete([2, 3]),
|
|
Box(-1.0, 1.0, (3, )),
|
|
]),
|
|
}),
|
|
},
|
|
},
|
|
# Need to set this to True to enable complex (prev.) actions
|
|
# as inputs to the attention net.
|
|
"_disable_action_flattening": True,
|
|
"model": {
|
|
"fcnet_hiddens": [10],
|
|
"use_attention": True,
|
|
"attention_dim": 16,
|
|
"attention_use_n_prev_actions": 3,
|
|
"attention_use_n_prev_rewards": 2,
|
|
},
|
|
"num_sgd_iter": 1,
|
|
"train_batch_size": 200,
|
|
"sgd_minibatch_size": 50,
|
|
"rollout_fragment_length": 100,
|
|
"num_workers": 1,
|
|
}
|
|
for _ in framework_iterator(config):
|
|
tune.run(
|
|
"PPO",
|
|
config=config,
|
|
stop={"training_iteration": 1},
|
|
verbose=1)
|
|
|
|
def test_ppo_attention_net_learning(self):
|
|
ModelCatalog.register_custom_model("attention_net", GTrXLNet)
|
|
config = dict(
|
|
self.config, **{
|
|
"num_workers": 0,
|
|
"entropy_coeff": 0.001,
|
|
"vf_loss_coeff": 1e-5,
|
|
"num_sgd_iter": 5,
|
|
"model": {
|
|
"custom_model": "attention_net",
|
|
"max_seq_len": 10,
|
|
"custom_model_config": {
|
|
"num_transformer_units": 1,
|
|
"attention_dim": 32,
|
|
"num_heads": 1,
|
|
"memory_inference": 5,
|
|
"memory_training": 5,
|
|
"head_dim": 32,
|
|
"position_wise_mlp_dim": 32,
|
|
},
|
|
},
|
|
})
|
|
tune.run("PPO", config=config, stop=self.stop, verbose=1)
|
|
|
|
# TODO: (sven) causes memory failures/timeouts on Travis.
|
|
# Re-enable this once we have fast attention in master branch.
|
|
def test_impala_attention_net_learning(self):
|
|
return
|
|
# ModelCatalog.register_custom_model("attention_net", GTrXLNet)
|
|
# config = dict(
|
|
# self.config, **{
|
|
# "num_workers": 4,
|
|
# "num_gpus": 0,
|
|
# "entropy_coeff": 0.01,
|
|
# "vf_loss_coeff": 0.001,
|
|
# "lr": 0.0008,
|
|
# "model": {
|
|
# "custom_model": "attention_net",
|
|
# "max_seq_len": 65,
|
|
# "custom_model_config": {
|
|
# "num_transformer_units": 1,
|
|
# "attention_dim": 64,
|
|
# "num_heads": 1,
|
|
# "memory_inference": 10,
|
|
# "memory_training": 10,
|
|
# "head_dim": 32,
|
|
# "position_wise_mlp_dim": 32,
|
|
# },
|
|
# },
|
|
# })
|
|
# tune.run("IMPALA", config=config, stop=self.stop, verbose=1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import pytest
|
|
import sys
|
|
sys.exit(pytest.main(["-v", __file__]))
|