def get_exploration_strategy(variant, env):
from rlkit.exploration_strategies.epsilon_greedy import EpsilonGreedy
from rlkit.exploration_strategies.gaussian_strategy import GaussianStrategy
from rlkit.exploration_strategies.ou_strategy import OUStrategy
exploration_type = variant['exploration_type']
exploration_noise = variant.get('exploration_noise', 0.1)
if exploration_type == 'ou':
es = OUStrategy(
action_space=env.action_space,
max_sigma=exploration_noise,
min_sigma=exploration_noise, # Constant sigma
)
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
max_sigma=exploration_noise,
min_sigma=exploration_noise, # Constant sigma
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
prob_random_action=exploration_noise,
)
else:
raise Exception("Invalid type: " + exploration_type)
return es
def experiment(variant):
num_agent = variant['num_agent']
from cartpole import CartPoleEnv
expl_env = CartPoleEnv(mode=4)
eval_env = CartPoleEnv(mode=4)
obs_dim = eval_env.observation_space.low.size
action_dim = eval_env.action_space.n
policy_n, qf1_n, target_qf1_n, qf2_n, target_qf2_n, eval_policy_n, expl_policy_n = \
[], [], [], [], [], [], []
for i in range(num_agent):
policy = SoftmaxMlpPolicy(input_size=obs_dim,
output_size=action_dim,
**variant['policy_kwargs'])
qf1 = FlattenMlp(input_size=(obs_dim * num_agent + action_dim *
(num_agent - 1)),
output_size=action_dim,
**variant['qf_kwargs'])
target_qf1 = copy.deepcopy(qf1)
qf2 = FlattenMlp(input_size=(obs_dim * num_agent + action_dim *
(num_agent - 1)),
output_size=action_dim,
**variant['qf_kwargs'])
target_qf2 = copy.deepcopy(qf1)
eval_policy = ArgmaxDiscretePolicy(policy)
expl_policy = PolicyWrappedWithExplorationStrategy(
EpsilonGreedy(expl_env.action_space),
eval_policy,
)
policy_n.append(policy)
qf1_n.append(qf1)
target_qf1_n.append(target_qf1)
qf2_n.append(qf2)
target_qf2_n.append(target_qf2)
eval_policy_n.append(eval_policy)
expl_policy_n.append(expl_policy)
eval_path_collector = MAMdpPathCollector(eval_env, eval_policy_n)
expl_path_collector = MAMdpPathCollector(expl_env, expl_policy_n)
replay_buffer = MAEnvReplayBuffer(variant['replay_buffer_size'],
expl_env,
num_agent=num_agent)
trainer = MASACDiscreteTrainer(env=expl_env,
qf1_n=qf1_n,
target_qf1_n=target_qf1_n,
qf2_n=qf2_n,
target_qf2_n=target_qf2_n,
policy_n=policy_n,
**variant['trainer_kwargs'])
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
log_path_function=get_generic_ma_path_information,
**variant['algorithm_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
expl_env = gym.make("CartPole-v0")
eval_env = gym.make("CartPole-v0")
obs_dim = expl_env.observation_space.low.size
action_dim = eval_env.action_space.n
qf = Mlp(hidden_sizes=[32, 32], input_size=obs_dim, output_size=action_dim)
target_qf = Mlp(hidden_sizes=[32, 32],
input_size=obs_dim,
output_size=action_dim)
qf_criterion = nn.MSELoss()
eval_policy = ArgmaxDiscretePolicy(qf)
expl_policy = PolicyWrappedWithExplorationStrategy(
EpsilonGreedy(expl_env.action_space), eval_policy)
eval_path_collector = MdpPathCollector(eval_env, eval_policy)
expl_path_collector = MdpPathCollector(expl_env, expl_policy)
trainer = DQNTrainer(qf=qf,
target_qf=target_qf,
qf_criterion=qf_criterion,
**variant["trainer_kwargs"])
replay_buffer = EnvReplayBuffer(variant["replay_buffer_size"], expl_env)
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
**variant["algorithm_kwargs"])
algorithm.to(ptu.device)
algorithm.train()
def tdm_td3_experiment(variant):
env = variant['env_class'](**variant['env_kwargs'])
tdm_normalizer = None
qf1 = TdmQf(
env=env,
vectorized=True,
tdm_normalizer=tdm_normalizer,
**variant['qf_kwargs']
)
qf2 = TdmQf(
env=env,
vectorized=True,
tdm_normalizer=tdm_normalizer,
**variant['qf_kwargs']
)
policy = TdmPolicy(
env=env,
tdm_normalizer=tdm_normalizer,
**variant['policy_kwargs']
)
exploration_type = variant['exploration_type']
if exploration_type == 'ou':
es = OUStrategy(action_space=env.action_space)
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
max_sigma=0.1,
min_sigma=0.1, # Constant sigma
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
prob_random_action=0.1,
)
else:
raise Exception("Invalid type: " + exploration_type)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
replay_buffer = variant['replay_buffer_class'](
env=env,
**variant['replay_buffer_kwargs']
)
qf_criterion = variant['qf_criterion_class']()
algo_kwargs = variant['algo_kwargs']
algo_kwargs['td3_kwargs']['qf_criterion'] = qf_criterion
algo_kwargs['tdm_kwargs']['tdm_normalizer'] = tdm_normalizer
algorithm = TdmTd3(
env,
qf1=qf1,
qf2=qf2,
replay_buffer=replay_buffer,
policy=policy,
exploration_policy=exploration_policy,
**algo_kwargs
)
algorithm.to(ptu.device)
algorithm.train()
def create_exploration_policy(
env,
policy,
exploration_version='identity',
repeat_prob=0.,
prob_random_action=0.,
exploration_noise=0.,
):
# TODO: merge with get_exploration_strategy
if exploration_version == 'identity':
return policy
elif exploration_version == 'occasionally_repeat':
return ActionRepeatPolicy(policy, repeat_prob=repeat_prob)
elif exploration_version == 'epsilon_greedy':
return PolicyWrappedWithExplorationStrategy(
exploration_strategy=EpsilonGreedy(
action_space=env.action_space,
prob_random_action=prob_random_action,
),
policy=policy)
elif exploration_version == 'epsilon_greedy_and_occasionally_repeat':
policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=EpsilonGreedy(
action_space=env.action_space,
prob_random_action=prob_random_action,
),
policy=policy)
return ActionRepeatPolicy(policy, repeat_prob=repeat_prob)
elif exploration_version == 'epsilon_greedy':
return PolicyWrappedWithExplorationStrategy(
exploration_strategy=EpsilonGreedy(
action_space=env.action_space,
prob_random_action=prob_random_action,
),
policy=policy)
elif exploration_version == 'ou':
return PolicyWrappedWithExplorationStrategy(
exploration_strategy=OUStrategy(
action_space=env.action_space,
max_sigma=exploration_noise,
min_sigma=exploration_noise,
),
policy=policy)
else:
raise ValueError(exploration_version)
def experiment(variant):
import sys
from traffic.make_env import make_env
expl_env = make_env(args.exp_name)
eval_env = make_env(args.exp_name)
obs_dim = eval_env.observation_space.low.size
action_dim = eval_env.action_space.n
gb = TrafficGraphBuilder(input_dim=4,
ego_init=torch.tensor([0.,1.]),
other_init=torch.tensor([1.,0.]),
edge_index=torch.tensor([[0,0,1,2],
[1,2,0,0]]))
qf = GNNNet(
pre_graph_builder = gb,
node_dim = 16,
output_dim = action_dim,
post_mlp_kwargs = variant['qf_kwargs'],
num_conv_layers=3)
target_qf = copy.deepcopy(qf)
eval_policy = ArgmaxDiscretePolicy(qf)
expl_policy = PolicyWrappedWithExplorationStrategy(
EpsilonGreedy(expl_env.action_space, variant['epsilon']),
eval_policy,
)
eval_path_collector = MdpPathCollector(
eval_env,
eval_policy,
)
expl_path_collector = MdpPathCollector(
expl_env,
expl_policy,
)
replay_buffer = PrioritizedReplayBuffer(
variant['replay_buffer_size'],
expl_env,
)
qf_criterion = nn.MSELoss()
trainer = DQNTrainer(
qf=qf,
target_qf=target_qf,
qf_criterion=qf_criterion,
replay_buffer=replay_buffer,
**variant['trainer_kwargs']
)
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
**variant['algorithm_kwargs']
)
algorithm.to(ptu.device)
algorithm.train()
def get_validation_returns(self, snapshot):
policy = snapshot['evaluation/policy']
policy = PolicyWrappedWithExplorationStrategy(
EpsilonGreedy(self.eval_env.action_space, 0.1), policy)
validation_envs = pickle.load(open(self.validation_envs_pkl, 'rb'))
returns = np.zeros(len(validation_envs['envs']))
for env_idx, env in enumerate(validation_envs['envs']):
path = rollout(env, policy, self.validation_rollout_length)
returns[env_idx] = path['rewards'].sum()
return {'returns': returns.mean()}
def experiment(variant):
if variant['multitask']:
env = CylinderXYPusher2DEnv(**variant['env_kwargs'])
env = MultitaskToFlatEnv(env)
else:
env = Pusher2DEnv(**variant['env_kwargs'])
if variant['normalize']:
env = NormalizedBoxEnv(env)
exploration_type = variant['exploration_type']
if exploration_type == 'ou':
es = OUStrategy(action_space=env.action_space)
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
max_sigma=0.1,
min_sigma=0.1, # Constant sigma
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
prob_random_action=0.1,
)
else:
raise Exception("Invalid type: " + exploration_type)
obs_dim = env.observation_space.low.size
action_dim = env.action_space.low.size
qf1 = ConcatMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[400, 300],
)
qf2 = ConcatMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[400, 300],
)
policy = TanhMlpPolicy(
input_size=obs_dim,
output_size=action_dim,
hidden_sizes=[400, 300],
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
algorithm = TD3(
env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
**variant['algo_kwargs']
)
algorithm.to(ptu.device)
algorithm.train()
def td3_experiment(variant):
env = variant['env_class'](**variant['env_kwargs'])
env = MultitaskToFlatEnv(env)
if variant.get('make_silent_env', True):
env = MultitaskEnvToSilentMultitaskEnv(env)
if variant['normalize']:
env = NormalizedBoxEnv(env)
exploration_type = variant['exploration_type']
if exploration_type == 'ou':
es = OUStrategy(action_space=env.action_space)
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
max_sigma=0.1,
min_sigma=0.1, # Constant sigma
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
prob_random_action=0.1,
)
else:
raise Exception("Invalid type: " + exploration_type)
obs_dim = env.observation_space.low.size
action_dim = env.action_space.low.size
qf1 = ConcatMlp(
input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs']
)
qf2 = ConcatMlp(
input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs']
)
policy = TanhMlpPolicy(
input_size=obs_dim,
output_size=action_dim,
**variant['policy_kwargs']
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
algorithm = TD3(
env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
**variant['algo_kwargs']
)
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
expl_env = gym.make('GoalGridworld-v0')
eval_env = gym.make('GoalGridworld-v0')
obs_dim = expl_env.observation_space.spaces['observation'].low.size
goal_dim = expl_env.observation_space.spaces['desired_goal'].low.size
action_dim = expl_env.action_space.n
qf = FlattenMlp(
input_size=obs_dim + goal_dim,
output_size=action_dim,
hidden_sizes=[400, 300],
)
target_qf = FlattenMlp(
input_size=obs_dim + goal_dim,
output_size=action_dim,
hidden_sizes=[400, 300],
)
eval_policy = ArgmaxDiscretePolicy(qf)
exploration_strategy = EpsilonGreedy(action_space=expl_env.action_space, )
expl_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=exploration_strategy,
policy=eval_policy,
)
replay_buffer = ObsDictRelabelingBuffer(env=eval_env,
**variant['replay_buffer_kwargs'])
observation_key = 'observation'
desired_goal_key = 'desired_goal'
eval_path_collector = GoalConditionedPathCollector(
eval_env,
eval_policy,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
expl_path_collector = GoalConditionedPathCollector(
expl_env,
expl_policy,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
trainer = DQNTrainer(qf=qf,
target_qf=target_qf,
**variant['trainer_kwargs'])
trainer = HERTrainer(trainer)
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
# Select a different success_function for different tasks.
expl_env = GymCraftingEnv(state_obs=True,
few_obj=True,
success_function=eval_eatbread)
eval_env = GymCraftingEnv(state_obs=True,
few_obj=True,
success_function=eval_eatbread)
obs_dim = expl_env.observation_space.low.size
action_dim = eval_env.action_space.n
qf = Mlp(
hidden_sizes=[32, 32],
input_size=obs_dim,
output_size=action_dim,
)
target_qf = Mlp(
hidden_sizes=[32, 32],
input_size=obs_dim,
output_size=action_dim,
)
qf_criterion = nn.MSELoss()
eval_policy = ArgmaxDiscretePolicy(qf)
expl_policy = PolicyWrappedWithExplorationStrategy(
EpsilonGreedy(expl_env.action_space),
eval_policy,
)
eval_path_collector = MdpPathCollector(
eval_env,
eval_policy,
)
expl_path_collector = MdpPathCollector(
expl_env,
expl_policy,
)
trainer = DQNTrainer(qf=qf,
target_qf=target_qf,
qf_criterion=qf_criterion,
**variant['trainer_kwargs'])
replay_buffer = EnvReplayBuffer(
variant['replay_buffer_size'],
expl_env,
)
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
**variant['algorithm_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
env = variant['env_class'](**variant['env_kwargs'])
if variant['normalize']:
env = NormalizedBoxEnv(env)
exploration_type = variant['exploration_type']
if exploration_type == 'ou':
es = OUStrategy(action_space=env.action_space)
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
max_sigma=0.1,
min_sigma=0.1, # Constant sigma
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
prob_random_action=0.1,
)
else:
raise Exception("Invalid type: " + exploration_type)
obs_dim = env.observation_space.low.size
action_dim = env.action_space.low.size
goal_dim = env.goal_dim
qf1 = ConcatMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
hidden_sizes=[400, 300],
)
qf2 = ConcatMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
hidden_sizes=[400, 300],
)
policy = TanhMlpPolicy(
input_size=obs_dim + goal_dim,
output_size=action_dim,
hidden_sizes=[400, 300],
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
replay_buffer = variant['replay_buffer_class'](
env=env, **variant['replay_buffer_kwargs'])
algorithm = HerTd3(env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
replay_buffer=replay_buffer,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
args = getArgs()
# expl_env = NormalizedBoxEnv(environment(args))
expl_env = environment(args, 'dqn')
eval_env = environment(args, 'dqn')
# expl_env.render()
obs_dim = expl_env.get_obsdim()
action_dim = expl_env.action_space.n
qf = Mlp(
hidden_sizes=[32, 32],
input_size=obs_dim,
output_size=action_dim,
)
target_qf = Mlp(
hidden_sizes=[32, 32],
input_size=obs_dim,
output_size=action_dim,
)
qf_criterion = nn.MSELoss()
eval_policy = ArgmaxDiscretePolicy(qf)
expl_policy = PolicyWrappedWithExplorationStrategy(
EpsilonGreedy(expl_env.action_space),
eval_policy,
)
eval_path_collector = MdpPathCollector(
eval_env,
eval_policy,
)
expl_path_collector = MdpPathCollector(
expl_env,
expl_policy,
)
trainer = DQNTrainer(qf=qf,
target_qf=target_qf,
qf_criterion=qf_criterion,
**variant['trainer_kwargs'])
replay_buffer = EnvReplayBuffer(
variant['replay_buffer_size'],
expl_env,
)
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
**variant['algorithm_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def __init__(self,
env_sampler,
qf,
policy=None,
learning_rate=1e-3,
use_hard_updates=False,
hard_update_period=1000,
tau=0.001,
epsilon=0.1,
qf_criterion=None,
**kwargs):
"""
:param env: Env.
:param qf: QFunction. Maps from state to action Q-values.
:param learning_rate: Learning rate for qf. Adam is used.
:param use_hard_updates: Use a hard rather than soft update.
:param hard_update_period: How many gradient steps before copying the
parameters over. Used if `use_hard_updates` is True.
:param tau: Soft target tau to update target QF. Used if
`use_hard_updates` is False.
:param epsilon: Probability of taking a random action.
:param kwargs: kwargs to pass onto TorchRLAlgorithm
"""
self.env_sampler = env_sampler
env, _ = env_sampler()
exploration_strategy = EpsilonGreedy(
action_space=env.action_space,
prob_random_action=epsilon,
)
self.policy = policy or ArgmaxDiscretePolicy(qf)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=exploration_strategy,
policy=self.policy,
)
super().__init__(env_sampler,
exploration_policy,
eval_policy=self.policy,
**kwargs)
self.qf = qf
self.target_qf = self.qf.copy()
self.learning_rate = learning_rate
self.use_hard_updates = use_hard_updates
self.hard_update_period = hard_update_period
self.tau = tau
self.qf_optimizer = optim.Adam(
self.qf.parameters(),
lr=self.learning_rate,
)
self.qf_criterion = qf_criterion or nn.MSELoss()
self.eval_statistics = None
def experiment(variant):
import sys
from traffic.make_env import make_env
expl_env = make_env(args.exp_name)
eval_env = make_env(args.exp_name)
obs_dim = eval_env.observation_space.low.size
action_dim = eval_env.action_space.n
qf = Mlp(
input_size=obs_dim,
output_size=action_dim,
**variant['qf_kwargs']
)
target_qf = copy.deepcopy(qf)
eval_policy = ArgmaxDiscretePolicy(qf)
expl_policy = PolicyWrappedWithExplorationStrategy(
EpsilonGreedy(expl_env.action_space, variant['epsilon']),
eval_policy,
)
eval_path_collector = MdpPathCollector(
eval_env,
eval_policy,
)
expl_path_collector = MdpPathCollector(
expl_env,
expl_policy,
)
qf_criterion = nn.MSELoss()
trainer = DoubleDQNTrainer(
qf=qf,
target_qf=target_qf,
qf_criterion=qf_criterion,
**variant['trainer_kwargs']
)
replay_buffer = EnvReplayBuffer(
variant['replay_buffer_size'],
expl_env,
)
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
log_path_function = get_traffic_path_information,
**variant['algorithm_kwargs']
)
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
from cartpole import CartPoleEnv
expl_env = CartPoleEnv(mode=2)
eval_env = CartPoleEnv(mode=2)
obs_dim = eval_env.observation_space.low.size
action_dim = eval_env.action_space.n
qf = Mlp(input_size=obs_dim,
output_size=action_dim,
**variant['qf_kwargs'])
target_qf = copy.deepcopy(qf)
eval_policy = ArgmaxDiscretePolicy(qf)
expl_policy = PolicyWrappedWithExplorationStrategy(
EpsilonGreedy(expl_env.action_space, variant['epsilon']),
eval_policy,
)
eval_path_collector = MdpPathCollector(
eval_env,
eval_policy,
)
expl_path_collector = MdpPathCollector(
expl_env,
expl_policy,
)
replay_buffer = PrioritizedReplayBuffer(
variant['replay_buffer_size'],
expl_env,
)
qf_criterion = nn.MSELoss()
trainer = DQNTrainer(qf=qf,
target_qf=target_qf,
qf_criterion=qf_criterion,
replay_buffer=replay_buffer,
**variant['trainer_kwargs'])
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
**variant['algorithm_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
env = CylinderXYPusher2DEnv(**variant['env_kwargs'])
if variant['normalize']:
env = NormalizedBoxEnv(env)
es = EpsilonGreedy(
action_space=env.action_space,
prob_random_action=0.1,
)
obs_dim = env.observation_space.low.size
action_dim = env.action_space.low.size
goal_dim = env.goal_dim
qf1 = ConcatMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
hidden_sizes=[400, 300],
)
qf2 = ConcatMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
hidden_sizes=[400, 300],
)
policy = TanhMlpPolicy(
input_size=obs_dim + goal_dim,
output_size=action_dim,
hidden_sizes=[400, 300],
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
replay_buffer = SimpleHerReplayBuffer(env=env,
**variant['replay_buffer_kwargs'])
algorithm = HerTd3(env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
replay_buffer=replay_buffer,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def get_exploration_strategy(variant, env):
from rlkit.exploration_strategies.epsilon_greedy import EpsilonGreedy
from rlkit.exploration_strategies.gaussian_strategy import GaussianStrategy
from rlkit.exploration_strategies.gaussian_and_epislon import \
GaussianAndEpislonStrategy
from rlkit.exploration_strategies.ou_strategy import OUStrategy
from rlkit.exploration_strategies.noop import NoopStrategy
exploration_type = variant['exploration_type']
# exploration_noise = variant.get('exploration_noise', 0.1)
es_kwargs = variant.get('es_kwargs', {})
if exploration_type == 'ou':
es = OUStrategy(
action_space=env.action_space,
# max_sigma=exploration_noise,
# min_sigma=exploration_noise, # Constant sigma
**es_kwargs)
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
# max_sigma=exploration_noise,
# min_sigma=exploration_noise, # Constant sigma
**es_kwargs)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
# prob_random_action=exploration_noise,
**es_kwargs)
elif exploration_type == 'gaussian_and_epsilon':
es = GaussianAndEpislonStrategy(
action_space=env.action_space,
# max_sigma=exploration_noise,
# min_sigma=exploration_noise, # Constant sigma
# epsilon=exploration_noise,
**es_kwargs)
elif exploration_type == 'noop':
es = NoopStrategy(action_space=env.action_space)
else:
raise Exception("Invalid type: " + exploration_type)
return es
def validate(self, snapshot):
"""
Collect list of stats for each validation env as dict of following format:
'pickup_wood': [0, 15, 20] means you picked up a wood object at timesteps 0, 15, and 20.
"""
policy = snapshot['evaluation/policy']
if hasattr(policy, 'policy'):
# if it's reset free, strip out the underlying policy from the exploration strategy
policy = policy.policy
policy = PolicyWrappedWithExplorationStrategy(
EpsilonGreedy(self.eval_env.action_space, 0.1), policy)
validation_envs = pickle.load(open(self.validation_envs_pkl, 'rb'))
stats = [{} for _ in range(len(validation_envs['envs']))]
for env_idx, env in enumerate(validation_envs['envs']):
path = rollout(env, policy, self.validation_rollout_length)
for typ in env.object_to_idx.keys():
if typ not in ['empty', 'wall', 'tree']:
key = 'pickup_%s' % typ
last_val = 0
pickup_idxs = []
for t, env_info in enumerate(path['env_infos']):
count = env_info[key] - last_val
pickup_idxs.extend([t for _ in range(count)])
last_val = env_info[key]
stats[env_idx][key] = pickup_idxs
for typ in env.interactions.values():
key = 'made_%s' % typ
last_val = 0
made_idxs = []
for t, env_info in enumerate(path['env_infos']):
count = env_info[key] - last_val
made_idxs.extend([t for _ in range(count)])
last_val = env_info[key]
stats[env_idx][key] = made_idxs
return stats
def her_td3_experiment(variant):
if 'env_id' in variant:
env = gym.make(variant['env_id'])
else:
env = variant['env_class'](**variant['env_kwargs'])
observation_key = variant['observation_key']
desired_goal_key = variant['desired_goal_key']
variant['algo_kwargs']['her_kwargs']['observation_key'] = observation_key
variant['algo_kwargs']['her_kwargs']['desired_goal_key'] = desired_goal_key
if variant.get('normalize', False):
raise NotImplementedError()
achieved_goal_key = desired_goal_key.replace("desired", "achieved")
replay_buffer = ObsDictRelabelingBuffer(
env=env,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
achieved_goal_key=achieved_goal_key,
**variant['replay_buffer_kwargs'])
obs_dim = env.observation_space.spaces['observation'].low.size
action_dim = env.action_space.low.size
goal_dim = env.observation_space.spaces['desired_goal'].low.size
exploration_type = variant['exploration_type']
if exploration_type == 'ou':
es = OUStrategy(action_space=env.action_space, **variant['es_kwargs'])
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
**variant['es_kwargs'],
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
**variant['es_kwargs'],
)
else:
raise Exception("Invalid type: " + exploration_type)
qf1 = FlattenMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
qf2 = FlattenMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
policy = TanhMlpPolicy(input_size=obs_dim + goal_dim,
output_size=action_dim,
**variant['policy_kwargs'])
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
algorithm = HerTd3(env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
replay_buffer=replay_buffer,
**variant['algo_kwargs'])
if variant.get("save_video", False):
rollout_function = rf.create_rollout_function(
rf.multitask_rollout,
max_path_length=algorithm.max_path_length,
observation_key=algorithm.observation_key,
desired_goal_key=algorithm.desired_goal_key,
)
video_func = get_video_save_func(
rollout_function,
env,
policy,
variant,
)
algorithm.post_epoch_funcs.append(video_func)
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
num_agent = variant['num_agent']
from cartpole import CartPoleEnv
from rlkit.envs.ma_wrappers import MAProbDiscreteEnv
expl_env = CartPoleEnv(mode=4)
eval_env = CartPoleEnv(mode=4)
obs_dim = eval_env.observation_space.low.size
action_dim = eval_env.action_space.n
qf_n, cactor_n, policy_n, target_qf_n, target_cactor_n, target_policy_n, eval_policy_n, expl_policy_n = \
[], [], [], [], [], [], [], []
for i in range(num_agent):
qf = FlattenMlp(
input_size=(obs_dim*num_agent+action_dim*num_agent),
output_size=1,
**variant['qf_kwargs']
)
cactor = GumbelSoftmaxMlpPolicy(
input_size=(obs_dim*num_agent+action_dim*(num_agent-1)),
output_size=action_dim,
**variant['cactor_kwargs']
)
policy = GumbelSoftmaxMlpPolicy(
input_size=obs_dim,
output_size=action_dim,
**variant['policy_kwargs']
)
target_qf = copy.deepcopy(qf)
target_cactor = copy.deepcopy(cactor)
target_policy = copy.deepcopy(policy)
eval_policy = ArgmaxDiscretePolicy(policy,use_preactivation=True)
expl_policy = PolicyWrappedWithExplorationStrategy(
EpsilonGreedy(expl_env.action_space),
eval_policy,
)
qf_n.append(qf)
cactor_n.append(cactor)
policy_n.append(policy)
target_qf_n.append(target_qf)
target_cactor_n.append(target_cactor)
target_policy_n.append(target_policy)
eval_policy_n.append(eval_policy)
expl_policy_n.append(expl_policy)
eval_path_collector = MAMdpPathCollector(eval_env, eval_policy_n)
expl_path_collector = MAMdpPathCollector(expl_env, expl_policy_n)
replay_buffer = MAEnvReplayBuffer(variant['replay_buffer_size'], expl_env, num_agent=num_agent)
trainer = PRGTrainer(
env=expl_env,
qf_n=qf_n,
target_qf_n=target_qf_n,
policy_n=policy_n,
target_policy_n=target_policy_n,
cactor_n=cactor_n,
target_cactor_n=target_cactor_n,
**variant['trainer_kwargs']
)
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
log_path_function=get_generic_ma_path_information,
**variant['algorithm_kwargs']
)
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
import sys
sys.path.append("./multiagent-particle-envs")
from make_env import make_env
from particle_env_wrapper import ParticleEnv
expl_env = ParticleEnv(make_env(args.exp_name,discrete_action_space=False,world_args=variant['world_args']))
eval_env = ParticleEnv(make_env(args.exp_name,discrete_action_space=False,world_args=variant['world_args']))
num_agent = expl_env.num_agent
obs_dim = eval_env.observation_space.low.size
action_dim = eval_env.action_space.low.size
from simple_spread_graph import SimpleSpreadGraphBuilder
og_builder_1 = SimpleSpreadGraphBuilder(
num_agents=expl_env.scenario.num_agents,
num_landmarks=expl_env.scenario.num_landmarks,
batch_size=variant['algorithm_kwargs']['batch_size'],
append_action=False,
single_observe=False,
contain_self_loop=True,
)
from rlkit.torch.networks.gnn_networks import GNNNet
from rlkit.torch.networks.layers import SelectLayer
og1 = nn.Sequential(
GNNNet(
og_builder_1,
node_dim=variant['graph_kwargs']['node_dim'],
conv_type='GSage',
num_conv_layers=variant['graph_kwargs']['num_layer'],
hidden_activation='lrelu0.2',
output_activation='lrelu0.2',
),
SelectLayer(dim=1, index=torch.arange(num_agent)),
)
target_og1 = copy.deepcopy(og1)
from rlkit.torch.networks.graph_builders import FullGraphBuilder
cg_builder_1 = FullGraphBuilder(
input_node_dim=variant['graph_kwargs']['node_dim']+action_dim,
num_node=num_agent,
batch_size=variant['algorithm_kwargs']['batch_size'],
contain_self_loop=False)
from rlkit.torch.networks.graph_context_network import GraphContextNet
cg1 = GraphContextNet(
cg_builder_1,
variant['graph_kwargs']['node_dim'],
action_dim,
output_activation='lrelu0.2',
**variant['graph_kwargs']
)
target_cg1 = copy.deepcopy(cg1)
from rlkit.torch.networks.networks import FlattenMlp
qf1 = FlattenMlp(input_size=variant['graph_kwargs']['node_dim']+action_dim,
output_size=1,
hidden_sizes=[variant['qf_kwargs']['hidden_dim']]*(variant['qf_kwargs']['num_layer']-1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
)
target_qf1 = copy.deepcopy(qf1)
og_builder_2 = SimpleSpreadGraphBuilder(
num_agents=expl_env.scenario.num_agents,
num_landmarks=expl_env.scenario.num_landmarks,
batch_size=variant['algorithm_kwargs']['batch_size'],
append_action=False,
single_observe=False,
contain_self_loop=True,
)
from rlkit.torch.networks.gnn_networks import GNNNet
og2 = nn.Sequential(
GNNNet(
og_builder_2,
node_dim=variant['graph_kwargs']['node_dim'],
conv_type='GSage',
num_conv_layers=variant['graph_kwargs']['num_layer'],
hidden_activation='lrelu0.2',
output_activation='lrelu0.2',
),
SelectLayer(dim=1, index=torch.arange(num_agent)),
)
target_og2 = copy.deepcopy(og2)
cg_builder_2 = FullGraphBuilder(
input_node_dim=variant['graph_kwargs']['node_dim']+action_dim,
num_node=num_agent,
batch_size=variant['algorithm_kwargs']['batch_size'],
contain_self_loop=False)
cg2 = GraphContextNet(
cg_builder_2,
variant['graph_kwargs']['node_dim'],
action_dim,
output_activation='lrelu0.2',
**variant['graph_kwargs']
)
target_cg2 = copy.deepcopy(cg2)
qf2 = FlattenMlp(input_size=variant['graph_kwargs']['node_dim']+action_dim,
output_size=1,
hidden_sizes=[variant['qf_kwargs']['hidden_dim']]*(variant['qf_kwargs']['num_layer']-1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
)
target_qf2 = copy.deepcopy(qf2)
og_builder_ca = SimpleSpreadGraphBuilder(
num_agents=expl_env.scenario.num_agents,
num_landmarks=expl_env.scenario.num_landmarks,
batch_size=variant['algorithm_kwargs']['batch_size'],
append_action=False,
single_observe=False,
contain_self_loop=True,
)
from rlkit.torch.networks.gnn_networks import GNNNet
ogca = nn.Sequential(
GNNNet(
og_builder_ca,
node_dim=variant['graph_kwargs']['node_dim'],
conv_type='GSage',
num_conv_layers=variant['graph_kwargs']['num_layer'],
hidden_activation='lrelu0.2',
output_activation='lrelu0.2',
),
SelectLayer(dim=1, index=torch.arange(num_agent)),
)
cg_builder_ca = FullGraphBuilder(
input_node_dim=variant['graph_kwargs']['node_dim']+action_dim,
num_node=num_agent,
batch_size=variant['algorithm_kwargs']['batch_size'],
contain_self_loop=False)
cgca = GraphContextNet(
cg_builder_ca,
variant['graph_kwargs']['node_dim'],
action_dim,
output_activation='lrelu0.2',
**variant['graph_kwargs']
)
from rlkit.torch.networks.layers import SplitLayer
from rlkit.torch.policies.tanh_gaussian_policy import TanhGaussianPolicy
cactor = nn.Sequential(
FlattenMlp(input_size=variant['graph_kwargs']['node_dim'],
output_size=variant['cactor_kwargs']['hidden_dim'],
hidden_sizes=[variant['cactor_kwargs']['hidden_dim']]*(variant['cactor_kwargs']['num_layer']-1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
output_activation=nn.LeakyReLU(negative_slope=0.2),
),
nn.LeakyReLU(negative_slope=0.2),
SplitLayer(layers=[nn.Linear(variant['policy_kwargs']['hidden_dim'],action_dim),
nn.Linear(variant['policy_kwargs']['hidden_dim'],action_dim)])
)
cactor = TanhGaussianPolicy(module=cactor)
policy_n, expl_policy_n, eval_policy_n = [], [], []
for i in range(num_agent):
policy = nn.Sequential(
FlattenMlp(input_size=obs_dim,
output_size=variant['policy_kwargs']['hidden_dim'],
hidden_sizes=[variant['policy_kwargs']['hidden_dim']]*(variant['policy_kwargs']['num_layer']-1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
output_activation=nn.LeakyReLU(negative_slope=0.2),
),
SplitLayer(layers=[nn.Linear(variant['policy_kwargs']['hidden_dim'],action_dim),
nn.Linear(variant['policy_kwargs']['hidden_dim'],action_dim)])
)
policy = TanhGaussianPolicy(module=policy)
from rlkit.torch.policies.make_deterministic import MakeDeterministic
eval_policy = MakeDeterministic(policy)
if variant['random_exploration']:
from rlkit.exploration_strategies.base import PolicyWrappedWithExplorationStrategy
from rlkit.exploration_strategies.epsilon_greedy import EpsilonGreedy
expl_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=EpsilonGreedy(expl_env.action_space, prob_random_action=1.0),
policy=policy,
)
else:
expl_policy = policy
policy_n.append(policy)
expl_policy_n.append(expl_policy)
eval_policy_n.append(eval_policy)
from rlkit.samplers.data_collector.ma_path_collector import MAMdpPathCollector
eval_path_collector = MAMdpPathCollector(eval_env, eval_policy_n)
expl_path_collector = MAMdpPathCollector(expl_env, expl_policy_n)
from rlkit.data_management.ma_env_replay_buffer import MAEnvReplayBuffer
replay_buffer = MAEnvReplayBuffer(variant['replay_buffer_size'], expl_env, num_agent=num_agent)
from rlkit.torch.r2g.r2g_gnn8 import R2GGNNTrainer
trainer = R2GGNNTrainer(
env=expl_env,
og1=og1,
target_og1=target_og1,
cg1=cg1,
target_cg1=target_cg1,
qf1=qf1,
target_qf1=target_qf1,
og2=og2,
target_og2=target_og2,
cg2=cg2,
target_cg2=target_cg2,
qf2=qf2,
target_qf2=target_qf2,
ogca=ogca,
cgca=cgca,
cactor=cactor,
policy_n=policy_n,
**variant['trainer_kwargs']
)
from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
log_path_function=get_generic_ma_path_information,
**variant['algorithm_kwargs']
)
algorithm.to(ptu.device)
# save init params
from rlkit.core import logger
snapshot = algorithm._get_snapshot()
file_name = osp.join(logger._snapshot_dir, 'itr_-1.pkl')
torch.save(snapshot, file_name)
algorithm.train()
def experiment(variant):
from multi_differential_game import MultiDifferentialGame
expl_env = MultiDifferentialGame(**variant['env_kwargs'])
eval_env = MultiDifferentialGame(**variant['env_kwargs'])
num_agent = expl_env.agent_num
obs_dim = eval_env.observation_space.low.size
action_dim = eval_env.action_space.low.size
from rlkit.torch.networks.graph_builders import FullGraphBuilder
graph_builder_1 = FullGraphBuilder(
input_node_dim=obs_dim + action_dim,
num_node=num_agent,
batch_size=variant['algorithm_kwargs']['batch_size'],
contain_self_loop=False)
from rlkit.torch.networks.graph_context_network import GraphContextNet
cg1 = GraphContextNet(graph_builder_1,
obs_dim,
action_dim,
output_activation='lrelu0.2',
**variant['graph_kwargs'])
target_cg1 = copy.deepcopy(cg1)
from rlkit.torch.networks.networks import FlattenMlp
qf1 = FlattenMlp(
input_size=variant['graph_kwargs']['node_dim'] + action_dim,
output_size=1,
hidden_sizes=[variant['qf_kwargs']['hidden_dim']] *
(variant['qf_kwargs']['num_layer'] - 1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
)
target_qf1 = copy.deepcopy(qf1)
graph_builder_2 = FullGraphBuilder(
input_node_dim=obs_dim + action_dim,
num_node=num_agent,
batch_size=variant['algorithm_kwargs']['batch_size'],
contain_self_loop=False)
cg2 = GraphContextNet(graph_builder_2,
obs_dim,
action_dim,
output_activation='lrelu0.2',
**variant['graph_kwargs'])
target_cg2 = copy.deepcopy(cg2)
qf2 = FlattenMlp(
input_size=variant['graph_kwargs']['node_dim'] + action_dim,
output_size=1,
hidden_sizes=[variant['qf_kwargs']['hidden_dim']] *
(variant['qf_kwargs']['num_layer'] - 1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
)
target_qf2 = copy.deepcopy(qf2)
graph_builder_ca = FullGraphBuilder(
input_node_dim=obs_dim + action_dim,
num_node=num_agent,
batch_size=variant['algorithm_kwargs']['batch_size'],
contain_self_loop=False)
from rlkit.torch.networks.gnn_networks import GNNNet
cgca = GNNNet(
pre_graph_builder=graph_builder_ca,
node_dim=variant['graph_kwargs']['node_dim'],
conv_type='GSage',
num_conv_layers=variant['graph_kwargs']['num_layer'],
hidden_activation='lrelu0.2',
output_activation='lrelu0.2',
)
from rlkit.torch.networks.layers import SplitLayer
from rlkit.torch.policies.tanh_gaussian_policy import TanhGaussianPolicy
cactor = nn.Sequential(
FlattenMlp(
input_size=variant['graph_kwargs']['node_dim'],
output_size=variant['cactor_kwargs']['hidden_dim'],
hidden_sizes=[variant['cactor_kwargs']['hidden_dim']] *
(variant['cactor_kwargs']['num_layer'] - 1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
output_activation=nn.LeakyReLU(negative_slope=0.2),
), nn.LeakyReLU(negative_slope=0.2),
SplitLayer(layers=[
nn.Linear(variant['policy_kwargs']['hidden_dim'], action_dim),
nn.Linear(variant['policy_kwargs']['hidden_dim'], action_dim)
]))
cactor = TanhGaussianPolicy(module=cactor)
policy_n, expl_policy_n, eval_policy_n = [], [], []
for i in range(num_agent):
policy = nn.Sequential(
FlattenMlp(
input_size=obs_dim,
output_size=variant['policy_kwargs']['hidden_dim'],
hidden_sizes=[variant['policy_kwargs']['hidden_dim']] *
(variant['policy_kwargs']['num_layer'] - 1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
output_activation=nn.LeakyReLU(negative_slope=0.2),
),
SplitLayer(layers=[
nn.Linear(variant['policy_kwargs']['hidden_dim'], action_dim),
nn.Linear(variant['policy_kwargs']['hidden_dim'], action_dim)
]))
policy = TanhGaussianPolicy(module=policy)
from rlkit.torch.policies.make_deterministic import MakeDeterministic
eval_policy = MakeDeterministic(policy)
if variant['random_exploration']:
from rlkit.exploration_strategies.base import PolicyWrappedWithExplorationStrategy
from rlkit.exploration_strategies.epsilon_greedy import EpsilonGreedy
expl_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=EpsilonGreedy(expl_env.action_space,
prob_random_action=1.0),
policy=policy,
)
else:
expl_policy = policy
policy_n.append(policy)
expl_policy_n.append(expl_policy)
eval_policy_n.append(eval_policy)
from rlkit.samplers.data_collector.ma_path_collector import MAMdpPathCollector
eval_path_collector = MAMdpPathCollector(eval_env, eval_policy_n)
expl_path_collector = MAMdpPathCollector(expl_env, expl_policy_n)
from rlkit.data_management.ma_env_replay_buffer import MAEnvReplayBuffer
replay_buffer = MAEnvReplayBuffer(variant['replay_buffer_size'],
expl_env,
num_agent=num_agent)
from rlkit.torch.r2g.r2g_gnn3 import R2GGNNTrainer
trainer = R2GGNNTrainer(env=expl_env,
cg1=cg1,
target_cg1=target_cg1,
qf1=qf1,
target_qf1=target_qf1,
cg2=cg2,
target_cg2=target_cg2,
qf2=qf2,
target_qf2=target_qf2,
cgca=cgca,
cactor=cactor,
policy_n=policy_n,
**variant['trainer_kwargs'])
from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
log_path_function=get_generic_ma_path_information,
**variant['algorithm_kwargs'])
algorithm.to(ptu.device)
# save init params
from rlkit.core import logger
snapshot = algorithm._get_snapshot()
file_name = osp.join(logger._snapshot_dir, 'itr_-1.pkl')
torch.save(snapshot, file_name)
algorithm.train()
def her_td3_experiment(variant):
env = variant['env_class'](**variant['env_kwargs'])
if 'history_len' in variant:
history_len = variant['history_len']
env = MultiTaskHistoryEnv(env, history_len=history_len)
if variant.get('make_silent_env', True):
env = MultitaskEnvToSilentMultitaskEnv(env)
if variant['normalize']:
env = NormalizedBoxEnv(env)
exploration_type = variant['exploration_type']
if exploration_type == 'ou':
es = OUStrategy(
action_space=env.action_space,
**variant['es_kwargs']
)
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
**variant['es_kwargs'],
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
**variant['es_kwargs'],
)
else:
raise Exception("Invalid type: " + exploration_type)
obs_dim = env.observation_space.low.size
action_dim = env.action_space.low.size
goal_dim = env.goal_space.low.size
qf1 = ConcatMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs']
)
qf2 = ConcatMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs']
)
policy = TanhMlpPolicy(
input_size=obs_dim + goal_dim,
output_size=action_dim,
**variant['policy_kwargs']
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
replay_buffer = variant['replay_buffer_class'](
env=env,
**variant['replay_buffer_kwargs']
)
algorithm = HerTd3(
env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
replay_buffer=replay_buffer,
**variant['algo_kwargs']
)
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
from cartpole import CartPoleEnv
expl_env = CartPoleEnv(mode=3)
eval_env = CartPoleEnv(mode=3)
num_agent = expl_env.num_agents
obs_dim = eval_env.observation_space.low.size
action_dim = eval_env.action_space.low.size
from rlkit.torch.networks.graph_builders import FullGraphBuilder
graph_builder_obs = FullGraphBuilder(
input_node_dim=obs_dim,
num_node=num_agent,
batch_size=variant['algorithm_kwargs']['batch_size'],
contain_self_loop=False)
from rlkit.torch.networks.gnn_networks import GNNNet
obs_gnn_1 = GNNNet(
graph_builder_obs,
hidden_activation='lrelu0.2',
output_activation='lrelu0.2',
**variant['graph_kwargs'],
)
graph_builder_eval = FullGraphBuilder(
input_node_dim=graph_builder_obs.output_node_dim,
num_node=num_agent,
batch_size=variant['algorithm_kwargs']['batch_size'],
contain_self_loop=False)
if variant['concat_emb']:
gnn_out_dim = int(obs_dim + variant['graph_kwargs']['node_dim'] *
variant['graph_kwargs']['num_conv_layers'])
else:
gnn_out_dim = variant['graph_kwargs']['node_dim']
from rlkit.torch.networks.networks import FlattenMlp
post_mlp1 = FlattenMlp(
input_size=gnn_out_dim,
output_size=1,
hidden_sizes=[variant['qf_kwargs']['hidden_dim']] *
(variant['qf_kwargs']['num_layer'] - 1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
)
from rlkit.torch.networks.graph_r2g_qnet2 import R2GQNet
qf1 = R2GQNet(
obs_gnn=obs_gnn_1,
pre_graph_builder=graph_builder_eval,
obs_dim=obs_dim,
action_dim=action_dim,
post_mlp=post_mlp1,
normalize_emb=False,
output_activation=None,
concat_emb=variant['concat_emb'],
**variant['graph_kwargs'],
)
target_qf1 = copy.deepcopy(qf1)
obs_gnn_2 = GNNNet(
graph_builder_obs,
hidden_activation='lrelu0.2',
output_activation='lrelu0.2',
**variant['graph_kwargs'],
)
post_mlp2 = FlattenMlp(
input_size=gnn_out_dim,
output_size=1,
hidden_sizes=[variant['qf_kwargs']['hidden_dim']] *
(variant['qf_kwargs']['num_layer'] - 1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
)
qf2 = R2GQNet(
obs_gnn=obs_gnn_2,
pre_graph_builder=graph_builder_eval,
obs_dim=obs_dim,
action_dim=action_dim,
post_mlp=post_mlp2,
normalize_emb=False,
output_activation=None,
concat_emb=variant['concat_emb'],
**variant['graph_kwargs'],
)
target_qf2 = copy.deepcopy(qf2)
graph_builder_ca = FullGraphBuilder(
input_node_dim=obs_dim + action_dim,
num_node=num_agent,
batch_size=variant['algorithm_kwargs']['batch_size'],
contain_self_loop=False)
from rlkit.torch.networks.gnn_networks import GNNNet
cgca = GNNNet(
graph_builder_ca,
hidden_activation='lrelu0.2',
output_activation='lrelu0.2',
**variant['graph_kwargs'],
)
from rlkit.torch.networks.networks import FlattenMlp
from rlkit.torch.networks.layers import SplitLayer
from rlkit.torch.policies.tanh_gaussian_policy import TanhGaussianPolicy
cactor = nn.Sequential(
cgca,
FlattenMlp(
input_size=variant['graph_kwargs']['node_dim'],
output_size=variant['cactor_kwargs']['hidden_dim'],
hidden_sizes=[variant['cactor_kwargs']['hidden_dim']] *
(variant['cactor_kwargs']['num_layer'] - 1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
output_activation=nn.LeakyReLU(negative_slope=0.2),
), nn.LeakyReLU(negative_slope=0.2),
SplitLayer(layers=[
nn.Linear(variant['policy_kwargs']['hidden_dim'], action_dim),
nn.Linear(variant['policy_kwargs']['hidden_dim'], action_dim)
]))
cactor = TanhGaussianPolicy(module=cactor)
graph_builder_policy = FullGraphBuilder(
input_node_dim=obs_dim,
num_node=num_agent,
batch_size=variant['algorithm_kwargs']['batch_size'],
contain_self_loop=False)
policy_n, expl_policy_n, eval_policy_n = [], [], []
for i in range(num_agent):
policy = nn.Sequential(
FlattenMlp(
input_size=variant['graph_kwargs']['node_dim'],
output_size=variant['policy_kwargs']['hidden_dim'],
hidden_sizes=[variant['policy_kwargs']['hidden_dim']] *
(variant['policy_kwargs']['num_layer'] - 1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
output_activation=nn.LeakyReLU(negative_slope=0.2),
),
SplitLayer(layers=[
nn.Linear(variant['policy_kwargs']['hidden_dim'], action_dim),
nn.Linear(variant['policy_kwargs']['hidden_dim'], action_dim)
]))
policy = TanhGaussianPolicy(module=policy)
from rlkit.torch.policies.make_deterministic import MakeDeterministic
eval_policy = MakeDeterministic(policy)
if variant['random_exploration']:
from rlkit.exploration_strategies.base import PolicyWrappedWithExplorationStrategy
from rlkit.exploration_strategies.epsilon_greedy import EpsilonGreedy
expl_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=EpsilonGreedy(expl_env.action_space,
prob_random_action=1.0),
policy=policy,
)
else:
expl_policy = policy
policy_n.append(policy)
expl_policy_n.append(expl_policy)
eval_policy_n.append(eval_policy)
from rlkit.samplers.data_collector.ma_path_collector import MAMdpPathCollector
eval_path_collector = MAMdpPathCollector(eval_env,
eval_policy_n,
shared_encoder=obs_gnn_1)
expl_path_collector = MAMdpPathCollector(expl_env,
expl_policy_n,
shared_encoder=obs_gnn_1)
from rlkit.data_management.ma_env_replay_buffer import MAEnvReplayBuffer
replay_buffer = MAEnvReplayBuffer(variant['replay_buffer_size'],
expl_env,
num_agent=num_agent)
from rlkit.torch.r2g.r2g_gnn12 import R2GGNNTrainer
trainer = R2GGNNTrainer(env=expl_env,
qf1=qf1,
target_qf1=target_qf1,
qf2=qf2,
target_qf2=target_qf2,
cactor=cactor,
policy_n=policy_n,
**variant['trainer_kwargs'])
from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
log_path_function=get_generic_ma_path_information,
**variant['algorithm_kwargs'])
algorithm.to(ptu.device)
# save init params
from rlkit.core import logger
snapshot = algorithm._get_snapshot()
file_name = osp.join(logger._snapshot_dir, 'itr_-1.pkl')
torch.save(snapshot, file_name)
algorithm.train()
def her_td3_experiment(variant):
env = variant['env_class'](**variant['env_kwargs'])
observation_key = variant.get('observation_key', 'observation')
desired_goal_key = variant.get('desired_goal_key', 'desired_goal')
replay_buffer = ObsDictRelabelingBuffer(env=env,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
**variant['replay_buffer_kwargs'])
obs_dim = env.observation_space.spaces['observation'].low.size
action_dim = env.action_space.low.size
goal_dim = env.observation_space.spaces['desired_goal'].low.size
if variant['normalize']:
env = NormalizedBoxEnv(env)
exploration_type = variant['exploration_type']
if exploration_type == 'ou':
es = OUStrategy(action_space=env.action_space,
max_sigma=0.1,
**variant['es_kwargs'])
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
max_sigma=0.1,
min_sigma=0.1, # Constant sigma
**variant['es_kwargs'],
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
prob_random_action=0.1,
**variant['es_kwargs'],
)
else:
raise Exception("Invalid type: " + exploration_type)
qf1 = ConcatMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
qf2 = ConcatMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
policy = TanhMlpPolicy(input_size=obs_dim + goal_dim,
output_size=action_dim,
**variant['policy_kwargs'])
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
algorithm = HerTd3(env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
replay_buffer=replay_buffer,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
**variant['algo_kwargs'])
if ptu.gpu_enabled():
qf1.to(ptu.device)
qf2.to(ptu.device)
policy.to(ptu.device)
algorithm.to(ptu.device)
algorithm.train()
from rlkit.exploration_strategies.ou_strategy import OUStrategy
from rlkit.policies.simple import ZeroPolicy
import numpy as np
print("making env")
# env = SawyerPushAndReachXYEasyEnv()
env = SawyerResetFreePushEnv(
hide_goal=False,
puck_limit='large',
)
# env = MultitaskToFlatEnv(env)
policy = ZeroPolicy(env.action_space.low.size)
es = OUStrategy(env.action_space, theta=1)
es = EpsilonGreedy(
action_space=env.action_space,
prob_random_action=0.1,
)
policy = exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
print("starting rollout")
import pygame
from pygame.locals import QUIT, KEYDOWN
pygame.init()
screen = pygame.display.set_mode((400, 300))
char_to_action = {
'w': np.array([0, -1, 0, 0]),
def experiment(variant):
import sys
sys.path.append("./multiagent-particle-envs")
from make_env import make_env
from particle_env_wrapper import ParticleEnv
expl_env = ParticleEnv(
make_env(args.exp_name,
discrete_action_space=False,
world_args=variant['world_args']))
eval_env = ParticleEnv(
make_env(args.exp_name,
discrete_action_space=False,
world_args=variant['world_args']))
num_agent = expl_env.num_agent
obs_dim = eval_env.observation_space.low.size
action_dim = eval_env.action_space.low.size
from simple_spread_graph import SimpleSpreadGraphBuilder
graph_builder_1 = SimpleSpreadGraphBuilder(
num_agents=expl_env.scenario.num_agents,
num_landmarks=expl_env.scenario.num_landmarks,
batch_size=variant['algorithm_kwargs']['batch_size'],
append_action=True,
single_observe=False,
contain_self_loop=True,
)
from rlkit.torch.networks.gnn_networks import GNNNet
gnn1 = GNNNet(
graph_builder_1,
hidden_activation='lrelu0.2',
output_activation='lrelu0.2',
**variant['graph_kwargs'],
)
from rlkit.torch.networks.networks import FlattenMlp
from rlkit.torch.networks.layers import SelectLayer
qf1 = nn.Sequential(
gnn1, SelectLayer(dim=1, index=torch.arange(num_agent)),
FlattenMlp(
input_size=variant['graph_kwargs']['node_dim'],
output_size=1,
hidden_sizes=[variant['qf_kwargs']['hidden_dim']] *
(variant['qf_kwargs']['num_layer'] - 1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
))
target_qf1 = copy.deepcopy(qf1)
graph_builder_2 = SimpleSpreadGraphBuilder(
num_agents=expl_env.scenario.num_agents,
num_landmarks=expl_env.scenario.num_landmarks,
batch_size=variant['algorithm_kwargs']['batch_size'],
append_action=True,
single_observe=False,
contain_self_loop=True,
)
gnn2 = GNNNet(
graph_builder_2,
hidden_activation='lrelu0.2',
output_activation='lrelu0.2',
**variant['graph_kwargs'],
)
qf2 = nn.Sequential(
gnn2, SelectLayer(dim=1, index=torch.arange(num_agent)),
FlattenMlp(
input_size=variant['graph_kwargs']['node_dim'],
output_size=1,
hidden_sizes=[variant['qf_kwargs']['hidden_dim']] *
(variant['qf_kwargs']['num_layer'] - 1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
))
target_qf2 = copy.deepcopy(qf2)
policy_n, eval_policy_n, expl_policy_n = [], [], []
for i in range(num_agent):
graph_builder_policy = SimpleSpreadGraphBuilder(
num_agents=expl_env.scenario.num_agents,
num_landmarks=expl_env.scenario.num_landmarks,
batch_size=variant['algorithm_kwargs']['batch_size'],
append_action=False,
single_observe=True,
contain_self_loop=True,
)
gnn_policy = GNNNet(
graph_builder_policy,
hidden_activation='lrelu0.2',
output_activation='lrelu0.2',
**variant['graph_kwargs'],
)
from rlkit.torch.networks.layers import SplitLayer, FlattenLayer
policy = nn.Sequential(
gnn_policy, SelectLayer(dim=1, index=0), FlattenLayer(),
FlattenMlp(
input_size=variant['graph_kwargs']['node_dim'],
output_size=variant['policy_kwargs']['hidden_dim'],
hidden_sizes=[variant['policy_kwargs']['hidden_dim']] *
(variant['policy_kwargs']['num_layer'] - 1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
output_activation=nn.LeakyReLU(negative_slope=0.2),
),
SplitLayer(layers=[
nn.Linear(variant['policy_kwargs']['hidden_dim'], action_dim),
nn.Linear(variant['policy_kwargs']['hidden_dim'], action_dim)
]))
from rlkit.torch.policies.tanh_gaussian_policy import TanhGaussianPolicy
policy = TanhGaussianPolicy(module=policy)
from rlkit.torch.policies.make_deterministic import MakeDeterministic
eval_policy = MakeDeterministic(policy)
from rlkit.exploration_strategies.base import PolicyWrappedWithExplorationStrategy
if variant['random_exploration']:
from rlkit.exploration_strategies.epsilon_greedy import EpsilonGreedy
expl_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=EpsilonGreedy(expl_env.action_space,
prob_random_action=1.0),
policy=policy,
)
else:
expl_policy = policy
policy_n.append(policy)
eval_policy_n.append(eval_policy)
expl_policy_n.append(expl_policy)
from rlkit.samplers.data_collector.ma_path_collector import MAMdpPathCollector
eval_path_collector = MAMdpPathCollector(eval_env, eval_policy_n)
expl_path_collector = MAMdpPathCollector(expl_env, expl_policy_n)
from rlkit.data_management.ma_env_replay_buffer import MAEnvReplayBuffer
replay_buffer = MAEnvReplayBuffer(variant['replay_buffer_size'],
expl_env,
num_agent=num_agent)
from rlkit.torch.masac.masac_gnn import MASACGNNTrainer
trainer = MASACGNNTrainer(env=expl_env,
qf1=qf1,
target_qf1=target_qf1,
qf2=qf2,
target_qf2=target_qf2,
policy_n=policy_n,
**variant['trainer_kwargs'])
from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
log_path_function=get_generic_ma_path_information,
**variant['algorithm_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
num_agent = variant['num_agent']
from differential_game import DifferentialGame
expl_env = DifferentialGame(game_name=args.exp_name)
eval_env = DifferentialGame(game_name=args.exp_name)
obs_dim = eval_env.observation_space.low.size
action_dim = eval_env.action_space.low.size
policy_n, eval_policy_n, expl_policy_n, qf1_n, target_qf1_n, qf2_n, target_qf2_n = \
[], [], [], [], [], [], []
for i in range(num_agent):
from rlkit.torch.layers import SplitLayer, ReshapeLayer
weight_head = nn.Linear(variant['policy_kwargs']['hidden_dim'],
variant['policy_kwargs']['m'])
mean_head = nn.Sequential(
nn.Linear(variant['policy_kwargs']['hidden_dim'],
action_dim * variant['policy_kwargs']['m']),
ReshapeLayer(shape=[variant['policy_kwargs']['m'], action_dim]))
logstd_head = nn.Sequential(
nn.Linear(variant['policy_kwargs']['hidden_dim'],
action_dim * variant['policy_kwargs']['m']),
ReshapeLayer(shape=[variant['policy_kwargs']['m'], action_dim]))
policy = nn.Sequential(
nn.Linear(obs_dim, variant['policy_kwargs']['hidden_dim']),
nn.ReLU(),
nn.Linear(variant['policy_kwargs']['hidden_dim'],
variant['policy_kwargs']['hidden_dim']), nn.ReLU(),
SplitLayer(layers=[weight_head, mean_head, logstd_head]))
from rlkit.torch.policies.mix_tanh_gaussian_policy import MixTanhGaussianPolicy
policy = MixTanhGaussianPolicy(module=policy)
from rlkit.torch.policies.make_deterministic import MakeDeterministic
eval_policy = MakeDeterministic(policy)
from rlkit.exploration_strategies.base import PolicyWrappedWithExplorationStrategy
if variant['random_exploration']:
from rlkit.exploration_strategies.epsilon_greedy import EpsilonGreedy
expl_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=EpsilonGreedy(expl_env.action_space,
prob_random_action=1.0),
policy=policy,
)
else:
expl_policy = policy
from rlkit.torch.networks import FlattenMlp
qf1 = FlattenMlp(
input_size=(obs_dim * num_agent + action_dim * num_agent),
output_size=1,
hidden_sizes=[variant['qf_kwargs']['hidden_dim']] * 2,
)
target_qf1 = copy.deepcopy(qf1)
qf2 = FlattenMlp(
input_size=(obs_dim * num_agent + action_dim * num_agent),
output_size=1,
hidden_sizes=[variant['qf_kwargs']['hidden_dim']] * 2,
)
target_qf2 = copy.deepcopy(qf2)
policy_n.append(policy)
eval_policy_n.append(eval_policy)
expl_policy_n.append(expl_policy)
qf1_n.append(qf1)
target_qf1_n.append(target_qf1)
qf2_n.append(qf2)
target_qf2_n.append(target_qf2)
from rlkit.samplers.data_collector.ma_path_collector import MAMdpPathCollector
eval_path_collector = MAMdpPathCollector(eval_env, eval_policy_n)
expl_path_collector = MAMdpPathCollector(expl_env, expl_policy_n)
from rlkit.data_management.ma_env_replay_buffer import MAEnvReplayBuffer
replay_buffer = MAEnvReplayBuffer(variant['replay_buffer_size'],
expl_env,
num_agent=num_agent)
from rlkit.torch.masac.masac import MASACTrainer
trainer = MASACTrainer(env=expl_env,
qf1_n=qf1_n,
target_qf1_n=target_qf1_n,
qf2_n=qf2_n,
target_qf2_n=target_qf2_n,
policy_n=policy_n,
**variant['trainer_kwargs'])
from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
log_path_function=get_generic_ma_path_information,
**variant['algorithm_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
from multi_differential_game import MultiDifferentialGame
expl_env = MultiDifferentialGame(**variant['env_kwargs'])
eval_env = MultiDifferentialGame(**variant['env_kwargs'])
num_agent = expl_env.agent_num
obs_dim = eval_env.observation_space.low.size
action_dim = eval_env.action_space.low.size
from rlkit.torch.networks.graph_builders import FullGraphBuilder
graph_builder_1 = FullGraphBuilder(
input_node_dim=obs_dim + action_dim,
num_node=num_agent,
batch_size=variant['algorithm_kwargs']['batch_size'],
contain_self_loop=False)
from rlkit.torch.networks.gnn_networks import GNNNet
gnn1 = GNNNet(
graph_builder_1,
hidden_activation='lrelu0.2',
output_activation='lrelu0.2',
**variant['graph_kwargs'],
)
from rlkit.torch.networks.networks import FlattenMlp
qf1 = nn.Sequential(
gnn1,
FlattenMlp(
input_size=variant['graph_kwargs']['node_dim'],
output_size=1,
hidden_sizes=[variant['qf_kwargs']['hidden_dim']] *
(variant['qf_kwargs']['num_layer'] - 1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
))
target_qf1 = copy.deepcopy(qf1)
from rlkit.torch.networks.graph_builders import FullGraphBuilder
graph_builder_2 = FullGraphBuilder(
input_node_dim=obs_dim + action_dim,
num_node=num_agent,
batch_size=variant['algorithm_kwargs']['batch_size'],
contain_self_loop=False)
from rlkit.torch.networks.gnn_networks import GNNNet
gnn2 = GNNNet(
graph_builder_2,
hidden_activation='lrelu0.2',
output_activation='lrelu0.2',
**variant['graph_kwargs'],
)
qf2 = nn.Sequential(
gnn2,
FlattenMlp(
input_size=variant['graph_kwargs']['node_dim'],
output_size=1,
hidden_sizes=[variant['qf_kwargs']['hidden_dim']] *
(variant['qf_kwargs']['num_layer'] - 1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
))
target_qf2 = copy.deepcopy(qf2)
graph_builder_policy = FullGraphBuilder(
input_node_dim=obs_dim,
num_node=num_agent,
batch_size=variant['algorithm_kwargs']['batch_size'],
contain_self_loop=False)
shared_gnn = GNNNet(
graph_builder_policy,
hidden_activation='lrelu0.2',
output_activation='lrelu0.2',
**variant['graph_kwargs'],
)
policy_n, eval_policy_n, expl_policy_n = [], [], []
for i in range(num_agent):
from rlkit.torch.networks.layers import SplitLayer
policy = nn.Sequential(
FlattenMlp(
input_size=variant['graph_kwargs']['node_dim'],
output_size=variant['policy_kwargs']['hidden_dim'],
hidden_sizes=[variant['policy_kwargs']['hidden_dim']] *
(variant['policy_kwargs']['num_layer'] - 1),
hidden_activation=nn.LeakyReLU(negative_slope=0.2),
output_activation=nn.LeakyReLU(negative_slope=0.2),
),
SplitLayer(layers=[
nn.Linear(variant['policy_kwargs']['hidden_dim'], action_dim),
nn.Linear(variant['policy_kwargs']['hidden_dim'], action_dim)
]))
from rlkit.torch.policies.tanh_gaussian_policy import TanhGaussianPolicy
policy = TanhGaussianPolicy(module=policy)
from rlkit.torch.policies.make_deterministic import MakeDeterministic
eval_policy = MakeDeterministic(policy)
from rlkit.exploration_strategies.base import PolicyWrappedWithExplorationStrategy
if variant['random_exploration']:
from rlkit.exploration_strategies.epsilon_greedy import EpsilonGreedy
expl_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=EpsilonGreedy(expl_env.action_space,
prob_random_action=1.0),
policy=policy,
)
else:
expl_policy = policy
policy_n.append(policy)
eval_policy_n.append(eval_policy)
expl_policy_n.append(expl_policy)
from rlkit.samplers.data_collector.ma_path_collector import MAMdpPathCollector
eval_path_collector = MAMdpPathCollector(eval_env,
eval_policy_n,
shared_encoder=shared_gnn)
expl_path_collector = MAMdpPathCollector(expl_env,
expl_policy_n,
shared_encoder=shared_gnn)
from rlkit.data_management.ma_env_replay_buffer import MAEnvReplayBuffer
replay_buffer = MAEnvReplayBuffer(variant['replay_buffer_size'],
expl_env,
num_agent=num_agent)
from rlkit.torch.masac.masac_gnn import MASACGNNTrainer
trainer = MASACGNNTrainer(env=expl_env,
qf1=qf1,
target_qf1=target_qf1,
qf2=qf2,
target_qf2=target_qf2,
policy_n=policy_n,
shared_gnn=shared_gnn,
**variant['trainer_kwargs'])
from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
log_path_function=get_generic_ma_path_information,
**variant['algorithm_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
