def avg_reward(action_script):
policy = plc.ScriptedPolicy(env.time_step_spec(),
action_script)
time_step = env.reset()
policy_state = policy.get_initial_state(env.batch_size)
while not time_step.is_last()[0]:
action_step = policy.action(time_step, policy_state)
policy_state = action_step.state
time_step = env.step(action_step.action)
return time_step.reward.numpy().mean()
policy = tf.compat.v2.saved_model.load(os.path.join(root_dir, policy_dir))
### simulate GKP stabilization with SBS
if 0:
env = env_init(control_circuit='gkp_qec_autonomous_sBs_osc_qb',
reward_kwargs={'reward_mode': 'zero'},
init='vac',
H=1,
T=2,
attn_step=1,
batch_size=1,
episode_length=10,
encoding='square')
from rl_tools.action_script import gkp_qec_autonomous_sBs_2round as action_script
policy = plc.ScriptedPolicy(env.time_step_spec(), action_script)
### simulate GKP state preparation with ECDC
if 0:
env = env_init(control_circuit='ECD_control',
reward_kwargs=dict(reward_mode='zero'),
init='vac',
T=11,
batch_size=1,
N=100,
episode_length=11)
from rl_tools.action_script import ECD_control_residuals_GKP_plusX_hex as action_script
policy = plc.ScriptedPolicy(env.time_step_spec(), action_script)
# N=40