text: # Copyright 2021 The Brax Authors.
  #
  # 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
  #
  #     http://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.
  
  """Trains a ur5e robot arm to move its end effector to a sequence of targets.
  
  The 6 joints have been placed faithfully to the actual joint locations of the
  ur5e robot arm model.  Because brax does not yet support meshes, the capsule
  locations for collider boundaries are only approximate.
  
  See https://www.universal-robots.com/products/ur5-robot/ for more details."""
  
  from typing import Tuple
  
  import dataclasses
  import jax
  import jax.numpy as jnp
  import brax
  from brax.envs import env
  from brax.physics import math
  
  from google.protobuf import text_format
  
  
  class Ur5e(env.Env):
    """Trains a UR5E robotic arm to touch a sequence of random targets."""
  
    def __init__(self, **kwargs):
      config = text_format.Parse(_SYSTEM_CONFIG, brax.Config())
      super().__init__(config, **kwargs)
      self.target_idx = self.sys.body_idx['Target']
      self.torso_idx = self.sys.body_idx['wrist_3_link']
      self.target_radius = .02
      self.target_distance = .5
  
    def reset(self, rng: jnp.ndarray) -> env.State:
      qp = self.sys.default_qp()
      rng, target = self._random_target(rng)
      pos = jax.ops.index_update(qp.pos, jax.ops.index[self.target_idx], target)
      qp = dataclasses.replace(qp, pos=pos)
      info = self.sys.info(qp)
      obs = self._get_obs(qp, info)
      reward, done, steps, zero = jnp.zeros(4)
      metrics = {
          'hits': zero,
          'weightedHits': zero,
          'movingToTarget': zero,
      }
      return env.State(rng, qp, info, obs, reward, done, steps, metrics)
  
    def step(self, state: env.State, action: jnp.ndarray) -> env.State:
      rng = state.rng
      qp, info = self.sys.step(state.qp, action)
      obs = self._get_obs(qp, info)
  
      # small reward for end effector moving towards target
      torso_delta = qp.pos[self.torso_idx] - state.qp.pos[self.torso_idx]
      target_rel = qp.pos[self.target_idx] - qp.pos[self.torso_idx]
      target_dist = jnp.linalg.norm(target_rel)
      target_dir = target_rel / (1e-6 + target_dist)
      moving_to_target = .1 * jnp.dot(torso_delta, target_dir)
  
      # big reward for reaching target
      target_hit = jnp.where(target_dist < self.target_radius, 1.0, 0.0)
      weighted_hit = target_hit
  
      reward = moving_to_target + weighted_hit
  
      steps = state.steps + self.action_repeat
      done = jnp.where(steps >= self.episode_length, 1.0, 0.0)
      metrics = {
          'hits': target_hit,
          'weightedHits': weighted_hit,
          'movingToTarget': moving_to_target,
      }
  
      # teleport any hit targets
      rng, target = self._random_target(rng)
      target = jnp.where(target_hit, target, qp.pos[self.target_idx])
      pos = jax.ops.index_update(qp.pos, jax.ops.index[self.target_idx], target)
      qp = dataclasses.replace(qp, pos=pos)
      return env.State(rng, qp, info, obs, reward, done, steps, metrics)
  
    def _get_obs(self, qp: brax.QP, info: brax.Info) -> jnp.ndarray:
      """Egocentric observation of target and arm body."""
      torso_fwd = math.rotate(jnp.array([1., 0., 0.]), qp.rot[self.torso_idx])
      torso_up = math.rotate(jnp.array([0., 0., 1.]), qp.rot[self.torso_idx])
  
      v_inv_rotate = jax.vmap(math.inv_rotate, in_axes=(0, None))
  
      pos_local = qp.pos - qp.pos[self.torso_idx]
      pos_local = v_inv_rotate(pos_local, qp.rot[self.torso_idx])
      vel_local = v_inv_rotate(qp.vel, qp.rot[self.torso_idx])
  
      target_local = pos_local[self.target_idx]
      target_local_mag = jnp.reshape(jnp.linalg.norm(target_local), -1)
      target_local_dir = target_local / (1e-6 + target_local_mag)
  
      pos_local = jnp.reshape(pos_local, -1)
      vel_local = jnp.reshape(vel_local, -1)
  
      contact_mag = jnp.sum(jnp.square(info.contact.vel), axis=-1)
      contacts = jnp.where(contact_mag > 0.00001, 1, 0)
  
      return jnp.concatenate([
          torso_fwd, torso_up, target_local_mag, target_local_dir, pos_local,
          vel_local, contacts
      ])
  
    def _random_target(self, rng: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
      """Returns a target location in a random circle slightly above xy plane."""
      rng, rng1, rng2 = jax.random.split(rng, 3)
      dist = self.target_radius + self.target_distance * jax.random.uniform(rng1)
      ang = jnp.pi * 2. * jax.random.uniform(rng2)
      target_x = dist * jnp.cos(ang)
      target_y = dist * jnp.sin(ang)
      target_z = .5
      target = jnp.array([target_x, target_y, target_z]).transpose()
      return rng, target
  
  _SYSTEM_CONFIG = """
  bodies {
    name: "floor"
    colliders {
      plane {
      }
    }
    inertia {
      x: 1.0
      y: 1.0
      z: 1.0
    }
    mass: 1.0
    frozen {
      position {
        x: 1.0
        y: 1.0
        z: 1.0
      }
      rotation {
        x: 1.0
        y: 1.0
        z: 1.0
      }
      all: true
    }
  }
  bodies {
    name: "shoulder_link"
    colliders {
      position {
        y: 0.06682991981506348
      }
      rotation {
        x: 90.0
      }
      capsule {
        radius: 0.05945208668708801
        length: 0.13365983963012695
      }
    }
    inertia {
      x: 1.0
      y: 1.0
      z: 1.0
    }
    mass: 1.0
  }
  bodies {
    name: "upper_arm_link"
    colliders {
      position {
        z: 0.21287038922309875
      }
      rotation {
      }
      capsule {
        radius: 0.05968618765473366
        length: 0.5446449518203735
      }
    }
    inertia {
      x: 1.0
      y: 1.0
      z: 1.0
    }
    mass: 1.0
  }
  bodies {
    name: "forearm_link"
    colliders {
      position {
        z: 0.1851803958415985
      }
      rotation {
      }
      capsule {
        radius: 0.05584339052438736
        length: 0.48926496505737305
      }
    }
    inertia {
      x: 1.0
      y: 1.0
      z: 1.0
    }
    mass: 1.0
  }
  bodies {
    name: "wrist_1_link"
    colliders {
      position {
        y: 0.10467606782913208
      }
      rotation {
        x: 90.0
      }
      capsule {
        radius: 0.038744933903217316
        length: 0.10467606782913208
      }
    }
    inertia {
      x: 1.0
      y: 1.0
      z: 1.0
    }
    mass: 1.0
  }
  bodies {
    name: "wrist_2_link"
    colliders {
      position {
        z: 0.052344050258398056
      }
      rotation {
      }
      capsule {
        radius: 0.03879201412200928
        length: 0.10468810051679611
      }
    }
    inertia {
      x: 1.0
      y: 1.0
      z: 1.0
    }
    mass: 1.0
  }
  bodies {
    name: "wrist_3_link"
    colliders {
      position {
        y: -0.04025782644748688
      }
      rotation {
        x: 90.0
      }
      capsule {
        radius: 0.01725015603005886
        length: 0.08051565289497375
      }
    }
    inertia {
      x: 1.0
      y: 1.0
      z: 1.0
    }
    mass: 1.0
  }
  bodies {
    name: "Target"
    colliders {
      sphere {
        radius: 0.1
      }
    }
    inertia {
      x: 1.0
      y: 1.0
      z: 1.0
    }
    mass: 1.0
    frozen {
      position {
        x: 1.0
        y: 1.0
        z: 1.0
      }
      rotation {
        x: 1.0
        y: 1.0
        z: 1.0
      }
      all: true
    }
  }
  joints {
    name: "shoulder_pan_joint"
    stiffness: 40000.0
    parent: "floor"
    child: "shoulder_link"
    parent_offset {
      z: 0.163
    }
    child_offset {
    }
    rotation {
      y: -90.0
    }
    angular_damping: 50.0
    angle_limit {
      min: -360.0
      max: 360.0
    }
  }
  joints {
    name: "shoulder_lift_joint"
    stiffness: 40000.0
    parent: "shoulder_link"
    child: "upper_arm_link"
    parent_offset {
      y: 0.138
    }
    child_offset {
    }
    rotation {
      z: 90.0
    }
    angular_damping: 50.0
    angle_limit {
      min: -360.0
      max: 360.0
    }
  }
  joints {
    name: "elbow_joint"
    stiffness: 40000.0
    parent: "upper_arm_link"
    child: "forearm_link"
    parent_offset {
      y: -0.13
      z: 0.425
    }
    child_offset {
    }
    rotation {
      z: 90.0
    }
    angular_damping: 50.0
    angle_limit {
      min: -360.0
      max: 360.0
    }
  }
  joints {
    name: "wrist_1_joint"
    stiffness: 40000.0
    parent: "forearm_link"
    child: "wrist_1_link"
    parent_offset {
      z: 0.3919999897480011
    }
    child_offset {
    }
    rotation {
      z: 90.0
    }
    angular_damping: 50.0
    angle_limit {
      min: -360.0
      max: 360.0
    }
  }
  joints {
    name: "wrist_2_joint"
    stiffness: 40000.0
    parent: "wrist_1_link"
    child: "wrist_2_link"
    parent_offset {
      y: 0.12700000405311584
    }
    child_offset {
    }
    rotation {
      y: -90.0
    }
    angular_damping: 50.0
    angle_limit {
      min: -360.0
      max: 360.0
    }
  }
  joints {
    name: "wrist_3_joint"
    stiffness: 40000.0
    parent: "wrist_2_link"
    child: "wrist_3_link"
    parent_offset {
      z: 0.1
    }
    child_offset {
    }
    rotation {
      z: 90.0
    }
    angular_damping: 50.0
    angle_limit {
      min: -360.0
      max: 360.0
    }
  }
  actuators {
    name: "shoulder_pan_joint"
    joint: "shoulder_pan_joint"
    strength: 100.0
    torque {
    }
  }
  actuators {
    name: "shoulder_lift_joint"
    joint: "shoulder_lift_joint"
    strength: 100.0
    torque {
    }
  }
  actuators {
    name: "elbow_joint"
    joint: "elbow_joint"
    strength: 100.0
    torque {
    }
  }
  actuators {
    name: "wrist_1_joint"
    joint: "wrist_1_joint"
    strength: 100.0
    torque {
    }
  }
  actuators {
    name: "wrist_2_joint"
    joint: "wrist_2_joint"
    strength: 100.0
    torque {
    }
  }
  actuators {
    name: "wrist_3_joint"
    joint: "wrist_3_joint"
    strength: 100.0
    torque {
    }
  }
  friction: 0.6
  gravity {
    z: -9.81
  }
  angular_damping: -0.05
  baumgarte_erp: 0.1
  dt: 0.02
  substeps: 8
  """
  