text: import numpy as np
  import scipy.linalg as la
  import cvxpy as cp
  import torch
  import torch.optim as optim
  import argparse
  import setproctitle
  import os
  from gym import spaces
  import tqdm
  
  import policy_models as pm
  import disturb_models as dm
  import robust_mpc as rmpc
  
  from envs.random_nldi_env import RandomNLDIEnv
  from envs.cartpole import CartPoleEnv
  from envs.quadrotor_env import QuadrotorEnv
  from envs.random_pldi_env import RandomPLDIEnv
  from envs.random_hinf_env import RandomHinfEnv
  from envs.microgrid import MicrogridEnv
  
  from constants import *
  
  from rl.ppo import PPO
  from rl.rarl_ppo import RARLPPO
  from rl.model import Policy
  from rl.storage import RolloutStorage
  from rl import trainer
  from rl import arguments
  from envs.rl_wrapper import RLWrapper
  
  # import ipdb
  # import sys
  # from IPython.core import ultratb
  # sys.excepthook = ultratb.FormattedTB(mode='Verbose',
  #      color_scheme='Linux', call_pdb=1)
  
  
  def main():
      parser = argparse.ArgumentParser(
          description='Run robust control experiments.')
      parser.add_argument('--baseLR', type=float, default=1e-3,
                          help='learning rate for non-projected DPS')
      parser.add_argument('--robustLR', type=float, default=1e-4,
                          help='learning rate for projected DPS')
      parser.add_argument('--alpha', type=float, default=0.001,
                          help='exponential stability coefficient')
      parser.add_argument('--gamma', type=float, default=20,
                          help='bound on L2 gain of disturbance-to-output map (for H_inf control)')
      parser.add_argument('--epochs', type=int, default=1000,
                          help='max epochs')
      parser.add_argument('--test_frequency', type=int, default=20,
                          help='frequency of testing during training')
      parser.add_argument('--T', type=float, default=2,
                          help='time horizon in seconds')
      parser.add_argument('--dt', type=float, default=0.01,
                          help='time increment')
      parser.add_argument('--testSetSz', type=int, default=50,
                          help='size of test set')
      parser.add_argument('--holdSetSz', type=int, default=50,
                          help='size of holdout set')
      parser.add_argument('--trainBatchSz', type=int, default=20,
                          help='batch size for training')
      parser.add_argument('--stepType', type=str,
                          choices=['euler', 'RK4', 'scipy'], default='RK4',
                          help='method for taking steps during training')
      parser.add_argument('--testStepType', type=str,
                          choices=['euler', 'RK4', 'scipy'], default='RK4',
                          help='method for taking steps during testing')
      parser.add_argument('--env', type=str,
                          choices=['random_nldi-d0', 'random_nldi-dnonzero', 'random_pldi_env',
                          'random_hinf_env', 'cartpole', 'quadrotor', 'microgrid'],
                          default='random_nldi-d0',
                          help='environment')
      parser.add_argument('--envRandomSeed', type=int, default=10,
                          help='random seed used to construct the environment')
      parser.add_argument('--save', type=str,
                          help='prefix to add to save path')
      parser.add_argument('--gpu', type=int, default=0,
                          help='prefix to add to save path')
      parser.add_argument('--evaluate', type=str,
                          help='instead of training, evaluate the models from a given directory'
                               ' (remember to use the same random seed)')
      args = parser.parse_args()
  
      dt = args.dt
      save_sub = '{}+alpha{}+gamma{}+testSz{}+holdSz{}+trainBatch{}+baselr{}+robustlr{}+T{}+stepType{}+testStepType{}+seed{}+dt{}'.format(
          args.env, args.alpha, args.gamma, args.testSetSz, args.holdSetSz,
          args.trainBatchSz, args.baseLR, args.robustLR, args.T,
          args.stepType, args.testStepType, args.envRandomSeed, dt)
      if args.save is not None:
          save = os.path.join('results', '{}+{}'.format(args.save, save_sub))
      else:
          save = os.path.join('results', save_sub)
      trained_model_dir = os.path.join(save, 'trained_models')
      if not os.path.exists(trained_model_dir):
          os.makedirs(trained_model_dir)
      setproctitle.setproctitle(save_sub)
      
      device = torch.device('cuda:%d' % args.gpu if torch.cuda.is_available() else 'cpu')
      if torch.cuda.is_available():
          torch.cuda.set_device(args.gpu)
      
      # Setup
      isD0 = (args.env == 'random_nldi-d0') or (args.env == 'quadrotor')  # no u dependence in disturbance bound
      problem_type = 'nldi'
      if 'random_nldi' in args.env:
          env = RandomNLDIEnv(isD0=isD0, random_seed=args.envRandomSeed, device=device)
      elif args.env == 'random_pldi_env':
          env = RandomPLDIEnv(random_seed=args.envRandomSeed, device=device)
          problem_type = 'pldi'
      elif args.env == 'random_hinf_env':
          env = RandomHinfEnv(T=args.T, random_seed=args.envRandomSeed, device=device)
          problem_type = 'hinf'
      elif args.env == 'cartpole':
          env = CartPoleEnv(random_seed=args.envRandomSeed, device=device)
      elif args.env == 'quadrotor':
          env = QuadrotorEnv(random_seed=args.envRandomSeed, device=device)
      elif args.env == 'microgrid':
          env = MicrogridEnv(random_seed=args.envRandomSeed, device=device)
      else:
          raise ValueError('No environment named %s' % args.env)
      evaluate_dir = args.evaluate
      evaluate = evaluate_dir is not None
  
      # Test and holdout set of states
      torch.manual_seed(17)
      x_test = env.gen_states(num_states=args.testSetSz, device=device)
      x_hold = env.gen_states(num_states=args.holdSetSz, device=device)
      num_episode_steps = int(args.T / dt)
  
      if problem_type == 'nldi':
          A, B, G, C, D, Q, R = env.get_nldi_linearization()
          state_dim = A.shape[0]
          action_dim = B.shape[1]
  
          # Get LQR solutions
          Kct, Pct = get_lqr_tensors(A, B, Q, R, args.alpha, device)
  
          Kr, Sr = get_robust_lqr_sol(*(v.cpu().numpy() for v in (A, B, G, C, D, Q, R)), args.alpha)
          Krt = torch.tensor(Kr, device=device, dtype=TORCH_DTYPE)
          Prt = torch.tensor(np.linalg.inv(Sr), device=device, dtype=TORCH_DTYPE)
          stable_projection = pm.StableNLDIProjection(Prt, A, B, G, C, D, args.alpha, isD0)
  
          disturb_model = dm.MultiNLDIDisturbModel(x_test.shape[0], C, D, state_dim, action_dim, env.wp)
          disturb_model.to(device=device, dtype=TORCH_DTYPE)
  
      elif problem_type == 'pldi':
          A, B, Q, R = env.get_pldi_linearization()
          state_dim = A.shape[1]
          action_dim = B.shape[2]
  
          # Get LQR solutions
          Kct, Pct = get_lqr_tensors(A.mean(0), B.mean(0), Q, R, args.alpha, device)
  
          Kr, Sr = get_robust_pldi_policy(*(v.cpu().numpy() for v in (A, B, Q, R)), args.alpha)
          Krt = torch.tensor(Kr, device=device, dtype=TORCH_DTYPE)
          Prt = torch.tensor(np.linalg.inv(Sr), device=device, dtype=TORCH_DTYPE)
          stable_projection = pm.StablePLDIProjection(Prt, A, B)
  
          disturb_model = dm.MultiPLDIDisturbModel(x_test.shape[0], state_dim, action_dim, env.L)
          disturb_model.to(device=device, dtype=TORCH_DTYPE)
      
      elif problem_type == 'hinf':
          A, B, G, Q, R = env.get_hinf_linearization()
          state_dim = A.shape[0]
          action_dim = B.shape[1]
  
          # Get LQR solutions
          Kct, Pct = get_lqr_tensors(A, B, Q, R, args.alpha, device)
  
          Kr, Sr, mu = get_robust_hinf_policy(*(v.cpu().numpy() for v in (A, B, G, Q, R)), args.alpha, args.gamma)
          Krt = torch.tensor(Kr, device=device, dtype=TORCH_DTYPE)
          Prt = torch.tensor(np.linalg.inv(Sr), device=device, dtype=TORCH_DTYPE)
          stable_projection = pm.StableHinfProjection(Prt, A, B, G, Q, R, args.alpha, args.gamma, 1/mu)
  
          disturb_model = dm.MultiHinfDisturbModel(x_test.shape[0], state_dim, action_dim, env.wp, args.T)
          disturb_model.to(device=device, dtype=TORCH_DTYPE)
  
      else:
          raise ValueError('No problem type named %s' % problem_type)
  
      adv_disturb_model = dm.MBAdvDisturbModel(env, None, disturb_model, dt, horizon=num_episode_steps//5, update_freq=num_episode_steps//20)
      env.adversarial_disturb_f = adv_disturb_model
  
      ###########################################################
      # LQR baselines
      ###########################################################
  
      ### Vanilla LQR (i.e., non-robust, exponentially stable)
      pi_custom_lqr = lambda x: x @ Kct.T
      adv_disturb_model.set_policy(pi_custom_lqr)
  
      custom_lqr_perf = eval_model(x_test, pi_custom_lqr, env,
                                 step_type=args.testStepType, T=args.T, dt=dt)
      write_results(custom_lqr_perf, 'LQR', save)
      custom_lqr_perf = eval_model(x_test, pi_custom_lqr, env,
                                   step_type=args.testStepType, T=args.T, dt=dt, adversarial=True)
      write_results(custom_lqr_perf, 'LQR-adv', save)
  
      ### Robust LQR
      pi_robust_lqr = lambda x: x @ Krt.T
      adv_disturb_model.set_policy(pi_robust_lqr)
  
      robust_lqr_perf = eval_model(x_test, pi_robust_lqr, env,
                                   step_type=args.testStepType, T=args.T, dt=dt)
      write_results(robust_lqr_perf, 'Robust LQR', save)
      robust_lqr_perf = eval_model(x_test, pi_robust_lqr, env,
                                   step_type=args.testStepType, T=args.T, dt=dt, adversarial=True)
      write_results(robust_lqr_perf, 'Robust LQR-adv', save)
  
  
      ###########################################################
      # Model-based planning methods
      ###########################################################
  
      ### Non-robust MBP (starting with robust LQR solution)
      pi_mbp = pm.MBPPolicy(Krt, state_dim, action_dim)
      pi_mbp.to(device=device, dtype=TORCH_DTYPE)
      adv_disturb_model.set_policy(pi_mbp)
  
      if evaluate:
          pi_mbp.load_state_dict(torch.load(os.path.join(evaluate_dir, 'mbp.pt')))
      else:
          pi_mbp_dict, train_losses, hold_losses, test_losses, test_losses_adv, stop_epoch = \
              train(pi_mbp, x_test, x_hold, env,
                    lr=args.baseLR, batch_size=args.trainBatchSz, epochs=args.epochs, T=args.T, dt=dt, step_type=args.stepType,
                    test_frequency=args.test_frequency, save_dir=save, model_name='mbp', device=device)
          save_results(train_losses, hold_losses, test_losses, test_losses_adv, save, 'mbp', pi_mbp_dict, epoch=stop_epoch,
                       is_final=True)
          torch.save(pi_mbp_dict, os.path.join(trained_model_dir, 'mbp.pt'))
  
      pi_mbp_perf = eval_model(x_test, pi_mbp, env,
                              step_type=args.testStepType, T=args.T, dt=dt)
      write_results(pi_mbp_perf, 'MBP', save)
      pi_mbp_perf = eval_model(x_test, pi_mbp, env,
                              step_type=args.testStepType, T=args.T, dt=dt, adversarial=True)
      write_results(pi_mbp_perf, 'MBP-adv', save)
  
  
      ### Robust MBP (starting with robust LQR solution)
      pi_robust_mbp = pm.StablePolicy(pm.MBPPolicy(Krt, state_dim, action_dim), stable_projection)
      pi_robust_mbp.to(device=device, dtype=TORCH_DTYPE)
      adv_disturb_model.set_policy(pi_robust_mbp)
  
      if evaluate:
          pi_robust_mbp.load_state_dict(torch.load(os.path.join(evaluate_dir, 'robust_mbp.pt')))
      else:
          pi_robust_mbp_dict, train_losses, hold_losses, test_losses, test_losses_adv, stop_epoch = \
              train(pi_robust_mbp, x_test, x_hold, env,
                    lr=args.robustLR, batch_size=args.trainBatchSz, epochs=args.epochs, T=args.T, dt=dt, step_type=args.stepType,
                    test_frequency=args.test_frequency, save_dir=save, model_name='robust_mbp', device=device)
          save_results(train_losses, hold_losses, test_losses, test_losses_adv, save, 'robust_mbp', pi_robust_mbp_dict, epoch=stop_epoch,
                       is_final=True)
          torch.save(pi_robust_mbp_dict, os.path.join(trained_model_dir, 'robust_mbp.pt'))
  
      pi_robust_mbp_perf = eval_model(x_test, pi_robust_mbp, env,
                                     step_type=args.testStepType, T=args.T, dt=dt)
      write_results(pi_robust_mbp_perf, 'Robust MBP', save)
      pi_robust_mbp_perf = eval_model(x_test, pi_robust_mbp, env,
                                     step_type=args.testStepType, T=args.T, dt=dt, adversarial=True)
      write_results(pi_robust_mbp_perf, 'Robust MBP-adv', save)
  
  
      ###########################################################
      # RL methods
      ###########################################################
  
      if 'random_nldi' in args.env:
          if isD0:
              rmax = 1000
          else:
              rmax = 1000
      elif args.env == 'random_pldi_env':
          rmax = 10
      elif args.env == 'random_hinf_env':
          rmax = 1000
      elif args.env == 'cartpole':
          rmax = 10
      elif args.env == 'quadrotor':
          rmax = 1000
      elif args.env == 'microgrid':
          rmax = 10
      else:
          raise ValueError('No environment named %s' % args.env)
  
      rl_args = arguments.get_args()
      linear_controller_K = Krt
      linear_controller_P = Prt
      linear_transform = lambda u, x: u + x @ linear_controller_K.T
  
  
      ### Vanilla and robust PPO
      base_ppo_perfs = []
      base_ppo_adv_perfs = []
      robust_ppo_perfs = []
      robust_ppo_adv_perfs = []
      for seed in range(1):
          for robust in [False, True]:
              torch.manual_seed(seed)
  
              if robust:
                  # stable_projection = pm.StableNLDIProjection(linear_controller_P, A, B, G, C, D, args.alpha, isD0=isD0)
                  action_transform = lambda u, x: stable_projection.project_action(linear_transform(u, x), x)
              else:
                  action_transform = linear_transform
  
              envs = RLWrapper(env, state_dim, action_dim, gamma=rl_args.gamma,
                               dt=dt, rmax=rmax, step_type='RK4', action_transform=action_transform,
                               num_envs=rl_args.num_processes, device=device)
              eval_envs = RLWrapper(env, state_dim, action_dim, gamma=rl_args.gamma,
                                    dt=dt, rmax=rmax, step_type='RK4', action_transform=action_transform,
                                    num_envs=args.testSetSz, device=device)
  
              actor_critic = Policy(
                  envs.observation_space.shape,
                  envs.action_space,
                  base_kwargs={'recurrent': False})
              actor_critic.to(device=device, dtype=TORCH_DTYPE)
              agent = PPO(
                  actor_critic,
                  rl_args.clip_param,
                  rl_args.ppo_epoch,
                  rl_args.num_mini_batch,
                  rl_args.value_loss_coef,
                  rl_args.entropy_coef,
                  lr=rl_args.lr,
                  eps=rl_args.rms_prop_eps,
                  max_grad_norm=rl_args.max_grad_norm,
                  use_linear_lr_decay=rl_args.use_linear_lr_decay)
              rollouts = RolloutStorage(num_episode_steps, rl_args.num_processes,
                                        envs.observation_space.shape, envs.action_space,
                                        actor_critic.recurrent_hidden_state_size)
  
              ppo_pi = lambda x: action_transform(actor_critic.act(x, None, None, deterministic=True)[1], x)
              adv_disturb_model.set_policy(ppo_pi)
  
              if evaluate:
                  actor_critic.load_state_dict(torch.load(os.path.join(evaluate_dir,
                                                                       'robust_ppo.pt' if robust else 'ppo.pt')))
              else:
                  hold_costs, test_costs, adv_test_costs =\
                      trainer.train(agent, envs, rollouts, device, rl_args,
                                    eval_envs=eval_envs, x_hold=x_hold, x_test=x_test, num_episode_steps=num_episode_steps,
                                    save_dir=os.path.join(save, 'robust_ppo' if robust else 'ppo'),
                                    save_extension='%d' % seed)
                  save_results(np.zeros_like(hold_costs), hold_costs, test_costs, adv_test_costs, save,
                               'robust_ppo' if robust else 'ppo', actor_critic.state_dict(),
                               epoch=rl_args.num_env_steps, is_final=True)
                  torch.save(actor_critic.state_dict(), os.path.join(trained_model_dir,
                                                                     'robust_ppo.pt' if robust else 'ppo.pt'))
  
              ppo_perf = eval_model(x_test, ppo_pi, env,
                                    step_type=args.testStepType, T=args.T, dt=dt)
              ppo_adv_perf = eval_model(x_test, ppo_pi, env,
                                        step_type=args.testStepType, T=args.T, dt=dt, adversarial=True)
  
              if robust:
                  robust_ppo_perfs.append(ppo_perf.item())
                  robust_ppo_adv_perfs.append(ppo_adv_perf.item())
              else:
                  base_ppo_perfs.append(ppo_perf.item())
                  base_ppo_adv_perfs.append(ppo_adv_perf.item())
  
      write_results(base_ppo_perfs, 'PPO', save)
      write_results(robust_ppo_perfs, 'Robust PPO', save)
      write_results(base_ppo_adv_perfs, 'PPO-adv', save)
      write_results(robust_ppo_adv_perfs, 'Robust PPO-adv', save)
  
  
      # RARL PPO baseline
      adv_ppo_perfs = []
      adv_ppo_adv_perfs = []
      seed = 0
      torch.manual_seed(seed)
  
      action_transform = linear_transform
  
      envs = RLWrapper(env, state_dim, action_dim, gamma=rl_args.gamma,
                       dt=dt, rmax=rmax, step_type='RK4', action_transform=action_transform,
                       num_envs=rl_args.num_processes, device=device, rarl=True)
      eval_envs = RLWrapper(env, state_dim, action_dim, gamma=rl_args.gamma,
                            dt=dt, rmax=rmax, step_type='RK4', action_transform=action_transform,
                            num_envs=args.testSetSz, device=device)
  
      protagornist_ac = Policy(
          envs.observation_space.shape,
          envs.action_space,
          base_kwargs={'recurrent': False})
      protagornist_ac.to(device=device, dtype=TORCH_DTYPE)
      adversary_ac = Policy(
          envs.observation_space.shape,
          envs.disturb_space,
          base_kwargs={'recurrent': False})
      adversary_ac.to(device=device, dtype=TORCH_DTYPE)
      agent = RARLPPO(
          protagornist_ac,
          adversary_ac,
          rl_args.clip_param,
          rl_args.ppo_epoch,
          rl_args.num_mini_batch,
          rl_args.value_loss_coef,
          rl_args.entropy_coef,
          lr=rl_args.lr,
          eps=rl_args.rms_prop_eps,
          max_grad_norm=rl_args.max_grad_norm,
          use_linear_lr_decay=rl_args.use_linear_lr_decay)
      action_space = spaces.Box(low=0, high=1,
                                shape=(envs.action_space.shape[0]+envs.disturb_space.shape[0],), dtype=NUMPY_DTYPE)
      rollouts = RolloutStorage(num_episode_steps, rl_args.num_processes,
                                envs.observation_space.shape, action_space,
                                protagornist_ac.recurrent_hidden_state_size + adversary_ac.recurrent_hidden_state_size,
                                rarl=True)
  
      ppo_pi = lambda x: action_transform(protagornist_ac.act(x, None, None, deterministic=True)[1], x)
      adv_disturb_model.set_policy(ppo_pi)
  
      if evaluate:
          agent.load(evaluate_dir)
      else:
          hold_costs, test_costs, adv_test_costs = \
              trainer.train(agent, envs, rollouts, device, rl_args,
                            eval_envs=eval_envs, x_hold=x_hold, x_test=x_test,
                            num_episode_steps=num_episode_steps,
                            save_dir=os.path.join(save, 'rarl_ppo'),
                            save_extension='%d' % seed)
          save_results(np.zeros_like(hold_costs), hold_costs, test_costs, adv_test_costs, save,
                       'rarl_ppo', protagornist_ac.state_dict(),
                       epoch=rl_args.num_env_steps, is_final=True)
          agent.save(trained_model_dir)
      env.disturb_f.disturbance = None
  
      ppo_perf = eval_model(x_test, ppo_pi, env,
                            step_type=args.testStepType, T=args.T, dt=dt)
      ppo_adv_perf = eval_model(x_test, ppo_pi, env,
                                step_type=args.testStepType, T=args.T, dt=dt, adversarial=True)
  
      adv_ppo_perfs.append(ppo_perf.item())
      adv_ppo_adv_perfs.append(ppo_adv_perf.item())
  
      write_results(adv_ppo_perfs, 'RARL PPO', save)
      write_results(adv_ppo_adv_perfs, 'RARL PPO-adv', save)
  
  
      ###########################################################
      # MPC baselines
      ###########################################################
  
      ### Robust MPC (not implemented for H_infinity settings)
      if problem_type != 'hinf':
          if problem_type == 'nldi':
              robust_mpc_model = rmpc.RobustNLDIMPC(A, B, G, C, D, Q, R, Krt, device)
          else:
              robust_mpc_model = rmpc.RobustPLDIMPC(A, B, Q, R, Krt, device)
  
          pi_robust_mpc = robust_mpc_model.get_action
          adv_disturb_model.set_policy(pi_robust_mpc)
  
          robust_mpc_perf = eval_model(x_test, pi_robust_mpc, env,
                                  step_type=args.testStepType, T=args.T, dt=dt, adversarial=True)
          write_results(robust_mpc_perf, 'Robust MPC-adv', save)
  
  
  
  def get_lqr_tensors(At, Bt, Qt, Rt, alpha, device):
      K, S = get_custom_lqr_sol(*(v.cpu().numpy() for v in (At, Bt, Qt, Rt)), alpha)
      Kt = torch.tensor(K, device=device, dtype=TORCH_DTYPE)
      Pt = torch.tensor(np.linalg.inv(S), device=device, dtype=TORCH_DTYPE)
  
      return Kt, Pt
  
  
  def get_custom_lqr_sol(A, B, Q, R, alpha):
      n, m = B.shape
      S = cp.Variable((n, n), symmetric=True)
      Y = cp.Variable((m, n))
  
      R_sqrt = la.sqrtm(R)
      f = cp.trace(S @ Q) + cp.matrix_frac(Y.T @ R_sqrt, S)
  
      # Exponential stability constraints from LMI book
      cons = [S >> np.eye(n)]  # make LMI non-homogeneous
      cons += [A @ S + S @ A.T + B @ Y + Y.T @ B.T << -alpha * S]
  
      cp.Problem(cp.Minimize(f), cons).solve()
      K = np.linalg.solve(S.value, Y.value.T).T
      S = S.value
  
      return np.array(K), np.array(S)
  
  
  def get_robust_lqr_sol(A, B, G, C, D, Q, R, alpha):
      n, m = B.shape
      wq = C.shape[0]
  
      S = cp.Variable((n, n), symmetric=True)
      Y = cp.Variable((m, n))
      mu = cp.Variable()
  
      R_sqrt = la.sqrtm(R)
      f = cp.trace(S @ Q) + cp.matrix_frac(Y.T @ R_sqrt, S)
  
      cons_mat = cp.bmat((
          (A @ S + S @ A.T + cp.multiply(mu, G @ G.T) + B @ Y + Y.T @ B.T + alpha * S, S @ C.T + Y.T @ D.T),
          (C @ S + D @ Y, -cp.multiply(mu, np.eye(wq)))
      ))
      cons = [S >> 0, mu >= 1e-2] + [cons_mat << 0]
  
      try:
          prob = cp.Problem(cp.Minimize(f), cons)
          prob.solve(solver=cp.SCS)
      except cp.error.SolverError as e:
          raise ValueError('Solver failed with error: %s \n Try another environment seed' % e)
      K = np.linalg.solve(S.value, Y.value.T).T
  
      return K, S.value
  
  
  def get_robust_pldi_policy(A, B, Q, R, alpha):
      L, n, m = B.shape
      S = cp.Variable((n, n), symmetric=True)
      Y = cp.Variable((m, n))
  
      R_sqrt = la.sqrtm(R)
  
      f = cp.trace(S @ Q) + cp.matrix_frac(Y.T @ R_sqrt, S)
      cons = [S >> np.eye(n)] + [A[i, :, :] @ S + B[i, :, :] @ Y + S @ A[i, :, :].T + Y.T @ B[i, :, :].T << -alpha * S for i in range(A.shape[0])]
      prob = cp.Problem(cp.Minimize(f), cons)
      prob.solve(solver=cp.MOSEK)
      K = np.linalg.solve(S.value, Y.value.T).T
      return K, S.value
  
  
  def get_robust_hinf_policy(A, B, G, Q, R, alpha, gamma):
      n, m = B.shape
      wq = G.shape[1]
  
      S = cp.Variable((n, n), symmetric=True)
      Y = cp.Variable((m, n))
      mu = cp.Variable()
  
      Q_sqrt = la.sqrtm(Q)
      R_sqrt = la.sqrtm(R)
      f = cp.trace(S @ Q) + cp.matrix_frac(Y.T @ R_sqrt, S)
  
      cons_mat = cp.bmat([[S @ A.T + A @ S + Y.T @ B.T + B @ Y + alpha * S + (mu / gamma ** 2) * G @ G.T,
                           cp.bmat([[S @ Q_sqrt, Y.T @ R_sqrt]])],
                          [cp.bmat([[Q_sqrt @ S], [R_sqrt @ Y]]), -mu * np.eye(m + n)]])
      cons = [S >> np.eye(n), mu >= 0] + [cons_mat << -1e-3 * np.eye(n+m+n)]
  
      try:
          prob = cp.Problem(cp.Minimize(f), cons)
          prob.solve(solver=cp.SCS)  #cp.MOSEK)
      except cp.error.SolverError as e:
          raise ValueError('Solver failed with error: %s \n Try another environment seed' % e)
      K = np.linalg.solve(S.value, Y.value.T).T
      
      assert np.all(np.linalg.eigvals(S.value) > 0)
      assert np.all(mu.value > 0)
      assert np.all(np.linalg.eigvals(cons_mat.value) <= 0)
      
      return K, S.value, mu.value
  
  
  def eval_model(x, pi, env, step_type='euler', T=10, dt=0.05, adversarial=False):
      if adversarial:
          env.adversarial_disturb_f.reset()
      loss = 0
      # maxes = torch.ones(6, dtype=TORCH_DTYPE) * -np.inf
      # mins = torch.ones(6, dtype=TORCH_DTYPE) * np.inf
      for t in tqdm.tqdm(range(int(T / dt)), desc='Evaluating agent%s' % (' adversarial' if adversarial else '')):
          u = pi(x)
          if adversarial:
              env.adversarial_disturb_f.update(x)
          x, cost = env.step(x, u, t, step_type=step_type, dt=dt, adversarial=adversarial)
          loss += cost
  
          # maxes = torch.max(maxes, torch.max(x, dim=0)[0])
          # mins = torch.min(mins, torch.min(x, dim=0)[0])
      return loss.mean()
  
  
  def train(model, x_test, x_hold, env, batch_size=20, epochs=1000, test_frequency=10, lr=1e-4, T=1,
            dt=0.05, step_type='euler', save_dir=None, model_name=None, device=None, hinf_loss=False):
      opt = optim.Adam(model.parameters(), lr=lr)
      losses = []
      hold_losses = []
      test_losses = []
      test_losses_adv = []
      num_episode_steps = int(T / dt)
  
      for i in range(epochs+1):
          opt.zero_grad()
          x = env.gen_states(batch_size, device=device)
          loss = 0
          for t in range(num_episode_steps):
              # train
              model.train()
              u = model(x)
              x, cost = env.step(x, u, t, dt=dt, step_type=step_type)
              loss += cost
  
          losses.append(loss.mean().item())
          print('Epoch {}. Loss: mean/median {:.3f}/{:.3f}, min/max {:.3f}/{:.3f}'
                .format(i, torch.mean(loss), torch.median(loss), torch.min(loss), torch.max(loss)))
  
          loss.mean().backward()
          opt.step()
  
          if i % test_frequency == 0:
              print('Testing...')
              env.adversarial_disturb_f.reset()
              xh = x_hold.detach()
              xt = x_test.detach()
              xta = x_test.detach()
              hold_loss = 0
              test_loss = 0
              test_loss_adv = 0
              hold_disturb_norm = 0
              test_disturb_norm = 0
              test_disturb_norm_adv = 0
              for t in range(num_episode_steps):
                  # holdout
                  model.eval()
                  u_hold = model(xh)
                  xh, cost_h = env.step(xh, u_hold, t, dt=dt, step_type=step_type)
                  hold_loss += cost_h
                  if hinf_loss:
                      hold_disturb_norm += torch.norm(env.disturb, p=2, dim=1)
  
                  # test
                  model.eval()
                  u_test = model(xt)
                  xt, cost_t = env.step(xt, u_test, t, dt=dt, step_type=step_type)
                  test_loss += cost_t
                  if hinf_loss:
                      test_disturb_norm += torch.norm(env.disturb, p=2, dim=1)
  
                  # test adversarial
                  env.adversarial_disturb_f.update(xta)
                  model.eval()
                  u_test_adv = model(xta)
                  xta, cost_ta = env.step(xta, u_test_adv, t, dt=dt, step_type=step_type, adversarial=True)
                  test_loss_adv += cost_ta
                  if hinf_loss:
                      test_disturb_norm_adv += torch.norm(env.disturb, p=2, dim=1)
  
              hold_losses.append(hold_loss.mean().item())
              test_losses.append(test_loss.mean().item())
              test_losses_adv.append(test_loss_adv.mean().item())
  
              print('Hold Loss: mean/median {:.3f}/{:.3f}, min/max {:.3f}/{:.3f}'.format(
                  torch.mean(hold_loss), torch.median(hold_loss),
                  torch.min(hold_loss), torch.max(hold_loss)))
              print('Test Loss: mean/median {:.3f}/{:.3f}, min/max {:.3f}/{:.3f}'.format(
                  torch.mean(test_loss), torch.median(test_loss),
                  torch.min(test_loss), torch.max(test_loss)))
              print('Test Loss Adv: mean/median {:.3f}/{:.3f}, min/max {:.3f}/{:.3f}'.format(
                  torch.mean(test_loss_adv), torch.median(test_loss_adv),
                  torch.min(test_loss_adv), torch.max(test_loss_adv)))
              print('')
  
          # Save intermediate results
          if i % 100 == 0:
              save_results(np.array(losses), np.array(hold_losses), np.array(test_losses), np.array(test_losses_adv),
                           save_dir, model_name, model.state_dict(), epoch=i)
  
      return model.state_dict(), losses, hold_losses, test_losses, test_losses_adv, i
  
  
  def save_results(train_losses, hold_losses, test_losses, test_losses_adv,
                   save_dir, model_name, model_dict, epoch, is_final=False):
      model_save_dir = os.path.join(save_dir, model_name)
      if not os.path.exists(model_save_dir):
          os.makedirs(model_save_dir)
  
      np.save(os.path.join(model_save_dir, 'train_losses.npy'), np.array(train_losses))
      np.save(os.path.join(model_save_dir, 'hold_losses.npy'), np.array(hold_losses))
      np.save(os.path.join(model_save_dir, 'test_losses.npy'), np.array(test_losses))
      np.save(os.path.join(model_save_dir, 'test_losses_adv.npy'), np.array(test_losses_adv))
      torch.save(model_dict, os.path.join(model_save_dir, 'model-{}.pt'.format(epoch)))
      if is_final:
          torch.save(model_dict, os.path.join(model_save_dir, 'model.pt'))
  
  
  def write_results(test_loss, model_name, save_dir):
      if not os.path.exists(save_dir): os.makedirs(save_dir)
      result_str = '{}: {}\n'.format(model_name, test_loss)
      print(result_str)
      with open(os.path.join(save_dir, 'results.txt'), 'a') as f:
          f.write(result_str)
  
  
  if __name__ == '__main__':
      main()
  