From b45701e448178e7a82ed609184eb69c0c0eaf0e6 Mon Sep 17 00:00:00 2001
From: Shambhuraj Sawant <shambhurajsawant10@gmail.com>
Date: Wed, 24 Jul 2024 11:29:15 +0200
Subject: [PATCH] faster dynamics for quad 2d

---
 .../envs/gym_pybullet_drones/base_aviary.py   | 30 +++++++++++--------
 .../envs/gym_pybullet_drones/quadrotor.py     |  9 ++++--
 2 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/safe_control_gym/envs/gym_pybullet_drones/base_aviary.py b/safe_control_gym/envs/gym_pybullet_drones/base_aviary.py
index 2c0636232..ddcaa89b0 100644
--- a/safe_control_gym/envs/gym_pybullet_drones/base_aviary.py
+++ b/safe_control_gym/envs/gym_pybullet_drones/base_aviary.py
@@ -271,7 +271,7 @@ def _advance_simulation(self, clipped_action, disturbance_force=None):
             # Between aggregate steps for certain types of update.
             if self.PYB_STEPS_PER_CTRL > 1 and self.PHYSICS in [
                 Physics.DYN, Physics.PYB_GND, Physics.PYB_DRAG,
-                Physics.PYB_DW, Physics.PYB_GND_DRAG_DW, Physics.RK4, Physics.DYN_2D
+                Physics.PYB_DW, Physics.PYB_GND_DRAG_DW, Physics.RK4 #, Physics.DYN_2D
             ]:
                 self._update_and_store_kinematic_information()
             # Step the simulation using the desired physics update.
@@ -314,8 +314,8 @@ def _advance_simulation(self, clipped_action, disturbance_force=None):
                 p.stepSimulation(physicsClientId=self.PYB_CLIENT)
             # Save the last applied action (e.g. to compute drag).
             self.last_clipped_action = clipped_action
-        # if self.PHYSICS == Physics.DYN_2D:
-        #     self._set_pybullet_information()
+        if self.PHYSICS == Physics.DYN_2D:
+            self._set_pybullet_information()
         # Update and store the drones kinematic information.
         self._update_and_store_kinematic_information()
 
@@ -715,16 +715,20 @@ def _dynamics_2d(self, rpm, nth_drone):
         ang_v = np.array([0, next_state[10], 0])
         ang_v = np.squeeze(ang_v)
 
-        # Set PyBullet's state.
-        p.resetBasePositionAndOrientation(self.DRONE_IDS[nth_drone],
-                                          pos,
-                                          p.getQuaternionFromEuler(rpy),
-                                          physicsClientId=self.PYB_CLIENT)
-        # Note: the base's velocity only stored and not used #
-        p.resetBaseVelocity(self.DRONE_IDS[nth_drone],
-                            vel,
-                            ang_v,  # ang_vel not computed by DYN
-                            physicsClientId=self.PYB_CLIENT)
+        # # Set PyBullet's state.
+        # p.resetBasePositionAndOrientation(self.DRONE_IDS[nth_drone],
+        #                                   pos,
+        #                                   p.getQuaternionFromEuler(rpy),
+        #                                   physicsClientId=self.PYB_CLIENT)
+        # # Note: the base's velocity only stored and not used #
+        # p.resetBaseVelocity(self.DRONE_IDS[nth_drone],
+        #                     vel,
+        #                     ang_v,  # ang_vel not computed by DYN
+        #                     physicsClientId=self.PYB_CLIENT)
+        self.pos[nth_drone, :] = pos.copy()
+        self.rpy[nth_drone, :] = rpy.copy()
+        self.vel[nth_drone, :] = vel.copy()
+        self.ang_v[nth_drone, :] = ang_v.copy()
         # Store the roll, pitch, yaw rates for the next step #
         self.rpy_rates[nth_drone, :] = X_dot[6:9]
 
diff --git a/safe_control_gym/envs/gym_pybullet_drones/quadrotor.py b/safe_control_gym/envs/gym_pybullet_drones/quadrotor.py
index f6f54bc48..1c30c184d 100644
--- a/safe_control_gym/envs/gym_pybullet_drones/quadrotor.py
+++ b/safe_control_gym/envs/gym_pybullet_drones/quadrotor.py
@@ -14,7 +14,7 @@
 
 from safe_control_gym.envs.benchmark_env import Cost, Task
 from safe_control_gym.envs.constraints import GENERAL_CONSTRAINTS
-from safe_control_gym.envs.gym_pybullet_drones.base_aviary import BaseAviary
+from safe_control_gym.envs.gym_pybullet_drones.base_aviary import BaseAviary, Physics
 from safe_control_gym.envs.gym_pybullet_drones.quadrotor_utils import (AttitudeControl, QuadType, cmd2pwm,
                                                                        pwm2rpm)
 from safe_control_gym.math_and_models.symbolic_systems import SymbolicModel
@@ -913,8 +913,11 @@ def _preprocess_control(self, action):
             # self.quat[0], np.array([0, pitch, 0])) # input thrust is in Newton
             # print(f"collective_thrust: {collective_thrust}, pitch: {pitch}")
 
-            _, quat = p.getBasePositionAndOrientation(self.DRONE_IDS[0], physicsClientId=self.PYB_CLIENT)
-            # quat = get_quaternion_from_euler(self.rpy[0, :])
+            if self.PHYSICS == Physics.DYN_2D:
+                # _, quat = p.getBasePositionAndOrientation(self.DRONE_IDS[0], physicsClientId=self.PYB_CLIENT)
+                quat = get_quaternion_from_euler(self.rpy[0, :])
+            else:
+                _, quat = p.getBasePositionAndOrientation(self.DRONE_IDS[0], physicsClientId=self.PYB_CLIENT)
             thrust_action = self.attitude_control._dslPIDAttitudeControl(collective_thrust / 4,
                                                                          quat, np.array([0, pitch, 0]))
             # input thrust is in Newton