Model bounces when calling the "reset_world" service in gazebo

edit

Hi, I am trying to reset a model to it's original position but when I call the "/reset_world" service (without a large time delay with time.sleep) the model bounces in the air. Any help to prevent this bounce/ wobble would be appreciated as many resets will be needed in the simulation (after every episode for Q-Learning) so it would save a lot of time. (Below is the code for controlling the model and resetting the simulation as well as the URDF)

for episode in range(5000):

        episode_rewards = 0
        rclpy.spin_once(self.imu_sub)
        self.observation = self.imu_sub.get_latest_observation()
        while(not self.is_episode_finished()):
            current_state = self.discretise_observation()
            action = self.choose_action(current_state)

            self.cmd_vel_pub.publish(action) #move_robot()
            rclpy.spin_once(self.imu_sub)
            self.observation = self.imu_sub.get_latest_observation()

            new_state = self.discretise_observation()
            reward = self.get_reward(new_state[0])
            episode_rewards += reward
            lr = self.get_learning_rate(episode+1)
            self.Q_table[current_state][action] = self.Q_table[current_state][action] + lr * (reward + self.DISCOUNT_FACTOR * np.max(self.Q_table[new_state]) - self.Q_table[current_state][action])


        self.get_logger().info(f"{episode}| Episode rewards: {episode_rewards}")
        self.epsilon = max(0.01, self.epsilon * self.epsilon_decay)
        self.all_rewards.append(episode_rewards)
        self.resetSim.send_request()

class ResetSim(Node):
def __init__(self):
    super().__init__("minimum_reset_client")
    self.client = self.create_client(Empty, "/reset_world")
    while not self.client.wait_for_service():
        self.get_logger().info("Service not available, waiting again...")
    self.req = Empty.Request()

def send_request(self):
    self.future = self.client.call_async(self.req)
    rclpy.spin_until_future_complete(self, self.future)
    return self.future.result()




<robot name="m2wr" xmlns:xacro="https://www.ros.org/wiki/xacro" >

<link name="link_chassis">
    <!-- pose and inertial -->
    <pose>0 0 0 0 0 0</pose>

    <inertial>
    <mass value="5"/>
    <origin rpy="0 0 0" xyz="0 0 0.1"/>
    <inertia ixx="0.0395416666667" ixy="0" ixz="0" iyy="0.106208333333" iyz="0" izz="0.106208333333"/>
    </inertial>

    <collision name="collision_chassis">
    <geometry>
        <box size="0.5 0.5 0.8"/>
    </geometry>
    </collision>

    <visual>
    <origin rpy="0 0 0" xyz="0 0 0"/>
    <geometry>
        <box size="0.5 0.5 0.8"/>
    </geometry>
    <material name="blue"/>
    </visual>
    </link>

<!-- Create wheel right -->  

<link name="link_right_wheel">    
    <inertial>
    <mass value="0.2"/>
    <origin rpy="0 1.5707 1.5707" xyz="0 0 0"/>
    <inertia ixx="0.00052666666" ixy="0" ixz="0" iyy="0.00052666666" iyz="0" izz="0.001"/>
    </inertial>

    <collision name="link_right_wheel_collision">
    <origin rpy="0 1.5707 1.5707" xyz="0 0 0" />
    <geometry>
        <cylinder length="0.1" radius="0.2"/>
    </geometry>
    </collision>

    <visual name="link_right_wheel_visual">
    <origin rpy="0 1.5707 1.5707" xyz="0 0 0"/>
    <geometry>
        <cylinder length="0.1" radius="0.2"/>
    </geometry>      
    </visual>

</link>

<!--   Joint for right wheel  -->
<joint name="joint_right_wheel" type="continuous">
    <origin rpy="0 0 0" xyz="0 0.3 -0.4"/>
    <child link="link_right_wheel" />
    <parent link="link_chassis"/>
    <axis rpy="0 0 0" xyz="0 -1 0"/>
    <limit effort="10000" velocity="1000 ...