text: #  Copyright 2020 Google LLC
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at
  #
  #      https://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  #  See the License for the specific language governing permissions and
  #  limitations under the License.
  
  from .hopper import HopperURDF
  
  from pybullet_utils import bullet_client
  import pybullet
  import time
  import gym, gym.utils.seeding, gym.spaces
  import numpy as np
  import math
  
  import os
  import inspect
  
  currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
  
  import torch
  from gan.wgan_models import Generator
  from gan import utils
  
  
  class HopperURDFEnvMB(gym.Env):
      metadata = {'render.modes': ['human', 'rgb_array'], 'video.frames_per_second': 50}
  
      def __init__(self,
                   render=True,
                   init_noise=True,
                   act_noise=True,
                   obs_noise=True,
                   control_skip=10,
                   using_torque_ctrl=True,
                   correct_obs_dx=True,  # if need to correct dx obs,
                   use_gen_dyn=False,
                   gen_dyn_path=None,
                   soft_floor_env=False,
                   low_torque_env=False
                   ):
  
          self.render = render
          self.init_noise = init_noise
          self.obs_noise = obs_noise
          self.act_noise = act_noise
          self.control_skip = int(control_skip)
          self._ts = 1. / 500.
          self.correct_obs_dx = correct_obs_dx
          self.use_gen_dyn = use_gen_dyn
  
          self.soft_floor_env = soft_floor_env
          self.low_torque_env = low_torque_env
  
          if self.render:
              self._p = bullet_client.BulletClient(connection_mode=pybullet.GUI)
          else:
              self._p = bullet_client.BulletClient()
  
          self.np_random = None
          self.robot = HopperURDF(init_noise=self.init_noise,
                                  time_step=self._ts,
                                  np_random=self.np_random)
          self.seed(0)  # used once temporarily, will be overwritten outside though superclass api
          self.viewer = None
          self.timer = 0
  
          self.floor_id = None
  
          self.obs = []
          self.reset()  # and update init obs
  
          self.action_dim = len(self.robot.ctrl_dofs)
          self.act = [0.0] * len(self.robot.ctrl_dofs)
          self.action_space = gym.spaces.Box(low=np.array([-1.] * self.action_dim),
                                             high=np.array([+1.] * self.action_dim))
          self.obs_dim = len(self.obs)
          obs_dummy = np.array([1.12234567] * self.obs_dim)
          self.observation_space = gym.spaces.Box(low=-np.inf * obs_dummy, high=np.inf * obs_dummy)
  
          self.gen_dyn = None
          if self.use_gen_dyn:
              self.gen_dyn = Generator()
              self.gen_dyn.load_state_dict(torch.load(gen_dyn_path))
              self.gen_dyn.eval()
  
      def reset(self):
          self._p.resetSimulation()
          self._p.setTimeStep(self._ts)
          self._p.setGravity(0, 0, -10)
          self.timer = 0
  
          self._p.setPhysicsEngineParameter(numSolverIterations=100)
          # self._p.setPhysicsEngineParameter(restitutionVelocityThreshold=0.000001)
  
          self.floor_id = self._p.loadURDF(os.path.join(currentdir, 'assets/plane.urdf'), [0, 0, 0.0], useFixedBase=1)
          self._p.changeDynamics(self.floor_id, -1, lateralFriction=1.0)  # TODO
          self._p.changeDynamics(self.floor_id, -1, restitution=.2)
  
          self.robot.reset(self._p)
          self.robot.update_x(reset=True)
          # # should be after reset!
  
          if self.soft_floor_env:
              self._p.changeDynamics(self.floor_id, -1, contactDamping=100.0, contactStiffness=600.0)
              for ind in range(self.robot.n_total_dofs):
                  self._p.changeDynamics(self.robot.hopper_id, ind, contactDamping=100.0, contactStiffness=600.0)
  
          if self.low_torque_env:
              self.robot.max_forces[2] = 200 / 1.6  # 1.6 for policy 4, 2.0 for policy 3
  
          #     self._p.changeDynamics(self.robot.hopper_id, ind, lateralFriction=1.0)
          #     self._p.changeDynamics(self.robot.hopper_id, ind, restitution=.2)
  
          # self._p.configureDebugVisualizer(pybullet.COV_ENABLE_PLANAR_REFLECTION, i)
  
          self._p.stepSimulation()
  
          self.update_extended_observation()
  
          return self.obs
  
      def step(self, a):
          if self.act_noise and a is not None:
              a = utils.perturb(a, 0.05, self.np_random)
  
          if not self.use_gen_dyn:
              for _ in range(self.control_skip):
                  # action is in not -1,1
                  if a is not None:
                      self.act = np.clip(a, -1.0, 1.0)
                      self.robot.apply_action(self.act)
                  self._p.stepSimulation()
                  if self.render:
                      time.sleep(self._ts * 0.5)
                  self.timer += 1
              self.robot.update_x()
              self.update_extended_observation()
              obs_unnorm = np.array(self.obs) / self.robot.obs_scaling
          else:
  
              # gen_input = list(self.obs) + list(a) + [0.0] * (11+3)
  
              gen_input = list(self.obs) + list(a) + list(np.random.normal(0, 1, 14))  # TODO
              gen_input = utils.wrap(gen_input, is_cuda=False)  # TODO
              gen_output = self.gen_dyn(gen_input)
              gen_output = utils.unwrap(gen_output, is_cuda=False)
  
              self.obs = gen_output
              obs_unnorm = np.array(self.obs) / self.robot.obs_scaling
              self.robot.last_x = self.robot.x
              self.robot.x += obs_unnorm[5] * (self.control_skip * self._ts)
  
              if self.render:
                  all_qs = [self.robot.x] + list(obs_unnorm[:5])
                  all_qs[1] -= 1.5
                  for ind in range(self.robot.n_total_dofs):
                      self._p.resetJointState(self.robot.hopper_id, ind, all_qs[ind], 0.0)
                  time.sleep(self._ts * 5.0)
  
              if self.obs_noise:
                  self.obs = utils.perturb(self.obs, 0.1, self.np_random)
              self.timer += self.control_skip
  
          reward = 2.0  # alive bonus
          reward += self.get_ave_dx()
          # print("v", self.get_ave_dx())
          reward += -0.1 * np.square(a).sum()
          # print("act norm", -0.1 * np.square(a).sum())
  
          q = np.array(obs_unnorm[2:5])
          pos_mid = 0.5 * (self.robot.ll + self.robot.ul)
          q_scaled = 2 * (q - pos_mid) / (self.robot.ul - self.robot.ll)
          joints_at_limit = np.count_nonzero(np.abs(q_scaled) > 0.97)
          reward += -2.0 * joints_at_limit
          # print("jl", -2.0 * joints_at_limit)
  
          dq = np.array(obs_unnorm[8:11])
          reward -= np.minimum(np.sum(np.abs(dq)) * 0.02, 5.0)  # almost like /23
          # print("vel pen", np.minimum(np.sum(np.abs(dq)) * 0.02, 5.0))
  
          height = obs_unnorm[0]
          # ang = self._p.getJointState(self.robot.hopper_id, 2)[0]
  
          # print(joints_dq)
          # print(height)
          # print("ang", ang)
          not_done = (np.abs(dq) < 50).all() and (height > .3) and (height < 1.8)
  
          return self.obs, reward, not not_done, {}
  
      def get_dist(self):
          return self.robot.x
  
      def get_ave_dx(self):
          if self.robot.last_x:
              return (self.robot.x - self.robot.last_x) / (self.control_skip * self._ts)
          else:
              return 0.0
  
      def update_extended_observation(self):
          self.obs = self.robot.get_robot_observation()
  
          if self.correct_obs_dx:
              dx = self.get_ave_dx() * self.robot.obs_scaling[5]
              self.obs[5] = dx
  
          if self.obs_noise:
              self.obs = utils.perturb(self.obs, 0.1, self.np_random)
  
      def seed(self, seed=None):
          self.np_random, seed = gym.utils.seeding.np_random(seed)
          self.robot.np_random = self.np_random  # use the same np_randomizer for robot as for env
          return [seed]
  
      def getSourceCode(self):
          s = inspect.getsource(type(self))
          s = s + inspect.getsource(type(self.robot))
          return s
  
      def cam_track_torso_link(self):
          distance = 5
          yaw = 0
          torso_x = self._p.getLinkState(self.robot.hopper_id, 2, computeForwardKinematics=1)[0]
          self._p.resetDebugVisualizerCamera(distance, yaw, -20, torso_x)
  