forked from v-sivak/quantum-control-rl
-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathT1_and_reward.py
62 lines (47 loc) · 2.51 KB
/
T1_and_reward.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
# -*- coding: utf-8 -*-
"""
Created on Tue Apr 7 16:16:45 2020
@author: Vladimir Sivak
"""
import os
os.environ["TF_FORCE_GPU_ALLOW_GROWTH"]='true'
os.environ["CUDA_VISIBLE_DEVICES"]="1"
import tensorflow as tf
from math import pi
from gkp.gkp_tf_env import helper_functions as hf
from gkp.gkp_tf_env import tf_env_wrappers as wrappers
import gkp.action_script as action_scripts
from gkp.gkp_tf_env import policy as plc
from gkp.gkp_tf_env import gkp_init
# env = gkp_init(simulate='phase_estimation_osc_qb_v2',
# reward_kwargs={'reward_mode':'fidelity', 'code_flips':True},
# init='X+', H=1, T=4, attn_step=1, batch_size=1000, episode_length=100,
# encoding='square')
env = gkp_init(simulate='gkp_qec_autonomous_sBs_osc_qb',
reward_kwargs={'reward_mode':'fidelity', 'code_flips':True},
init='X+', H=1, T=2, attn_step=1, batch_size=500, episode_length=60,
encoding='square')
# action_script = 'gkp_qec_autonomous_BsB_2round'
# action_scale = {'beta':1, 'phi':pi, 'epsilon':1}
# to_learn = {'beta':False, 'phi':False, 'epsilon':True}
# action_script = action_scripts.__getattribute__(action_script)
# env = wrappers.ActionWrapper(env, action_script, action_scale, to_learn)
# root_dir = r'E:\data\gkp_sims\PPO\examples\gkp_qec_autonomous_BsB'
# policy_dir = r'policy\000200'
# policy = tf.compat.v2.saved_model.load(os.path.join(root_dir,policy_dir))
# action_script = 'gkp_qec_autonomous_sBs_2round'
# action_scale = {'beta':1, 'phi':pi, 'eps1':1, 'eps2':1}
# to_learn = {'beta':False, 'phi':False, 'eps1':True, 'eps2':True}
# action_script = action_scripts.__getattribute__(action_script)
# env = wrappers.ActionWrapper(env, action_script, action_scale, to_learn)
# root_dir = r'E:\data\gkp_sims\PPO\examples\test'
# policy_dir = r'policy\000300'
# policy = tf.compat.v2.saved_model.load(os.path.join(root_dir,policy_dir))
# from gkp.action_script import gkp_qec_autonomous_BsB_2round as action_script
from gkp.action_script import gkp_qec_autonomous_sBs_2round as action_script
# from gkp.action_script import v2_phase_estimation_with_trim_4round as action_script
policy = plc.ScriptedPolicy(env.time_step_spec(), action_script)
#-----------------------------------------------------------------------------
#-----------------------------------------------------------------------------
#-----------------------------------------------------------------------------
fit_params = hf.fit_logical_lifetime(env, policy, plot=True, reps=1, states=['X+','Y+','Z+'])