Skip to content

Commit 093dae7

Browse files
salaastsalaast
authored
Add gradient ascent optimization algorithms (#247)
This is part of the effort to port the work by Krzysztof Choromanski, Mark Rowland, Vikas Sindhwani, Richard E. Turner, Adrian Weller: "Structured Evolution with Compact Architectures for Scalable Policy Optimization", https://arxiv.org/abs/1804.02395
1 parent aeb1eb2 commit 093dae7

File tree

2 files changed

+278
-0
lines changed

2 files changed

+278
-0
lines changed

‎compiler_opt/es/gradient_ascent_optimization_algorithms.py

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,189 @@
1+
# coding=utf-8
2+
# Copyright 2020 Google LLC
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
###############################################################################
17+
#
18+
# This is a port of the work by: Krzysztof Choromanski, Mark Rowland,
19+
# Vikas Sindhwani, Richard E. Turner, Adrian Weller: "Structured Evolution
20+
# with Compact Architectures for Scalable Policy Optimization",
21+
# https://arxiv.org/abs/1804.02395
22+
#
23+
###############################################################################
24+
r"""Library of gradient ascent algorithms.
25+
26+
Library of stateful gradient ascent algorithms taking as input the gradient and
27+
current parameters, and output the new parameters.
28+
"""
29+
30+
import abc
31+
import numpy as np
32+
33+
34+
# TODO(kchoro): Borrow JAXs optimizer library here. Integrated into Blackbox-v2.
35+
class GAOptimizer(metaclass=abc.ABCMeta):
36+
"""Abstract class for general gradient ascent optimizers.
37+
38+
Class is responsible for encoding different gradient ascent optimization
39+
techniques.
40+
"""
41+
42+
@abc.abstractmethod
43+
def run_step(self, current_input, gradient):
44+
"""Conducts a single step of gradient ascent optimization.
45+
46+
Conduct a single step of gradient ascent optimization procedure, given the
47+
current parameters and the raw gradient.
48+
49+
Args:
50+
current_input: the current parameters.
51+
gradient: the raw gradient.
52+
53+
Returns:
54+
New parameters by conducting a single step of gradient ascent.
55+
"""
56+
raise NotImplementedError("Abstract method")
57+
58+
@abc.abstractmethod
59+
def get_state(self):
60+
"""Returns the state of the optimizer.
61+
62+
Returns the state of the optimizer.
63+
64+
Args:
65+
66+
Returns:
67+
The state of the optimizer.
68+
"""
69+
raise NotImplementedError("Abstract method")
70+
71+
@abc.abstractmethod
72+
def set_state(self, state):
73+
"""Sets up the internal state of the optimizer.
74+
75+
Sets up the internal state of the optimizer.
76+
77+
Args:
78+
state: state to be set up
79+
80+
Returns:
81+
"""
82+
raise NotImplementedError("Abstract method")
83+
84+
85+
class MomentumOptimizer(GAOptimizer):
86+
"""Class implementing momentum gradient ascent optimizer.
87+
88+
Setting momentum coefficient to zero is equivalent to vanilla gradient
89+
ascent.
90+
91+
the state is the moving average as a list
92+
"""
93+
94+
def __init__(self, step_size, momentum):
95+
self.step_size = step_size
96+
self.momentum = momentum
97+
98+
self.moving_average = np.asarray([], dtype=np.float32)
99+
super().__init__()
100+
101+
def run_step(self, current_input, gradient):
102+
if self.moving_average.size == 0:
103+
# Initialize the moving average
104+
self.moving_average = np.zeros(len(current_input), dtype=np.float32)
105+
elif len(self.moving_average) != len(current_input):
106+
raise ValueError(
107+
"Dimensions of the parameters and moving average do not match")
108+
109+
if not isinstance(gradient, np.ndarray):
110+
gradient = np.asarray(gradient, dtype=np.float32)
111+
112+
self.moving_average = self.momentum * self.moving_average + (
113+
1 - self.momentum) * gradient
114+
step = self.step_size * self.moving_average
115+
116+
return current_input + step
117+
118+
def get_state(self):
119+
return self.moving_average.tolist()
120+
121+
def set_state(self, state):
122+
self.moving_average = np.asarray(state, dtype=np.float32)
123+
124+
125+
class AdamOptimizer(GAOptimizer):
126+
"""Class implementing ADAM gradient ascent optimizer.
127+
128+
The state is the first moment moving average, the second moment moving average,
129+
and t (current step number) combined in that order into one list
130+
"""
131+
132+
def __init__(self, step_size, beta1=0.9, beta2=0.999, epsilon=1e-07):
133+
self.step_size = step_size
134+
self.beta1 = beta1
135+
self.beta2 = beta2
136+
self.epsilon = epsilon
137+
138+
self.first_moment_moving_average = np.asarray([], dtype=np.float32)
139+
self.second_moment_moving_average = np.asarray([], dtype=np.float32)
140+
self.t = 0
141+
super().__init__()
142+
143+
def run_step(self, current_input, gradient):
144+
if self.first_moment_moving_average.size == 0:
145+
# Initialize the moving averages
146+
self.first_moment_moving_average = np.zeros(
147+
len(current_input), dtype=np.float32)
148+
self.second_moment_moving_average = np.zeros(
149+
len(current_input), dtype=np.float32)
150+
# Initialize the step counter
151+
self.t = 0
152+
elif len(self.first_moment_moving_average) != len(current_input):
153+
raise ValueError(
154+
"Dimensions of the parameters and moving averages do not match")
155+
156+
if not isinstance(gradient, np.ndarray):
157+
gradient = np.asarray(gradient, dtype=np.float32)
158+
159+
self.first_moment_moving_average = (
160+
self.beta1 * self.first_moment_moving_average +
161+
(1 - self.beta1) * gradient)
162+
self.second_moment_moving_average = (
163+
self.beta2 * self.second_moment_moving_average + (1 - self.beta2) *
164+
(gradient * gradient))
165+
166+
self.t += 1
167+
scale = np.sqrt(1 - self.beta2**self.t) / (1 - self.beta1**self.t)
168+
169+
step = self.step_size * scale * self.first_moment_moving_average / (
170+
np.sqrt(self.second_moment_moving_average) + self.epsilon)
171+
172+
return current_input + step
173+
174+
def get_state(self):
175+
return (self.first_moment_moving_average.tolist() +
176+
self.second_moment_moving_average.tolist() + [self.t])
177+
178+
def set_state(self, state):
179+
total_len = len(state)
180+
if total_len % 2 != 1:
181+
raise ValueError("The dimension of the state should be odd")
182+
dim = total_len // 2
183+
184+
self.first_moment_moving_average = np.asarray(state[:dim], dtype=np.float32)
185+
self.second_moment_moving_average = np.asarray(
186+
state[dim:2 * dim], dtype=np.float32)
187+
self.t = int(state[-1])
188+
if self.t < 0:
189+
raise ValueError("The step counter should be non-negative")

‎compiler_opt/es/gradient_ascent_optimization_algorithms_test.py

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
# coding=utf-8
2+
# Copyright 2020 Google LLC
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
###############################################################################
17+
#
18+
# This is a port of the work by: Krzysztof Choromanski, Mark Rowland,
19+
# Vikas Sindhwani, Richard E. Turner, Adrian Weller: "Structured Evolution
20+
# with Compact Architectures for Scalable Policy Optimization",
21+
# https://arxiv.org/abs/1804.02395
22+
#
23+
###############################################################################
24+
"""Tests for google3.learning.brain.contrib.blackbox.gradient_ascent_optimization_algorithms."""
25+
26+
import numpy as np
27+
28+
# from google3.learning.brain.contrib.blackbox import gradient_ascent_optimization_algorithms
29+
# from google3.testing.pybase import googletest
30+
31+
# from google3.testing.pybase import parameterized
32+
import gradient_ascent_optimization_algorithms
33+
from absl.testing import absltest
34+
35+
from absl.testing import parameterized
36+
37+
38+
class GradientAscentOptimizationAlgorithmsTest(parameterized.TestCase):
39+
40+
@parameterized.parameters((np.asarray([1., 2., 3.], dtype=np.float32),),
41+
(np.asarray([1.1, 2.2, 3.3], dtype=np.float32),))
42+
def test_momentum_set_state(self, state):
43+
optimizer = gradient_ascent_optimization_algorithms.MomentumOptimizer(
44+
0.1, 0.9)
45+
optimizer.set_state(state)
46+
recovered_state = optimizer.get_state()
47+
np.testing.assert_array_almost_equal(state, recovered_state)
48+
49+
@parameterized.parameters(
50+
(np.asarray([1., 2., 3., 4., 5., 6., 7], dtype=np.float32),),
51+
(np.asarray([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7], dtype=np.float32),))
52+
def test_adam_set_state(self, state):
53+
optimizer = gradient_ascent_optimization_algorithms.AdamOptimizer(0.1)
54+
optimizer.set_state(state)
55+
recovered_state = optimizer.get_state()
56+
np.testing.assert_array_almost_equal(state, recovered_state)
57+
58+
@parameterized.parameters(
59+
(0.1, 0.9, np.asarray([1.1, 0.0], dtype=np.float32),
60+
np.asarray([1.0, 1.0],
61+
dtype=np.float32), np.asarray([1.0, 1.0], dtype=np.float32),
62+
np.asarray([1.129, 0.029], dtype=np.float32)))
63+
def test_momentum_step(self, step_size, momentum, ini_parameter, gradient1,
64+
gradient2, final_parameter):
65+
optimizer = gradient_ascent_optimization_algorithms.MomentumOptimizer(
66+
step_size, momentum)
67+
parameter = ini_parameter
68+
parameter = optimizer.run_step(parameter, gradient1)
69+
parameter = optimizer.run_step(parameter, gradient2)
70+
np.testing.assert_array_almost_equal(parameter, final_parameter)
71+
72+
@parameterized.parameters(
73+
(0.1, 0.2, 0.5, np.asarray([1.1, 0.0], dtype=np.float32),
74+
np.asarray([1.0, 1.0],
75+
dtype=np.float32), np.asarray([1.0, 1.0], dtype=np.float32),
76+
np.asarray([1.3, 0.2], dtype=np.float32)))
77+
def test_adam_step(self, step_size, beta1, beta2, ini_parameter, gradient1,
78+
gradient2, final_parameter):
79+
optimizer = gradient_ascent_optimization_algorithms.AdamOptimizer(
80+
step_size, beta1, beta2)
81+
parameter = ini_parameter
82+
parameter = optimizer.run_step(parameter, gradient1)
83+
parameter = optimizer.run_step(parameter, gradient2)
84+
np.testing.assert_array_almost_equal(parameter, final_parameter)
85+
86+
87+
if __name__ == '__main__':
88+
# googletest.main()
89+
absltest.main()

0 commit comments

Comments
 (0)