-
Notifications
You must be signed in to change notification settings - Fork 5.6k
/
per_worker_gaussian_noise.py
49 lines (43 loc) · 1.74 KB
/
per_worker_gaussian_noise.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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
from gymnasium.spaces import Space
from typing import Optional
from ray.rllib.utils.annotations import OldAPIStack
from ray.rllib.utils.exploration.gaussian_noise import GaussianNoise
from ray.rllib.utils.schedules import ConstantSchedule
@OldAPIStack
class PerWorkerGaussianNoise(GaussianNoise):
"""A per-worker Gaussian noise class for distributed algorithms.
Sets the `scale` schedules of individual workers to a constant:
0.4 ^ (1 + [worker-index] / float([num-workers] - 1) * 7)
See Ape-X paper.
"""
def __init__(
self,
action_space: Space,
*,
framework: Optional[str],
num_workers: Optional[int],
worker_index: Optional[int],
**kwargs
):
"""
Args:
action_space: The gym action space used by the environment.
num_workers: The overall number of workers used.
worker_index: The index of the Worker using this
Exploration.
framework: One of None, "tf", "torch".
"""
scale_schedule = None
# Use a fixed, different epsilon per worker. See: Ape-X paper.
if num_workers > 0:
if worker_index > 0:
num_workers_minus_1 = float(num_workers - 1) if num_workers > 1 else 1.0
exponent = 1 + (worker_index / num_workers_minus_1) * 7
scale_schedule = ConstantSchedule(0.4**exponent, framework=framework)
# Local worker should have zero exploration so that eval
# rollouts run properly.
else:
scale_schedule = ConstantSchedule(0.0, framework=framework)
super().__init__(
action_space, scale_schedule=scale_schedule, framework=framework, **kwargs
)