-
Notifications
You must be signed in to change notification settings - Fork 717
/
piecewise_stochastic_environment.py
229 lines (191 loc) · 8.87 KB
/
piecewise_stochastic_environment.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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
# coding=utf-8
# Copyright 2018 The TF-Agents Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Bandit piecewise linear stationary environment."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import gin
import tensorflow as tf # pylint: disable=g-explicit-tensorflow-version-import
from tf_agents.bandits.environments import non_stationary_stochastic_environment as nsse
from tf_agents.specs import tensor_spec
def _raise_batch_shape_error(tensor_name, batch_shape):
raise ValueError('`{tensor_name}` must have batch shape with length 1; '
'got {batch_shape}.'.format(
tensor_name=tensor_name,
batch_shape=batch_shape))
@gin.configurable
class PiecewiseStationaryDynamics(nsse.EnvironmentDynamics):
"""A piecewise stationary environment dynamics.
This is a piecewise stationary environment which computes rewards as:
rewards(t) = observation(t) * observation_to_reward(i) + additive_reward(i)
where t is the environment time (env_time) and i is the index of each piece.
The environment time is incremented after the reward is computed while the
piece index is incremented at the end of the time interval. The parameters
observation_to_reward(i), additive_reward(i), and the length of interval, are
drawn from given distributions at the beginning of each temporal interval.
"""
def __init__(self,
observation_distribution,
interval_distribution,
observation_to_reward_distribution,
additive_reward_distribution):
"""Initialize the parameters of the piecewise dynamics.
Args:
observation_distribution: A distribution from tfp.distributions with shape
`[batch_size, observation_dim]` Note that the values of `batch_size` and
`observation_dim` are deduced from the distribution.
interval_distribution: A scalar distribution from `tfp.distributions`. The
value is casted to `int64` to update the time range.
observation_to_reward_distribution: A distribution from
`tfp.distributions` with shape `[observation_dim, num_actions]`. The
value `observation_dim` must match the second dimension of
`observation_distribution`.
additive_reward_distribution: A distribution from `tfp.distributions` with
shape `[num_actions]`. This models the non-contextual behavior of the
bandit.
"""
self._observation_distribution = observation_distribution
self._interval_distribution = interval_distribution
self._observation_to_reward_distribution = (
observation_to_reward_distribution)
self._additive_reward_distribution = additive_reward_distribution
observation_batch_shape = observation_distribution.batch_shape
reward_batch_shape = additive_reward_distribution.batch_shape
if observation_batch_shape.rank != 2:
_raise_batch_shape_error(
'observation_distribution', observation_batch_shape)
if reward_batch_shape.rank != 1:
_raise_batch_shape_error(
'additive_reward_distribution', reward_batch_shape)
if additive_reward_distribution.dtype != tf.float32:
raise ValueError('Reward must have dtype float32; got {}'.format(
self._reward.dtype))
expected_observation_to_reward_shape = [
tf.compat.dimension_value(
self._observation_distribution.batch_shape[1:]),
tf.compat.dimension_value(
self._additive_reward_distribution.batch_shape[0])]
observation_to_reward_shape = [
tf.compat.dimension_value(x)
for x in observation_to_reward_distribution.batch_shape]
if (observation_to_reward_shape !=
expected_observation_to_reward_shape):
raise ValueError(
'Observation to reward has {} as expected shape; got {}'.format(
observation_to_reward_shape,
expected_observation_to_reward_shape))
self._current_interval = tf.compat.v2.Variable(
tf.cast(interval_distribution.sample(), dtype=tf.int64),
dtype=tf.int64, name='interval')
self._current_observation_to_reward = tf.compat.v2.Variable(
observation_to_reward_distribution.sample(),
dtype=tf.float32,
name='observation_to_reward')
self._current_additive_reward = tf.compat.v2.Variable(
additive_reward_distribution.sample(),
dtype=tf.float32,
name='additive_reward')
@property
def batch_size(self):
return tf.compat.dimension_value(
self._observation_distribution.batch_shape[0])
@property
def observation_spec(self):
return tensor_spec.TensorSpec(
shape=self._observation_distribution.batch_shape[1:],
dtype=self._observation_distribution.dtype,
name='observation_spec')
@property
def action_spec(self):
return tensor_spec.BoundedTensorSpec(
shape=(),
dtype=tf.int32,
minimum=0,
maximum=tf.compat.dimension_value(
self._additive_reward_distribution.batch_shape[0]) - 1,
name='action')
def observation(self, unused_t):
return self._observation_distribution.sample()
def reward(self, observation, t):
def same_interval_parameters():
"""Returns the parameters of the current piece.
Returns:
The pair of `tf.Tensor` `(observation_to_reward, additive_reward)`.
"""
return [self._current_observation_to_reward,
self._current_additive_reward]
def new_interval_parameters():
"""Update and returns the piece parameters.
Returns:
The pair of `tf.Tensor` `(observation_to_reward, additive_reward)`.
"""
tf.compat.v1.assign_add(
self._current_interval,
tf.cast(self._interval_distribution.sample(), dtype=tf.int64))
tf.compat.v1.assign(self._current_observation_to_reward,
self._observation_to_reward_distribution.sample())
tf.compat.v1.assign(self._current_additive_reward,
self._additive_reward_distribution.sample())
return [self._current_observation_to_reward,
self._current_additive_reward]
observation_to_reward, additive_reward = tf.cond(
t < self._current_interval,
same_interval_parameters,
new_interval_parameters)
reward = (tf.matmul(observation, observation_to_reward) +
tf.reshape(additive_reward, [1, -1]))
return reward
@gin.configurable
def compute_optimal_reward(self, observation):
deterministic_reward = tf.matmul(
observation, self._current_observation_to_reward)
optimal_action_reward = tf.reduce_max(deterministic_reward, axis=-1)
return optimal_action_reward
@gin.configurable
def compute_optimal_action(self, observation):
deterministic_reward = tf.matmul(
observation, self._current_observation_to_reward)
optimal_action = tf.argmax(
deterministic_reward, axis=-1, output_type=tf.int32)
return optimal_action
@gin.configurable
class PiecewiseStochasticEnvironment(nsse.NonStationaryStochasticEnvironment):
"""Implements a piecewise stationary linear environment."""
def __init__(self,
observation_distribution,
interval_distribution,
observation_to_reward_distribution,
additive_reward_distribution):
"""Initialize the environment with the dynamics parameters.
Args:
observation_distribution: A distribution from `tfp.distributions` with
shape `[batch_size, observation_dim]`. Note that the values of
`batch_size` and `observation_dim` are deduced from the distribution.
interval_distribution: A scalar distribution from `tfp.distributions`. The
value is casted to `int64` to update the time range.
observation_to_reward_distribution: A distribution from
`tfp.distributions` with shape `[observation_dim, num_actions]`. The
value `observation_dim` must match the second dimension of
`observation_distribution`.
additive_reward_distribution: A distribution from `tfp.distributions` with
shape `[num_actions]`. This models the non-contextual behavior of the
bandit.
"""
super(PiecewiseStochasticEnvironment, self).__init__(
PiecewiseStationaryDynamics(
observation_distribution,
interval_distribution,
observation_to_reward_distribution,
additive_reward_distribution))