Hello! 

I am working on this ML project, and I just cant figure it out. My agent just wont learn. I am using stable-baselines3 SAC, and I am training a 2 wheeled balancing robot. My neural network is 9 nodes, which are 

    linear_velocity * 0.2, roll, pitch * 10, roll_rate * 0.5, rot_rate * 0.2, wheel_velocity_1, wheel_velocity_2, target_velocity, target_rot_rate

where they are scaled by a normalization factor to be roughly -1 to 1. The last two are user input controls, which I am trying to train the agent to move forward or backward, or to turn left and right accordingly. My reward function can be seen on the second slide. As you can see from the following two slides, it is not getting any better at balancing, since each episode is terminated when the robot hits a certain tilt angle. You can also see that my robot is not even getting good at maximizing its reward. For this test, I have a random linear velocity set, a normal distribution centered around zero, with SD of 1m/s (which I am just realizing now could be the problem). When I train it with target\_velocity target\_rotation = 0, and just train it to balance starting from a random angle, its actually not bad. I have run about 4000 episodes, with a small input for target\_velocity and target\_rotation and it just does not seem to be working. I will post all my code below, just in case anyone would like to comb through it. Thank you for your help! This is for a capstone project coming up. 



  
Here is my pybullet environment: this handles the physics, and simulating my robot with friction, gravity, and setting motor controls and extracting state values:

    import pybullet as p
    import pybullet_data
    import numpy as np
    import csv
    MAXFORCE = 1.56
    
    
    def create_env():
        #Create physicsClient
        physicsClient = p.connect(p.GUI)
        #physicsClient = p.connect(p.DIRECT)
        p.setAdditionalSearchPath(pybullet_data.getDataPath())
    
    
        #Load ground plane
        planeId = p.loadURDF("plane.urdf")
        
        #Load robot from URDF, arbitrary start position, and standing starting orientation
        startPos = [0, 0, .1]
        startOrientation = p.getQuaternionFromEuler([0, 0, 0.7])
        boxId = p.loadURDF("C:\\Users\\dylan\\Documents\\Semester_9\\CAPSTONE\\Robot_RL\\Balro3.urdf", startPos, startOrientation)
        
        #Set friction and gravity
        p.changeDynamics(planeId, -1, lateralFriction=10.0)
        p.changeDynamics(boxId, 1, lateralFriction=3.0)
        p.changeDynamics(boxId, 0, lateralFriction=3.0)
        p.setGravity(0,0,-9.8) 
        
        return boxId
        
    def step_sim(velocity1, velocity2, boxId, target_velocity, target_rot_rate):
        #Set motor speeds and step simulation
        p.setJointMotorControl2(bodyUniqueId = boxId, jointIndex=0, controlMode=p.VELOCITY_CONTROL, targetVelocity = velocity1 * 12.775, force= MAXFORCE)
        p.setJointMotorControl2(bodyUniqueId = boxId, jointIndex=1, controlMode=p.VELOCITY_CONTROL, targetVelocity = velocity2 * 12.775, force= MAXFORCE)
        p.stepSimulation()
        
        #get velocity of robot (velocity of forward movement)
        linear_velocity = get_linear_velocity(boxId)
        
        #Collect orientation data from base link
        roll, pitch = get_roll_pitch(boxId)
        
        #Collect roll rate and rotation rate data
        roll_rate, rot_rate = get_pitch_rate(boxId)
        
        #Collect wheel velocity data from joints
        wheel_velocity_1 = p.getJointState(boxId, 0)[1] / 12.775
        wheel_velocity_2 = p.getJointState(boxId, 1)[1] / 12.775
        
        return linear_velocity * 0.2, roll, pitch * 10, roll_rate * 0.5, rot_rate * 0.2, wheel_velocity_1, wheel_velocity_2, target_velocity, target_rot_rate
    
    
    def get_linear_velocity(boxId):
        linear_vel_world, angular_vel_world = p.getBaseVelocity(boxId)
        pos, orn = p.getBasePositionAndOrientation(boxId)
        rot_matrix = np.array(p.getMatrixFromQuaternion(orn)).reshape(3, 3)
        
        x, y, z = rot_matrix.T @ np.array(linear_vel_world)
        
        return y
    
    
    
    def get_roll_pitch(boxId):
        _, orn = p.getBasePositionAndOrientation(boxId)
        rot_matrix = np.array(p.getMatrixFromQuaternion(orn)).reshape(3, 3)
    
    
        g_world = np.array([0, 0, -1])
        g_body = rot_matrix.T @ g_world
    
    
        # Use -g_body[2] so upright = 0
        roll = np.arctan2(g_body[1], -g_body[2])
        pitch = np.arctan2(-g_body[0], np.sqrt(g_body[1]**2 + g_body[2]**2))
    
    
        return roll, pitch
    
    
    
    def get_pitch_rate(robot_id):
        
        #Get angular velocity in world frame
        _, ang_vel_world = p.getBaseVelocity(robot_id)
        ang_vel_world = np.array(ang_vel_world)
    
    
        #Get orientation quaternion and rotation matrix
        _, orn = p.getBasePositionAndOrientation(robot_id)
        rot_matrix = np.array(p.getMatrixFromQuaternion(orn)).reshape(3, 3)
    
    
        #Transform angular velocity to body frame
        #Body-frame angular velocity = R^T * world-frame angular velocity
        ang_vel_body = rot_matrix.T @ ang_vel_world
    
    
        #Extract pitch rate and roll rate
        roll_rate = float(ang_vel_body[0])
        rot_rate  = float(ang_vel_body[2])
    
    
        return roll_rate, rot_rate
    
    
        
    
       
                
    def env_disconnect():
        p.disconnect()



Here is my Gymnasium environment: 

    import gymnasium as gym
    from gymnasium import spaces
    import numpy as np
    import pybullet as p
    import csv
    from pyb_env import create_env, step_sim, env_disconnect
    
    
    
    class GymEnv(gym.Env):
        def __init__(self):
            super().__init__()
            
            # Action space: 2 motor velocities (normalized between -1 and 1)
            self.action_space = spaces.Box(low=-1, high=1, shape=(2,), dtype=np.float32)
            
            # Observation space: roll, pitch, roll_rate, rot_rate, wheel_vel_1, wheel_vel_2, targ_vel, targ_rotation
            self.observation_space = spaces.Box(
        low=np.array([
            -5.0,   # linear_velocity (m/s)
            -1.6,   # roll (≈ -40°)
            -0.7,   # pitch (≈ -40°)
            -2.0,  # roll_rate (rad/s)
            -2.0,  # rot_rate (rad/s)
            -3.0,  # wheel_velocity_1 (rad/s)
            -3.0,  # wheel_velocity_2 (rad/s)
            -5.0,   # target_velocity
            -2.0    # target_rot_rate
        ], dtype=np.float32),
        high=np.array([
            5.0, 1.6, 0.7, 2.0, 2.0, 3.0, 3.0, 5.0, 2.0
        ], dtype=np.float32),
        dtype=np.float32
    )
            
            self.boxId = None
            self.step_count = 0
            rand = np.random.uniform(low = -100, high = 100)
            self.max_steps = 1000 + rand
            self.episode_reward = 0
            self.target_velocity = np.random.uniform(low = -.1, high = .1)
    
    
        def reset(self, *, seed=None, options=None):
            super().reset(seed=seed)
            if self.boxId is None: 
                self.boxId = create_env()
            angle = np.random.normal(loc=0.055, scale=0.055, size=1)
            is_negative = np.random.choice([True, False])
            if is_negative:
                angle *= -1
                
            p.resetBasePositionAndOrientation(self.boxId, [0, 0, .1], p.getQuaternionFromEuler([angle, 0, 0.7]) )
            p.resetBaseVelocity(self.boxId, [0,0,0], [0,0,0])
            if self.step_count != 0:
                log_data("Tests/rot_vel_test1.csv", self.step_count)
                print(self.step_count)
            if self.step_count != 0:
                log_data("Tests/rot_vel_test1_rewards.csv", self.episode_reward)
            #log_data("0_degree_rewards.csv", self.episode_reward)
            self.step_count = 0
            self.episode_reward = 0
            self.ctrl_set = False
            rand = np.random.normal(loc = 0, scale = 5, size = 1)
            self.max_steps = 5000 + rand
            
            self.target_velocity = 0
            self.target_rot = 0
            state = np.array(step_sim(0, 0, self.boxId, self.target_velocity, self.target_rot), dtype=np.float32)
            
            
            info = {}
            
            return state, info
    
    
        def step(self, action):
            # Denormalize actions if needed
            velocity1 = float(action[0])
            velocity2 = float(action[1])
    
    
            if self.step_count > 200 and self.ctrl_set == False:
                self.target_velocity = np.random.normal(loc=0, scale=0.2)
                self.target_rot = np.random.normal(loc=0, scale=0.02)
                self.ctrl_set = True
            state = np.array(step_sim(velocity1, velocity2, self.boxId, self.target_velocity, self.target_rot), dtype=np.float32)
    
    
            reward = compute_reward(state)
            self.episode_reward += reward
            terminated = abs(state[1]) > 0.35 or abs(state[2]) > 0.2
            
                
            truncated = False
            self.step_count += 1
            if self.step_count >= self.max_steps:
                truncated = True
            info = {}
            return state, reward, terminated, truncated, info
        
    
    
    
        def close(self):
            env_disconnect()
            
    
    
    def compute_reward(state):
        # state = linear_velocity, roll, pitch, roll_rate, rot_rate, wheel_velocity_1, wheel_velocity_2, target_velocity, target_rot_rate
        roll = state[1]
        pitch = state[2]
        target_velocity = state[7]
        target_rotation_rate = state[8]
        rot_rate = state[4]
        linear_velocity = state[0]
        
        upright_bonus = 1 - (abs(roll)/0.12) - abs(pitch)   # close to 1 when upright
        velocity_compliance = 1 - (abs(target_velocity - linear_velocity)/0.05)
        rotation_compliance = 1 - (abs(target_rotation_rate - rot_rate)/0.4)
        
        
        
        return np.clip(reward, -10, 10)

And you can see how I call my agent for training on the last slide. Thanks again if you made it this far



Why Wont My Machine Learn! Please Help!