def testReadPartialSlots(self):
spec = [
specs.TensorSpec([3], tf.float32, 'action'),
[
specs.TensorSpec([5], tf.float32, 'camera'),
specs.TensorSpec([3, 2], tf.float32, 'lidar')
]
]
replay_table = table.Table(spec, capacity=4)
batch_size = 2
action = 1 * np.ones([batch_size] + spec[0].shape.as_list())
camera = 2 * np.ones([batch_size] + spec[1][0].shape.as_list())
lidar = 3 * np.ones([batch_size] + spec[1][1].shape.as_list())
values = [action, [camera, lidar]]
tensors = nest.map_structure(
lambda x: tf.convert_to_tensor(x, dtype=tf.float32), values)
write_op = replay_table.write(list(range(batch_size)), tensors)
read_op = replay_table.read(list(range(batch_size)),
slots=['lidar', ['action']])
self.evaluate(tf.global_variables_initializer())
self.evaluate(write_op)
read_value_ = self.evaluate(read_op)
expected_values = [lidar, [action]]
nest.map_structure(self.assertAllClose, read_value_, expected_values)
def testReadWriteString(self):
spec = [
specs.TensorSpec([3], tf.float32, 'action'),
[
specs.TensorSpec([], tf.string, 'camera'),
specs.TensorSpec([3, 2], tf.float32, 'lidar')
]
]
replay_table = table.Table(spec, capacity=3)
variables = replay_table.variables()
self.assertEqual(3, len(variables))
self.assertAllEqual(
['Table/action:0', 'Table/camera:0', 'Table/lidar:0'],
[v.name for v in variables])
expected_values = [
1 * np.ones(spec[0].shape.as_list()),
[b'foo', 3 * np.ones(spec[1][1].shape.as_list())]
]
tensors = nest.map_structure(
lambda x, dtype: tf.convert_to_tensor(x, dtype=dtype),
expected_values, [tf.float32, [tf.string, tf.float32]])
write_op = replay_table.write(0, tensors)
read_op = replay_table.read(0)
self.evaluate(tf.global_variables_initializer())
self.evaluate(write_op)
read_value_ = self.evaluate(read_op)
self.assertAllClose(read_value_[0], expected_values[0])
self.assertEqual(read_value_[1][0], expected_values[1][0])
self.assertAllClose(read_value_[1][1], expected_values[1][1])
def _quantum_circuit_spec(self):
spec = {
'alpha': specs.TensorSpec(shape=[1], dtype=tf.float32),
'phi_g': specs.TensorSpec(shape=[1], dtype=tf.float32),
'phi_e': specs.TensorSpec(shape=[1], dtype=tf.float32)
}
return spec
def testReadWriteBatch(self):
spec = [
specs.TensorSpec([3], tf.float32, 'action'),
[
specs.TensorSpec([5], tf.float32, 'camera'),
specs.TensorSpec([3, 2], tf.float32, 'lidar')
]
]
replay_table = table.Table(spec, capacity=4)
batch_size = 2
expected_values = [
1 * np.ones([batch_size] + spec[0].shape.as_list()),
[
2 * np.ones([batch_size] + spec[1][0].shape.as_list()),
3 * np.ones([batch_size] + spec[1][1].shape.as_list())
]
]
tensors = tf.nest.map_structure(
lambda x: tf.convert_to_tensor(value=x, dtype=tf.float32),
expected_values)
write_op = replay_table.write(list(range(batch_size)), tensors)
read_op = replay_table.read(list(range(batch_size)))
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(write_op)
read_value_ = self.evaluate(read_op)
tf.nest.map_structure(self.assertAllClose, read_value_,
expected_values)
def testReadWriteDict(self):
spec = {
'action': specs.TensorSpec([3], tf.float32, 'action'),
'camera': specs.TensorSpec([5], tf.float32, 'camera'),
'lidar': specs.TensorSpec([3, 2], tf.float32, 'lidar')
}
replay_table = table.Table(spec, capacity=3)
variables = replay_table.variables()
self.assertEqual(3, len(variables))
self.assertAllEqual(
['Table/action:0', 'Table/camera:0', 'Table/lidar:0'],
[v.name for v in variables])
expected_values = {
'action': 1 * np.ones(spec['action'].shape.as_list()),
'camera': 2 * np.ones(spec['camera'].shape.as_list()),
'lidar': 3 * np.ones(spec['lidar'].shape.as_list())
}
tensors = nest.map_structure(
lambda x: tf.convert_to_tensor(x, dtype=tf.float32),
expected_values)
write_op = replay_table.write(0, tensors)
read_op = replay_table.read(0)
self.evaluate(tf.global_variables_initializer())
self.evaluate(write_op)
read_value_ = self.evaluate(read_op)
nest.map_structure(self.assertAllClose, read_value_, expected_values)
def _quantum_circuit_spec(self):
spec = {
'beta': specs.TensorSpec(shape=[2], dtype=tf.float32),
'epsilon': specs.TensorSpec(shape=[2], dtype=tf.float32),
'phi': specs.TensorSpec(shape=[1], dtype=tf.float32)
}
return spec
def testWritePartialSlots(self):
spec = [
specs.TensorSpec([3], tf.float32, 'action'),
[
specs.TensorSpec([5], tf.float32, 'camera'),
specs.TensorSpec([3, 2], tf.float32, 'lidar')
]
]
replay_table = table.Table(spec, capacity=4)
batch_size = 2
action1 = 1 * np.ones([batch_size] + spec[0].shape.as_list())
camera1 = 2 * np.ones([batch_size] + spec[1][0].shape.as_list())
lidar1 = 3 * np.ones([batch_size] + spec[1][1].shape.as_list())
write_op1 = replay_table.write(list(range(batch_size)),
[action1, [camera1, lidar1]])
lidar2 = 10 * np.ones([batch_size] + spec[1][1].shape.as_list())
action2 = 20 * np.ones([batch_size] + spec[0].shape.as_list())
write_op2 = replay_table.write(list(range(batch_size)),
[lidar2, [action2]],
['lidar', ['action']])
read_op = replay_table.read(list(range(batch_size)))
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(write_op1)
self.evaluate(write_op2)
read_value_ = self.evaluate(read_op)
expected_values = [action2, [camera1, lidar2]]
tf.nest.map_structure(self.assertAllClose, read_value_,
expected_values)
def testReadWriteNamedTuple(self):
# pylint: disable=invalid-name
Observation = collections.namedtuple('Observation',
['action', 'camera', 'lidar'])
# pylint: enable=invalid-name
spec = Observation(action=specs.TensorSpec([3], tf.float32, 'action'),
camera=specs.TensorSpec([5], tf.float32, 'camera'),
lidar=specs.TensorSpec([3, 2], tf.float32, 'lidar'))
replay_table = table.Table(spec, capacity=3)
variables = replay_table.variables()
self.assertEqual(3, len(variables))
self.assertAllEqual(
['Table/action:0', 'Table/camera:0', 'Table/lidar:0'],
[v.name for v in variables])
expected_values = Observation(
action=1 * np.ones(spec.action.shape.as_list()),
camera=2 * np.ones(spec.camera.shape.as_list()),
lidar=3 * np.ones(spec.lidar.shape.as_list()))
tensors = tf.nest.map_structure(
lambda x: tf.convert_to_tensor(value=x, dtype=tf.float32),
expected_values)
write_op = replay_table.write(0, tensors)
read_op = replay_table.read(0)
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(write_op)
read_value_ = self.evaluate(read_op)
tf.nest.map_structure(self.assertAllClose, read_value_,
expected_values)
def test_summary_no_exception(self): """Tests that Network.summary() does not throw an exception.""" observation_spec = specs.TensorSpec([1], tf.float32, 'observation') action_spec = specs.TensorSpec([2], tf.float32, 'action') net = MockNetwork(observation_spec, action_spec) net.create_variables() net.summary()
def _control_circuit_spec(self):
spec = {
'alpha': specs.TensorSpec(shape=[2], dtype=tf.float32),
'theta': specs.TensorSpec(shape=[7], dtype=tf.float32),
'phi': specs.TensorSpec(shape=[7], dtype=tf.float32)
}
return spec
def testReadWriteSingle(self):
spec = [
specs.TensorSpec([3], tf.float32, 'action'),
[
specs.TensorSpec([5], tf.float32, 'camera'),
specs.TensorSpec([3, 2], tf.float32, 'lidar')
]
]
replay_table = table.Table(spec, capacity=3)
variables = replay_table.variables()
self.assertEqual(3, len(variables))
self.assertAllEqual(
['Table/action:0', 'Table/camera:0', 'Table/lidar:0'],
[v.name for v in variables])
expected_values = [
1 * np.ones(spec[0].shape.as_list()),
[
2 * np.ones(spec[1][0].shape.as_list()),
3 * np.ones(spec[1][1].shape.as_list())
]
]
tensors = tf.nest.map_structure(
lambda x: tf.convert_to_tensor(value=x, dtype=tf.float32),
expected_values)
write_op = replay_table.write(0, tensors)
read_op = replay_table.read(0)
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(write_op)
read_value_ = self.evaluate(read_op)
tf.nest.map_structure(self.assertAllClose, read_value_,
expected_values)
def _control_circuit_spec(self):
spec = {
'alpha': specs.TensorSpec(shape=[2], dtype=tf.float32),
'beta': specs.TensorSpec(shape=[2], dtype=tf.float32),
'epsilon': specs.TensorSpec(shape=[2], dtype=tf.float32)
}
return spec
def __init__(self, time_step_spec, action_script):
"""
Input:
time_step_spec -- see tf-agents docs
action_script -- module or class with attributes 'alpha', 'beta',
'epsilon', 'phi' and 'period'.
"""
self.period = action_script.period # periodicity of the protocol
# load the script of actions and convert to tensors
self.script = action_script.script
for a, val in self.script.items():
self.script[a] = tf.constant(val, dtype=tf.float32)
# Calculate specs and call init of parent class
action_spec = {
a: specs.TensorSpec(shape=C.shape[1:], dtype=tf.float32)
for a, C in self.script.items()
}
policy_state_spec = specs.TensorSpec(shape=[], dtype=tf.int32)
super(ScriptedPolicy, self).__init__(time_step_spec,
action_spec,
policy_state_spec,
automatic_state_reset=True)
self._policy_info = ()
def _data_spec(self):
return [
specs.TensorSpec([3], tf.float32, 'action'),
[
specs.TensorSpec([5], tf.float32, 'lidar'),
specs.TensorSpec([3, 2], tf.float32, 'camera')
]
]
def test_variables_calls_build(self):
observation_spec = specs.TensorSpec([1], tf.float32, 'observation')
action_spec = specs.TensorSpec([2], tf.float32, 'action')
net = MockNetwork(observation_spec, action_spec)
self.assertFalse(net.built)
variables = net.variables
self.assertTrue(net.built)
self.assertLen(variables, 2)
def __init__(
self,
*args,
# Optional kwargs
H=1,
T=4,
attn_step=1,
episode_length=20,
batch_size=50,
init="vac",
reward_kwargs={'reward_mode': 'zero'},
encoding='square',
phase_space_rep='wigner',
**kwargs):
"""
Args:
H (int, optional): Horizon for history returned in observations. Defaults to 1.
T (int, optional): Periodicity of the 'clock' observation. Defaults to 4.
attn_step (int, optional): step size for hard-coded attention to
measurement outcomes. For example, set to 4 to return history
of measurement oucomes separated by 4 steps -- when the same
stabilizer is measured in the square code. In hexagonal code
this can be 2. Defaults to 1.
episode_length (int, optional): Number of iterations in training episode. Defaults to 20.
batch_size (int, optional): Vectorized minibatch size. Defaults to 50.
init (str, optional): Initial quantum state of system. Defaults to "vac".
reward_kwargs (dict, optional): optional dictionary of parameters
for the reward function of RL agent.
encoding (str, optional): Type of GKP lattice. Defaults to "square".
phase_space_rep (str, optional): phase space representation to use
for rendering ('wigner' or 'CF')
"""
# Default simulation parameters
self.H = H
self.T = T
self.attn_step = attn_step
self.episode_length = episode_length
self.batch_size = batch_size
self.init = init
self.phase_space_rep = phase_space_rep
self.setup_reward(reward_kwargs)
self.define_stabilizer_code(encoding)
self._epoch = 0
# Define action and observation specs
self.control_circuit = self._control_circuit
action_spec = self._control_circuit_spec
observation_spec = {
'msmt': specs.TensorSpec(shape=[self.H], dtype=tf.float32),
'clock': specs.TensorSpec(shape=[self.T], dtype=tf.float32),
'const': specs.TensorSpec(shape=[1], dtype=tf.float32)
}
time_step_spec = ts.time_step_spec(observation_spec)
super().__init__(time_step_spec, action_spec, self.batch_size)
def _control_circuit_spec(self):
spec = {
'beta': specs.TensorSpec(shape=[2], dtype=tf.float32),
'eps1': specs.TensorSpec(shape=[2], dtype=tf.float32),
'eps2': specs.TensorSpec(shape=[2], dtype=tf.float32),
'phi': specs.TensorSpec(shape=[1], dtype=tf.float32),
'theta': specs.TensorSpec(shape=[1], dtype=tf.float32)
}
return spec
def _get_mock_spec(self):
spec = [
specs.TensorSpec([3], tf.float32, 'action'),
[
specs.TensorSpec([5], tf.float32, 'lidar'),
specs.TensorSpec([3, 2], tf.float32, 'camera')
]
]
return spec
def test_create_variables(self):
observation_spec = specs.TensorSpec([1], tf.float32, 'observation')
action_spec = specs.TensorSpec([2], tf.float32, 'action')
net = MockNetwork(observation_spec, action_spec)
self.assertFalse(net.built)
with self.assertRaises(ValueError):
net.variables # pylint: disable=pointless-statement
net.create_variables()
self.assertTrue(net.built)
self.assertLen(net.variables, 2)
self.assertLen(net.trainable_variables, 1)
def testNetworkCreate(self):
observation_spec = specs.TensorSpec([1], tf.float32, 'observation')
action_spec = specs.TensorSpec([2], tf.float32, 'action')
net = MockNetwork(observation_spec, action_spec)
self.assertFalse(net.built)
with self.assertRaises(ValueError):
net.variables # pylint: disable=pointless-statement
output_spec = network.create_variables(net)
# MockNetwork adds some variables to observation, which has shape [bs, 1]
self.assertEqual(output_spec, tf.TensorSpec([1], dtype=tf.float32))
self.assertTrue(net.built)
self.assertLen(net.variables, 2)
self.assertLen(net.trainable_variables, 1)
def testSaveRestore(self):
spec = [
specs.TensorSpec([3], tf.float32),
specs.TensorSpec([5], tf.float32, 'lidar'),
specs.TensorSpec([3, 2], tf.float32, 'lidar')
]
replay_table = table.Table(spec, capacity=3)
self.evaluate(tf.compat.v1.global_variables_initializer())
directory = self.get_temp_dir()
prefix = os.path.join(directory, 'table')
root = tf.train.Checkpoint(table=replay_table)
save_path = root.save(prefix)
root.restore(save_path).assert_consumed().run_restore_ops()
def test_create_variables(self):
observation_spec = specs.TensorSpec([1], tf.float32, 'observation')
action_spec = specs.TensorSpec([1], tf.float32, 'action')
input_spec = (observation_spec, action_spec)
state_spec = tf.TensorSpec([], tf.float32, 'state')
net = MockStateFullNetwork(input_spec, state_spec)
self.assertFalse(net.built)
with self.assertRaises(ValueError):
net.variables # pylint: disable=pointless-statement
output_spec = net.create_variables()
self.assertEqual(output_spec, tf.TensorSpec([1, 1], dtype=tf.float32))
self.assertTrue(net.built)
self.assertLen(net.variables, 2)
self.assertLen(net.trainable_variables, 1)
def get_policy(self):
def policy_fn(observation, dtype=tf.int32):
if tf.rank(observation) < 1:
observation = [observation]
if self._latent_policy:
embed = self._embed_state(observation)
else:
embed = tf.one_hot(observation, self._num_states)
distribution = tf.matmul(
embed, tf.nn.softmax(self._embed_policy_logits, axis=-1))
policy_info = {'distribution': distribution}
return (tfp.distributions.Categorical(probs=distribution,
dtype=dtype), policy_info)
policy_info_spec = {
'log_probability':
specs.TensorSpec([], tf.float32),
'distribution':
specs.BoundedTensorSpec([self._num_actions],
tf.float32,
minimum=0.0,
maximum=1.0)
}
return policy_fn, policy_info_spec
def testSampleSingleCorrectProbabilityAsDataset(self, buffer_batch_size):
max_length = 3
spec = specs.TensorSpec([], tf.int32, 'action')
replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
spec, batch_size=buffer_batch_size, max_length=max_length)
actions = tf.stack([tf.Variable(0).count_up_to(9)] * buffer_batch_size)
add_op = replay_buffer.add_batch(actions)
ds = replay_buffer.as_dataset()
itr = ds.make_initializable_iterator()
_, buffer_info = itr.get_next()
probabilities = buffer_info.probabilities
with self.cached_session() as sess:
tf.global_variables_initializer().run()
itr.initializer.run()
num_adds = 5
for i in range(1, num_adds):
sess.run(add_op)
probabilities_ = sess.run(probabilities)
expected_probability = (
1. /
min(i * buffer_batch_size, max_length * buffer_batch_size))
self.assertAllClose(expected_probability, probabilities_)
def testSampleSingleCorrectProbabilityAsDataset(self, buffer_batch_size):
max_length = 3
spec = specs.TensorSpec([], tf.int32, 'action')
replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
spec, batch_size=buffer_batch_size, max_length=max_length)
actions = tf.stack([tf.Variable(0).count_up_to(9)] * buffer_batch_size)
self.evaluate(tf.compat.v1.global_variables_initializer())
ds = replay_buffer.as_dataset()
if tf.executing_eagerly():
add_op = lambda: replay_buffer.add_batch(actions)
itr = iter(ds)
sample = lambda: next(itr)
else:
add_op = replay_buffer.add_batch(actions)
itr = tf.compat.v1.data.make_initializable_iterator(ds)
self.evaluate(itr.initializer)
sample = itr.get_next()
num_adds = 5
for i in range(1, num_adds):
self.evaluate(add_op)
probabilities_ = self.evaluate(sample)[1].probabilities
expected_probability = (
1. /
min(i * buffer_batch_size, max_length * buffer_batch_size))
self.assertAllClose(expected_probability, probabilities_)
def testSampleSingleCorrectProbability(self, buffer_batch_size):
max_length = 3
spec = specs.TensorSpec([], tf.int32, 'action')
replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
spec, batch_size=buffer_batch_size, max_length=max_length)
actions = tf.stack([tf.Variable(0).count_up_to(9)] * buffer_batch_size)
@common.function
def add(actions):
replay_buffer.add_batch(actions)
@common.function
def probabilities():
_, buffer_info = replay_buffer.get_next()
return buffer_info.probabilities
self.evaluate(tf.compat.v1.global_variables_initializer())
num_adds = 5
for i in range(1, num_adds):
self.evaluate(add(actions))
probabilities_ = self.evaluate(probabilities())
expected_probability = (
1. /
min(i * buffer_batch_size, max_length * buffer_batch_size))
self.assertAllClose(expected_probability, probabilities_)
def testMultiStepStackedBatchedSampling(self, batch_size):
spec = specs.TensorSpec([], tf.int64, 'action')
replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
spec, batch_size=batch_size)
@common.function(autograph=True)
def add_data():
for i in tf.range(10, dtype=tf.int64):
replay_buffer.add_batch(
tf.ones((batch_size, ), dtype=tf.int64) * i)
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(add_data())
if tf.executing_eagerly():
steps = lambda: replay_buffer._get_next(
3, # pylint: disable=g-long-lambda
num_steps=2,
time_stacked=True)[0]
else:
steps, _ = replay_buffer._get_next(3,
num_steps=2,
time_stacked=True)
self.assertEqual(self.evaluate(steps).shape, (3, 2))
for _ in range(100):
steps_ = self.evaluate(steps)
self.assertAllEqual((steps_[:, 0] + 1) % 10, steps_[:, 1])
def __init__(self, time_step_spec):
sim_dir = r'E:\VladGoogleDrive\Qulab\GKP\sims'
# name = 'Benchmarking_HybridMarkovian4Rounds\supervised_lstm\lstm.hdf5'
name = 'Benchmarking_HybridMarkovian4Rounds\supervised_dnn\dnn.hdf5'
# name = 'Benchmarking_HybridMarkovian4Rounds\supervised_linear\shallow.hdf5'
self.model = keras.models.load_model(os.path.join(sim_dir, name))
action_spec = specs.TensorSpec(shape=[5], dtype=tf.float32)
super(SupervisedNeuralNet, self).__init__(time_step_spec, action_spec)
def _control_circuit_spec(self):
spec = { # SBS params
'beta': specs.TensorSpec(shape=[4, 2], dtype=tf.float32),
'phi': specs.TensorSpec(shape=[4, 2], dtype=tf.float32),
'flip': specs.TensorSpec(shape=[4, 2], dtype=tf.float32),
'detune': specs.TensorSpec(shape=[4, 2], dtype=tf.float32),
# Murch params
'Murch_phi': specs.TensorSpec(shape=[2], dtype=tf.float32),
'Murch_amp': specs.TensorSpec(shape=[2], dtype=tf.float32),
'Murch_detune_MHz': specs.TensorSpec(shape=[2], dtype=tf.float32),
# misc
'cavity_phase': specs.TensorSpec(shape=[1], dtype=tf.float32),
'Kerr_drive_amp': specs.TensorSpec(shape=[1], dtype=tf.float32)
}
return spec
def testDuplicateSpecNames(self):
spec = [
specs.TensorSpec([3], tf.float32, 'lidar'),
specs.TensorSpec([5], tf.float32, 'lidar'),
specs.TensorSpec([3, 2], tf.float32, 'lidar')
]
replay_table = table.Table(spec, capacity=3)
variables = replay_table.variables()
self.assertEqual(3, len(variables))
self.assertAllEqual(
['Table/lidar:0', 'Table/lidar_1:0', 'Table/lidar_2:0'],
[v.name for v in variables])
expected_slots = ['lidar', 'lidar_1', 'lidar_2']
self.assertAllEqual(replay_table.slots, expected_slots)
tensors = replay_table.read(0, expected_slots)
tf.nest.map_structure(lambda x, y: self.assertEqual(x.shape, y.shape),
spec, tensors)
