-
Notifications
You must be signed in to change notification settings - Fork 0
/
trainers.py
32 lines (23 loc) · 942 Bytes
/
trainers.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
"""
RLLib trainers modified to include a projection step after updating model parameters.
"""
from ray.rllib.agents import ppo, pg
from ray.rllib.utils.annotations import override
class ProjectedPGPolicy(pg.pg_torch_policy.PGTorchPolicy):
@override(pg.pg_torch_policy.PGTorchPolicy)
def apply_gradients(self, gradients):
super().apply_gradients(gradients)
self.model.project()
class ProjectedPPOPolicy(ppo.ppo_torch_policy.PPOTorchPolicy):
@override(ppo.ppo_torch_policy.PPOTorchPolicy)
def apply_gradients(self, gradients):
super().apply_gradients(gradients)
self.model.project()
class ProjectedPGTrainer(pg.PGTrainer):
@override(pg.PGTrainer)
def get_default_policy_class(self, config):
return ProjectedPGPolicy
class ProjectedPPOTrainer(ppo.PPOTrainer):
@override(ppo.PPOTrainer)
def get_default_policy_class(self, config):
return ProjectedPPOPolicy