From 55eecdc8d7ade931b5f5dd336b2221ab02a5d7ef Mon Sep 17 00:00:00 2001
From: Rune Harlyk <rune@harlyk.dk>
Date: Tue, 20 Aug 2024 21:41:35 +0200
Subject: [PATCH] =?UTF-8?q?=F0=9F=9B=B9=20Adds=20static=20gui=20to=20env?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 simulation/simulation/environment.py |  3 ++
 simulation/simulation/gui.py         | 71 ++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+)
 create mode 100644 simulation/simulation/gui.py

diff --git a/simulation/simulation/environment.py b/simulation/simulation/environment.py
index 7927558..10357d6 100644
--- a/simulation/simulation/environment.py
+++ b/simulation/simulation/environment.py
@@ -2,6 +2,7 @@ import pybullet as p
 import pybullet_data
 import numpy as np
 from simulation.robot import QuadrupedRobot
+from simulation.gui import GUI
 
 
 roll_pitch_reward_weight = 0.1
@@ -16,6 +17,7 @@ class QuadrupedEnv:
         self.urdf_path = urdf_path
 
         self.setupWorld()
+        self.gui = GUI(self.robot.robot_id)
 
         self.envStartState = p.saveState()
 
@@ -32,6 +34,7 @@ class QuadrupedEnv:
         return self.robot.get_observation()
 
     def step(self, action):
+        self.gui.update()
         self.robot.apply_action(action)
         p.stepSimulation()
         obs = self.robot.get_observation()
diff --git a/simulation/simulation/gui.py b/simulation/simulation/gui.py
new file mode 100644
index 0000000..c7c593c
--- /dev/null
+++ b/simulation/simulation/gui.py
@@ -0,0 +1,71 @@
+import pybullet as pb
+import numpy as np
+import sys
+
+
+class GUI:
+    def __init__(self, quadruped):
+        self.cameraYaw = 0
+        self.cameraPitch = -7
+        self.cameraDist = 0.66
+
+        self.rollId = pb.addUserDebugParameter("roll", -np.pi / 4, np.pi / 4, 0.0)
+        self.pitchId = pb.addUserDebugParameter("pitch", -np.pi / 4, np.pi / 4, 0.0)
+        self.yawId = pb.addUserDebugParameter("yaw", -np.pi / 4, np.pi / 4, 0.0)
+        self.xId = pb.addUserDebugParameter("x", -0.10, 0.10, 0.0)
+        self.yId = pb.addUserDebugParameter("y", -0.10, 0.10, 0.0)
+        self.zId = pb.addUserDebugParameter("z", -0.055, 0.17, 0.0)
+
+        self.quadruped = quadruped
+
+    def updateCamera(self):
+        quadruped_position, _ = pb.getBasePositionAndOrientation(self.quadruped)
+        pb.resetDebugVisualizerCamera(
+            cameraDistance=self.cameraDist,
+            cameraYaw=self.cameraYaw,
+            cameraPitch=self.cameraPitch,
+            cameraTargetPosition=quadruped_position,
+        )
+
+    def handleInput(self):
+        keys = pb.getKeyboardEvents()
+        # Keys to change camera
+        if keys.get(100):  # D
+            self.cameraYaw += 1
+        if keys.get(97):  # A
+            self.cameraYaw -= 1
+        if keys.get(99):  # C
+            self.cameraPitch += 1
+        if keys.get(102):  # F
+            self.cameraPitch -= 1
+        if keys.get(122):  # Z
+            self.cameraDist += 0.01
+        if keys.get(120):  # X
+            self.cameraDist -= 0.01
+        if keys.get(27):  # ESC
+            pb.disconnect()
+            sys.exit()
+
+    def update(self):
+        self.updateCamera()
+        self.handleInput()
+
+        position = np.array(
+            [
+                pb.readUserDebugParameter(self.xId),
+                pb.readUserDebugParameter(self.yId),
+                pb.readUserDebugParameter(self.zId),
+            ]
+        )
+        orientation = np.array(
+            [
+                pb.readUserDebugParameter(self.rollId),
+                pb.readUserDebugParameter(self.pitchId),
+                pb.readUserDebugParameter(self.yawId),
+            ]
+        )
+
+        return (
+            position,
+            orientation,
+        )