code: import gym
  from gym.utils import seeding
  import numpy as np
  import json
  from gym_let_mpc.simulator import ControlSystem
  from gym_let_mpc.controllers import ETMPC, AHMPC
  import collections.abc
  import matplotlib.pyplot as plt
  from gym_let_mpc.utils import str_replace_whole_words
  import copy
  
  
  class LetMPCEnv(gym.Env):
      def __init__(self, config_path):
          with open(config_path) as file_object:
              config = json.load(file_object)
  
          if config["mpc"]["model"] == "plant":
              config["mpc"]["model"] = copy.deepcopy(config["plant"]["model"])
          elif config["mpc"]["model"].get("parameters", None) == "plant":
              config["mpc"]["model"]["parameters"] = copy.deepcopy(config["plant"]["model"]["parameters"])
  
          if config["lqr"]["model"] == "plant":
              config["lqr"]["model"] = copy.deepcopy(config["plant"]["model"])
          elif config["lqr"]["model"] == "mpc":
              config["lqr"]["model"] = copy.deepcopy(config["mpc"]["model"])
          elif config["lqr"]["model"].get("parameters", None) == "plant":
              config["lqr"]["model"]["parameters"] = copy.deepcopy(config["plant"]["model"]["parameters"])
          elif config["lqr"]["model"].get("parameters", None) == "mpc":
              config["lqr"]["model"]["parameters"] = copy.deepcopy(config["mpc"]["model"]["parameters"])
  
          self.config = config
          assert "max_steps" in self.config["environment"]
          self.max_steps = self.config["environment"]["max_steps"]
  
          assert "randomize" in self.config["environment"]
          assert "state" in self.config["environment"]["randomize"] and "reference" in self.config["environment"]["randomize"]
          assert "render" in self.config["environment"]
          if config["mpc"]["type"] == "ETMPC":
              assert len(config["environment"]["action"]["variables"]) == 1 and \
                     config["environment"]["action"]["variables"][0]["name"] == "mpc_compute"
              controller = ETMPC(config["mpc"], config["lqr"])
              self.action_space = gym.spaces.Discrete(2)
          elif config["mpc"]["type"] == "AHMPC":
              assert len(config["environment"]["action"]["variables"]) == 1 and \
                     config["environment"]["action"]["variables"][0]["name"] == "mpc_horizon"
              controller = AHMPC(config["mpc"])
              self.action_space = gym.spaces.Box(low=np.array([1]), high=np.array([50]), dtype=np.float32)
          else:
              raise ValueError
          self.control_system = ControlSystem(config["plant"], controller=controller)
          self.history = None
          self.steps_count = None
          self.np_random = None
          self.min_constraint_delta = 0.25  # TODO: how and where to set
  
          obs_high = []
          obs_low = []
          for obs_var in self.config["environment"]["observation"]["variables"]:
              for var_transform in obs_var.get("transform", ["none"]):
                  for lim_i, lim in enumerate(obs_var.get("limits", [None, None])):
                      if lim is None:
                          if lim_i == 0:
                              obs_low.append(-np.finfo(np.float32).max)
                          else:
                              obs_high.append(np.finfo(np.float32).max)
                      else:
                          if var_transform == "none":
                              if lim_i == 0:
                                  obs_low.append(lim)
                              else:
                                  obs_high.append(lim)
                          elif var_transform == "absolute":
                              if lim_i == 0:
                                  obs_low.append(0)
                              else:
                                  obs_high.append(lim)
                          elif var_transform == "square":
                              if lim_i == 0:
                                  obs_low.append(0)
                              else:
                                  obs_high.append(lim ** 2)
                          else:
                              raise NotImplementedError
          self.observation_space = gym.spaces.Box(low=np.array(obs_low, dtype=np.float32),
                                                  high=np.array(obs_high, dtype=np.float32),
                                                  dtype=np.float32)
  
          self.value_function_is_set = False
  
          self.viewer = None
  
      def seed(self, seed=None):
          """
          Seed the random number generator of the control system.
          :param seed: (int) seed for random state
          """
          self.np_random, seed = gym.utils.seeding.np_random(seed)
          self.control_system.seed(seed)
          return [seed]
  
      def reset(self, state=None, reference=None, constraint=None, model=None, process_noise=None, tvp=None):
          """
          Reset state of environment. Note that the simulator is reset, the MPC solution is computed and the first
          MPC action is applied to the plant.
  
          :param state: (dict) initial conditions (value) for state name (key).
          :param reference: (dict) reference value (value) for reference name (key).
          :param constraint: (dict) constraint values (value) for constraint names (key).
          :param model: (dict) dictionary of dictionary where first key is model that it applies to ["plant", "mpc", "lqr"],
          first value is dictionary of model parameters where second value is the specified model parameter value.
          :param process_noise: (dict) process noise values (value) as ndarray for state name (key). The process noise at
          each time step loops through the provided array.
          :param tvp: (dict) values of time-varying parameters. New values are generated if values arent specified
          for all time steps elapsed.
          :return: ([float]) observation vector
          """
          def update_dict_recursively(d, u):
              for k, v in u.items():
                  if isinstance(v, collections.abc.Mapping):
                      d[k] = update_dict_recursively(d.get(k, {}), v)
                  else:
                      d[k] = v
              return d
  
          sampled_state = self.sample_state()
          sampled_reference = self.sample_reference()
          sampled_constraint = self.sample_constraints()
          sampled_model = self.sample_model()
  
          if state is not None:
              sampled_state.update(state)
          elif len(sampled_state) == 0:
              sampled_state = None
          if reference is not None:
              sampled_reference.update(reference)
          elif len(sampled_reference) == 0:
              sampled_reference = None
          if constraint is not None:
              sampled_constraint.update(constraint)
          elif len(sampled_constraint) == 0:
              sampled_constraint = None
          if model is not None:
              sampled_model = update_dict_recursively(sampled_model, model)
          elif len(sampled_model) == 0:
              sampled_model = None
          self.control_system.reset(state=sampled_state, reference=sampled_reference, constraint=sampled_constraint,
                                    model=sampled_model, process_noise=process_noise, tvp=tvp)
          if self.config["mpc"]["type"] == "ETMPC":
              self.control_system.step(action=np.array([1]))
          obs = self.get_observation()
          self.history = {"obs": [obs], "actions": [], "rewards": []}
          self.steps_count = 0
  
          return obs
  
      def step(self, action):
          a_dict = {a_props["name"]: action[a_i]
                                          for a_i, a_props in enumerate(self.config["environment"]["action"]["variables"])}
  
          self.control_system.step(np.round(a_dict["mpc_horizon"]).astype(np.int32))#np.atleast_1d(int(a_dict["mpc_compute"])))
          self.history["actions"].append(a_dict)
          self.steps_count += 1
  
          info = {}
          obs = self.get_observation()
          done = False
          if self.steps_count >= self.max_steps:
              done = True
              info["termination"] = "steps"
          elif len(self.config["environment"].get("end_on_constraint_violation", [])) > 0:
              for c_name, c_d in self.control_system.get_constraint_distances().items():
                  if c_name.split("-")[1] in self.config["environment"]["end_on_constraint_violation"] and c_d > 0:
                      done = True
                      info["termination"] = "constraint"
                      break
  
          rew = self.get_reward(done=done)
          for category, v in self.config["environment"].get("info", {}).items():
              if category == "reward":
                  for rew_name, rew_expr in v.items():
                      info["reward/{}".format(rew_name)] = self.get_reward(rew_expr, done=done)
              else:
                  raise NotImplementedError
  
          if self.value_function_is_set:
              step_vf_data = {"mpc_state": self.control_system.get_state_vector(self.control_system.history["state"][-2]),
                              "mpc_next_state": self.control_system.controller.mpc_state_preds[:, -1, -1]}
              step_vf_data["mpc_n_horizon"] = self.control_system.controller.history["mpc_horizon"][-1]
              info["mpc_value_fn"] = (self.control_system.controller.value_function.eval([step_vf_data["mpc_next_state"].reshape(1, -1)])[0][0, 0]).astype(np.float64)
              step_vf_data["mpc_rewards"] = self.control_system.controller.mpc.opt_f_num.toarray()[0, 0] - \
                                            self.config["mpc"]["objective"].get("discount_factor") ** (step_vf_data["mpc_n_horizon"] + 1) * info["mpc_value_fn"]
              info["mpc_computation_time"] = sum([v for k, v in self.control_system.controller.mpc.solver_stats.items() if k.startswith("t_proc")])
              info["data"] = step_vf_data
              info["mpc_avg_stage_cost"] = step_vf_data["mpc_rewards"] / step_vf_data["mpc_n_horizon"]
  
          info.update({k: v.astype(np.float64) if hasattr(v, "dtype") else v for k, v in a_dict.items()})
  
          self.history["obs"].append(obs)
          self.history["rewards"].append(rew)
  
          return obs, rew, done, info
  
      def render(self, mode='human', save_path=None):  # TODO: add env renders
          figure, axes = None, None
          if self.viewer is None:
              env_plots = [plot_name for plot_name, make_plot in self.config["environment"]["render"].items() if make_plot]
              if len(env_plots) > 0:
                  figure, axes = plt.subplots(self.control_system.render_n_axes + len(env_plots), sharex=True,
                                              figsize=(9, 16))
              self.viewer = self.control_system.render(figure=figure, axes=axes, return_viewer=True)
              for i, plot in enumerate(env_plots):
                  self.viewer["axes"][plot] = axes[-(i + 1)]
          else:
              self.viewer = self.control_system.render(figure=figure, axes=axes, return_viewer=True)
          for plot_name, make_plot in self.config["environment"]["render"].items():
              if make_plot:
                  self.viewer["axes"][plot_name].set_ylabel("-".join(plot_name.split("_")[1:]))
                  x_data = np.array(range(self.steps_count)) * self.control_system.config["params"]["t_step"]
                  self.viewer["axes"][plot_name].clear()
                  if plot_name == "plot_action":
                      for a_var in self.config["environment"]["action"]["variables"]:
                          y_data = [step_a[a_var["name"]] for step_a in self.history["actions"]]
                          self.viewer["axes"][plot_name].plot(x_data, y_data, label=a_var["name"], drawstyle="steps")
                  elif plot_name == "plot_reward":
                      self.viewer["axes"][plot_name].plot(x_data, self.history["rewards"], label="reward")
                      self.viewer["axes"][plot_name].text(max(x_data) + self.control_system.config["params"]["t_step"],
                                                          self.history["rewards"][-1],
                                                          "{:.3f}".format(np.sum(self.history["rewards"])))
                  else:
                      raise ValueError
          for axis in self.viewer["axes"].values():
              axis.legend()
          if save_path is not None:
              self.viewer["figure"].savefig(save_path, bbox_inches="tight", format="png")
              plt.close(self.viewer["figure"])
          else:
              self.viewer["figure"].show()
  
      def get_observation(self):
          obs = []
          for var in self.config["environment"]["observation"]["variables"]:
              var_val = self._get_variable_value(var)
              for transform in var.get("transform", ["none"]):
                  if transform == "none":
                      obs.append(var_val)
                  elif transform == "absolute":
                      obs.append(abs(var_val))
                  elif transform == "square":
                      obs.append(var_val ** 2)
                  else:
                      raise ValueError
  
          return np.array(obs)
  
      def get_reward(self, rew_expr=None, done=False):
          if rew_expr is None:
              rew_expr = self.config["environment"]["reward"]["expression"]
  
          rew_expr = str_replace_whole_words(rew_expr, "done", int(done))
  
          for var in sorted(self.config["environment"]["reward"]["variables"], key=lambda x: len(x), reverse=True):
              var_val = self._get_variable_value(var)
              if isinstance(var_val, list) or isinstance(var_val, np.ndarray):  # TODO: needs to be better way to do this
                  var_val = var_val[0]
              rew_expr = str_replace_whole_words(rew_expr, var["name"], var_val)
  
          return eval(rew_expr)
  
      def _get_variable_value(self, var):
          if var["type"] == "state":
              val = self.control_system.current_state[var["name"]]
          elif var["type"] == "input":
              if var.get("value_type", "absolute") == "absolute":
                  val = self.control_system.controller.current_input[var["name"]]
              elif var.get("value_type") == "delta":
                  val = self.control_system.controller.history["inputs"][-2][var["name"]] - \
                        self.control_system.controller.current_input[var["name"]]
              else:
                  raise ValueError
          elif var["type"] == "reference":
              val = self.control_system.controller.current_reference[var["name"]]
          elif var["type"] == "tvp":
              val = self.control_system.tvps[var["name"]].get_values(self.steps_count)
          elif var["type"] == "error":
              val = self.control_system.controller.history["errors"][-1][var["name"]]
              if np.isnan(val):
                  val = 0
          elif var["type"] == "epsilon":
              val = self.control_system.controller.history["epsilons"][-1][var["name"]]
              if np.isnan(val):
                  val = 0
          elif var["type"] == "constraint":
              if var.get("value_type") == "distance":
                  val = self.control_system.get_constraint_distances((var["name"],))[var["name"]]
              else:
                  raise ValueError
          elif var["type"] == "action":
              if var.get("value_type", "agent") == "agent":
                  val = self.history["actions"][-1][var["name"]]
              elif var.get("value_type") == "controller":
                  val = self.control_system.controller.history[var["name"]][-1]
              else:
                  raise ValueError
          elif var["type"] == "time":
              if var.get("value_type") == "fraction":
                  val = self.control_system.controller.steps_since_mpc_computation / self.control_system.controller.mpc.n_horizon
              elif var.get("value_type") == "absolute":
                  val = self.control_system.controller.steps_since_mpc_computation
              else:
                  raise ValueError
          elif var["type"] == "parameter":
              if var["value_type"] in ["plant", "mpc", "lqr"]:
                  val = self.config[var["value_type"]]["model"]["parameters"][var["name"]]
              else:
                  raise ValueError
          else:
              raise ValueError
  
          if isinstance(val, np.ndarray):
              val = val[0]
          if "limits" in var:
              val = np.clip(val, var["limits"][0], var["limits"][1])
  
          return val
  
      def sample_constraints(self):
          constraints = {}
          for c_name, c_props in self.config["environment"].get("randomize", {}).get("constraints", {}).items():
              constraint_val = getattr(self.np_random, c_props["type"])(**c_props["kw"])
              if c_name.split("-")[1] in [k.split("-")[1] for k in constraints.keys()]:
                  other_bound_type = "u" if c_name.split("-")[2] == "l" else "l"
                  other_bound_val = constraints[c_name[:-1] + other_bound_type]
                  if other_bound_type == "u":
                      constraint_val = min(other_bound_val - self.min_constraint_delta, constraint_val)
                  else:
                      constraint_val = max(other_bound_val + self.min_constraint_delta, constraint_val)
              constraints[c_name] = constraint_val
          return constraints
  
      def sample_state(self):
          state = {}
          for s_name, s_props in self.config["environment"].get("randomize", {}).get("state", {}).items():
              state[s_name] = getattr(self.np_random, s_props["type"])(**s_props["kw"])
  
          return state
  
      def sample_reference(self):
          reference = {}
          for r_name, r_props in self.config["environment"].get("randomize", {}).get("reference", {}).items():
              reference[r_name] = getattr(self.np_random, r_props["type"])(**r_props["kw"])
  
          return reference
  
      def sample_model(self):
          model = {}
          for s_name, s_props in self.config["environment"].get("randomize", {}).get("model", {}).get("states", {}).items():
              model["states"] = {s_name: {}}
              for component_name, component_props in s_props.items():
                  model["states"][s_name][component_name] = \
                      {comp_v_name: getattr(self.np_random, v_prop["type"])(**v_prop["kw"])
                                      for comp_v_name, v_prop in component_props.items()}
  
          model = {dest: model for dest in self.config["environment"].get("randomize", {}).get("model", {}).get("apply", [])}
          return model
  
      def stop(self):
          pass
  
      def create_dataset(self, n_scenarios):
          dataset = []
          self.reset()
          for i in range(n_scenarios):
              process_noise = np.array([self.control_system._get_process_noise() for i in range(self.max_steps)])
              ep_dict = {"state": self.sample_state(), "reference": self.sample_reference(),
                              "constraint": self.sample_constraints(), "model": self.sample_model(),
                         "process_noise": {}, "tvp": {}}
              s_i = 0
              for s_name, s_props in self.config["plant"]["model"]["states"].items():
                  if "W" in s_props:
                      ep_dict["process_noise"][s_name] = process_noise[:, s_i]
                      s_i += 1
              for tvp_name, tvp_obj in self.control_system.tvps.items():
                  tvp_obj.generate_values(self.max_steps)
                  ep_dict["tvp"][tvp_name] = tvp_obj.values
              dataset.append(ep_dict)
              self.reset()
  
          return dataset
  
      def set_value_function(self, input_ph, output_ph, tf_session):
          self.control_system.controller.set_value_function(input_ph, output_ph, tf_session)
          self.value_function_is_set = True
  
      def set_learning_status(self, status):
          if self.value_function_is_set:
              self.control_system.controller.value_function.set_enabled(status)
  
  
  if __name__ == "__main__":  # TODO: constraints on pendulum and end episode if constraints violated
      env = LetMPCEnv("configs/cart_pendulum_horizon.json")
      env.seed(0)
  
      """
      from tensorflow_casadi import TensorFlowEvaluator, MLP
      import tensorflow as tf
      a = tf.placeholder(shape=(None, 4), dtype=tf.float32)
      mlp = MLP(a)
      sess = tf.Session()
      val_fun = TensorFlowEvaluator([mlp.input_ph], [mlp.output], sess)
      env.set_value_function(mlp.input_ph, mlp.output, sess)
      """
  
      import pickle
      with open("../../lmpc-horizon/datasets/cart_pendulum_10.pkl", "rb") as f:
          test_set = pickle.load(f)
  
      rews = {}
  
      for i in range(1):
          import time
          obs = env.reset(**test_set[5])
  
          done = False
          t_before = time.process_time()
          horizon = 10
          while not done:
              t_step = time.process_time()
              if env.steps_count % 1 == 0 and False:
                  horizon = 25 if horizon == 50 else 50
              obs, rew, done, info = env.step([horizon])#[np.random.randint(1, 10)])
              for rew_comp, v in info.items():
                  if rew_comp.startswith("reward/"):
                      if rew_comp not in rews:
                          rews[rew_comp] = []
                      rews[rew_comp].append(v)
              if time.process_time() - t_step > 1:
                  print(env.control_system.controller.mpc.solver_stats)
              print(env.steps_count)
  
          for k, v in rews.items():
              print("{}: {}".format(k, sum(v)))
          print("Elapsed time {}".format(time.process_time() - t_before))
          env.render()
  
  
          
  
  
  
  
apis: numpy.sum numpy.clip matplotlib.pyplot.close numpy.isnan numpy.array numpy.finfo numpy.round