Upload 2 files

learning_machines_robobo/examples/full_project_setup/catkin_ws/src/learning_machines/src/learning_machines/task2_robobo_env_t3rework.py

@@ -0,0 +1,131 @@
+import os
+import cv2
+import numpy as np
+import gymnasium as gym
+from gymnasium import spaces
+from stable_baselines3 import PPO
+from stable_baselines3.common.vec_env import DummyVecEnv
+from stable_baselines3.common.callbacks import EvalCallback, CheckpointCallback, EveryNTimesteps, CallbackList
+from datetime import datetime
+from robobo_interface import IRobobo, Position, Orientation
+from data_files import FIGRURES_DIR
+
+class RoboboEnv(gym.Env):
+    def __init__(self, rob: IRobobo, ep_steps=1024, simulation=True):
+        super(RoboboEnv, self).__init__()
+        self.robot = rob
+        self.simulation = simulation
+        self.ep_steps = ep_steps
+        
+        # Define action and observation space
+        self.action_space = spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32)
+        self.observation_space = spaces.Box(low=0, high=255, shape=(64, 64, 3), dtype=np.uint8)
+        
+        self.has_reset = False
+        self.current_step = 0
+
+    def reset(self, seed=None, options=None):
+        super().reset(seed=seed)  # Reset parent class
+        if self.simulation:
+            self.robot.stop_simulation()
+            self.robot.set_position(Position(0, 0, 0), Orientation(0, 0, 0))
+            self.robot.play_simulation()
+        
+        self.robot.set_phone_tilt(240, 50)  # tilted down
+        self.robot.sleep(3)
+        self.current_step = 0
+        self.has_reset = True
+        return self._get_obs(), {}
+
+    def step(self, action):
+        if not self.has_reset:
+            raise RuntimeError("Environment must be reset before stepping")
+        
+        self.current_step += 1
+        
+        action[2] = (((action[2] + 1) / 2) * 0.75) + 0.15
+        left_speed, right_speed, duration_scale = action * np.array([50, 50, 800])
+        
+        self.robot.move(left_speed, right_speed, duration_scale)
+        self.robot.sleep(action[2])
+        
+        obs = self._get_obs()
+        reward = self._calculate_reward(obs)
+        print(self.current_step,reward)
+        
+        done = self.current_step >= self.ep_steps
+        
+        return obs, reward, done, False, {}
+
+    def _get_obs(self):
+        frame = self.robot.get_image_front()
+        frame = cv2.resize(frame, (64, 64))
+        return frame
+    
+
+    def _calculate_reward(self, obs):
+        food_collected = self.robot.nr_food_collected()
+        
+        # Convert to HSV for easier color detection
+        hsv = cv2.cvtColor(obs, cv2.COLOR_BGR2HSV)
+        
+        # Detect green area (representing food)
+        lower_green = np.array([40, 50, 50])
+        upper_green = np.array([80, 255, 255])
+        green_mask = cv2.inRange(hsv, lower_green, upper_green)
+        green_area = np.sum(green_mask > 0)
+        
+        # Calculate centroid of the green area
+        green_centroid = np.mean(np.argwhere(green_mask > 0), axis=0) if green_area > 0 else np.array([32, 32])
+        
+        # Distance to the green centroid
+        distance_to_green = np.linalg.norm(np.array([32, 32]) - green_centroid)
+        
+        # Reward for collecting food
+        reward = food_collected * 10  # Reward for collecting food
+        
+        # Additional reward for getting closer to the green area
+        if green_area > 0:
+            reward += 5 / (distance_to_green + 1)  # Inverse of the distance to the green area for reward
+        
+        # Less reward if no food detected by base
+        if not self.robot.base_detects_food():
+            reward *= 0.5
+        
+        return reward
+
+def run(rob: IRobobo):
+    # load_model = False
+    env = RoboboEnv(rob)
+    env = DummyVecEnv([lambda: env])
+    
+    current_datetime = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    log_dir = os.path.join(FIGRURES_DIR, "models", f"PPO_{current_datetime}")
+    os.makedirs(log_dir, exist_ok=True)
+    
+    steps = 1024
+    episodes = 1000
+    total_timesteps = steps * episodes
+    
+    checkpoint_callback = CheckpointCallback(save_freq=total_timesteps // 1000, save_path=log_dir, name_prefix="ppo_model")
+    event_callback = EveryNTimesteps(n_steps=total_timesteps // 1000, callback=checkpoint_callback)
+    
+    model = PPO("CnnPolicy", env, verbose=1, tensorboard_log=log_dir,
+                learning_rate=1e-4, 
+                n_steps=512, 
+                batch_size=128, 
+                n_epochs=5, 
+                gamma=0.99, 
+                ent_coef=0.01, 
+                use_sde=True, 
+                sde_sample_freq=4, 
+                target_kl=0.015)
+    
+    eval_callback = EvalCallback(env, best_model_save_path=log_dir,
+                                 log_path=log_dir, eval_freq=steps * 11,
+                                 deterministic=True, render=False)
+    
+    callbacks = CallbackList([checkpoint_callback, event_callback, eval_callback])
+    model.learn(total_timesteps=total_timesteps, callback=callbacks)
+    model.save(os.path.join(log_dir, "final_model"))
+

learning_machines_robobo/examples/full_project_setup/catkin_ws/src/learning_machines/src/learning_machines/task3_rob_env_irs.py

@@ -0,0 +1,695 @@
+import os
+import cv2
+import numpy as np
+import gymnasium as gym
+from gymnasium import spaces
+from stable_baselines3 import PPO
+from stable_baselines3.common.vec_env import DummyVecEnv
+from stable_baselines3.common.callbacks import EvalCallback, CheckpointCallback, EveryNTimesteps, CallbackList
+from datetime import datetime
+from robobo_interface import IRobobo, Position, Orientation
+from data_files import FIGRURES_DIR
+
+stream = True
+
+class RoboboEnv(gym.Env):
+    def __init__(self, rob: IRobobo, ep_steps=512, simulation=True):
+        super(RoboboEnv, self).__init__()
+        self.robot = rob
+        self.simulation = simulation
+        self.ep_steps = ep_steps
+        
+        # Define action and observation space
+        self.action_space = spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32)
+        
+        # Modified observation space with only specified IR sensors
+        self.observation_space = spaces.Dict({
+            'image_segments': spaces.Box(low=0, high=2, shape=(3,), dtype=np.uint8),
+            'ir_sensors': spaces.Box(low=0, high=1, shape=(3,), dtype=np.float32)  # Normalized IR sensors
+        })
+        
+        self.has_reset = False
+        self.current_step = 0
+        self.episode_reward = 0
+        self.best_completion_steps = float('inf')
+        self.consecutive_red_sightings = 0
+        self.previous_action = np.zeros(3)
+        
+        # Curriculum variables
+        self.curriculum_stage = 4
+        # Updated curriculum stages
+        # Updated curriculum stages
+        # Reordered curriculum stages
+        self.stages = [
+            "find_red_box",
+            "smooth_movement",
+            "push_red_box",
+            "find_green_box",
+            "push_to_green_box"
+        ]
+        self.lost_red_box = False
+        
+        # Adjust thresholds to match the new order (you may need to fine-tune these)
+        self.stage_threshold_rewards = [
+            4096,   # Stage 0: find_red_box
+            3072,   # Stage 1: smooth_movement
+            5120,   # Stage 2: push_red_box
+            6144,   # Stage 3: find_green_box
+            10640   # Stage 4: push_to_green_box
+        ]
+        
+        self.stage_success_count = 0
+        self.recent_episode_rewards = []
+        
+        self.consecutive_wall_detections = 0
+        self.max_wall_penalty = 10000000 
+        
+        self.visited_places = set()
+        self.grid_size = 0.1  # Size of each grid cell in meters
+        self.exploration_reward_scale = 10  # Scale factor for exploration reward
+        
+        # Time tracking for final stage
+        self.start_time = None
+        self.best_completion_time = float('inf')
+        
+        if stream:
+            cv2.namedWindow("Robot View", cv2.WINDOW_NORMAL)
+            
+        # Initialize previous_front_distance
+        self.previous_front_distance = None
+
+    def reset(self, seed=None, options=None):
+        super().reset(seed=seed)
+        if self.simulation:
+            self.robot.stop_simulation()
+            self.robot.set_position(Position(-2.4000000000000026, 0.07699995934963236, 0.03970504179596901), Orientation(-1.57191039448275, -1.5144899542893442, -1.5719101704888712))
+            self.robot.play_simulation()
+        
+        self.robot.set_phone_tilt(240, 50)  # tilted down
+        self.robot.sleep(2)  
+        self.current_step = 0
+        self.episode_reward = 0
+        self.consecutive_red_sightings = 0
+        self.has_reset = True
+        self.previous_action = np.zeros(3)
+        self.consecutive_wall_detections = 0
+        
+        self.visited_places = set()
+        
+        
+        
+        self.previous_front_distance = None
+        
+        if self.curriculum_stage == len(self.stages) - 1:
+            self.start_time = datetime.now()
+        
+        self.action_history = []
+        return self._get_obs(), {}
+
+    def step(self, action):
+        if not self.has_reset:
+            raise RuntimeError("Environment must be reset before stepping")
+        
+        self.current_step += 1
+        self.action_history.append(action)
+        if len(self.action_history) > 5:
+            self.action_history.pop(0)
+        
+        # Adjust action scaling for faster movements
+        action[2] = ((action[2]+1)/8)+0.1  # Reduced upper limit of duration
+        
+        
+        # Convert normalized actions to robot commands
+        left_speed, right_speed, duration_scale = action * np.array([60, 60, 1000])  # Increased speed range, reduced duration
+        
+        self.robot.move(left_speed, right_speed, duration_scale)
+        self.robot.sleep(action[2]/4)
+        #print actions and sleep
+        # print("----------------------")
+        # print("Left Speed:", left_speed, "Right Speed:", right_speed, "Duration Scale:", duration_scale, "Action[2]:", action[2])
+        # print("----------------------")
+        
+        obs = self._get_obs()
+        # print("Observation:")
+        # print(f"  Image Segments: {obs['image_segments']}")
+        # print(f"  IR Sensors: {obs['ir_sensors']}")
+        
+        if stream:
+            self._stream_frame(obs['image_segments'])
+            
+        # exploration_reward = self._calculate_exploration_reward()
+        # reward += exploration_reward
+        
+        reward = self._calculate_reward(obs, action)
+        self.episode_reward += reward
+        print(f"Step: {self.current_step}, Reward: {reward}, Stage: {self.stages[self.curriculum_stage]}")
+        
+        done = self._check_done(obs)
+        
+        if done:
+            self._check_curriculum_progress()
+        
+        self.previous_action = action
+        
+        return obs, reward, done, False, {}
+
+    def _get_obs(self):
+        frame = self.robot.get_image_front()
+        height, width = frame.shape[:2]
+        crop_height = int(height * 0.6)
+        frame = frame[height - crop_height:height, :]
+        frame = cv2.resize(frame, (64, 64))
+        
+        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+        
+        lower_red = np.array([0, 50, 50])
+        upper_red = np.array([10, 255, 255])
+        lower_green = np.array([40, 50, 50])
+        upper_green = np.array([80, 255, 255])
+        
+        red_mask = cv2.inRange(hsv, lower_red, upper_red)
+        green_mask = cv2.inRange(hsv, lower_green, upper_green)
+        
+        segment_width = 21  # 64 // 3
+        image_segments = np.zeros(3, dtype=np.uint8)
+        
+        for i in range(3):
+            start = i * segment_width
+            end = (i + 1) * segment_width
+            red_area = np.sum(red_mask[:, start:end] > 0)
+            green_area = np.sum(green_mask[:, start:end] > 0)
+            
+            if red_area > 100 and green_area > 100:
+                image_segments[i] = 3
+            elif red_area > 100:
+                image_segments[i] = 1
+            elif green_area > 100:
+                image_segments[i] = 2
+        
+        all_ir_sensors = self.robot.read_irs()
+        ir_sensors = np.array([all_ir_sensors[4], all_ir_sensors[5], all_ir_sensors[7]], dtype=np.float32) / 1000.0
+        
+        return {
+            'image_segments': image_segments,
+            'ir_sensors': ir_sensors
+        }
+        
+    def _stream_frame(self, image_segments):
+        # Create a visual representation of the image segments
+        frame = np.zeros((64, 64, 3), dtype=np.uint8)
+        segment_width = frame.shape[1] // 3
+        for i, color in enumerate(image_segments):
+            if color == 1:  # Red
+                frame[:, i*segment_width:(i+1)*segment_width] = [0, 0, 255]
+            elif color == 2:  # Green
+                frame[:, i*segment_width:(i+1)*segment_width] = [0, 255, 0]
+            elif color == 3:  # Both red and green
+                frame[:, i*segment_width:(i+1)*segment_width] = [0, 255, 255]  # Yellow to represent both
+        cv2.imshow("Robot View", cv2.resize(frame, (320, 320)))
+        cv2.waitKey(1)
+        
+    def close(self):
+        if stream:
+            cv2.destroyAllWindows()
+            
+    def get_stage_threshold(self):
+        return self.stage_threshold_rewards[self.curriculum_stage]
+    
+    def _calculate_exploration_reward(self):
+        # Get current position
+        position = self.robot.get_position()
+        x, y = position.x, position.y
+
+        # Discretize the position to grid coordinates
+        grid_x = int(x / self.grid_size)
+        grid_y = int(y / self.grid_size)
+
+        # Check if this grid cell has been visited before
+        grid_pos = (grid_x, grid_y)
+        if grid_pos not in self.visited_places:
+            self.visited_places.add(grid_pos)
+            return 10*self.exploration_reward_scale**2
+        return -100
+
+    
+    
+    def _calculate_reward(self, obs, action):
+        image_segments = obs['image_segments']
+        ir_sensors = obs['ir_sensors']
+        
+        reward = 0
+        
+        # if wall_detected: artifact for training till 30712
+        #         reward -= 10
+        
+        # Common penalties for all stages
+        movement_penalty = np.sum(np.abs(action - self.previous_action))
+        reward -= movement_penalty * 0.75  # Reduced penalty for unnecessary movement
+
+        if action[0] == action[1]:
+            reward += 100
+        else:
+            reward -= 50
+        
+        # Penalty for spinning
+        if np.abs(action[0] - action[1]) > 0.3:
+            reward -= 1001 # Reduced penalty for spinning
+        
+        # New penalty for fast turns
+        turn_speed = abs(action[0] - action[1])  # Difference between left and right wheel speeds
+        fast_turn_threshold = 0.3  # Adjust this value as needed
+        if turn_speed > fast_turn_threshold:
+            fast_turn_penalty = -200 * (turn_speed - fast_turn_threshold)**2  # Quadratic penalty
+            reward += fast_turn_penalty
+        
+        red_visible = 1 in image_segments
+        red_centered = image_segments[1] == 1
+        
+        # Improved object detection
+        object_close = ir_sensors[0] > 0.1 and ir_sensors[1] < 0.01 and ir_sensors[2] < 0.01
+        
+        # New wall avoidance logic with controlled penalties
+        front_distance = ir_sensors[0]
+        left_distance = ir_sensors[1]
+        right_distance = ir_sensors[2]
+        
+        wall_penalty = 0
+        
+        # Quadratic penalty for getting close to walls
+        wall_threshold = 0.1
+        for distance in [left_distance, right_distance]:
+            if distance > wall_threshold:
+                wall_penalty -= min(1000, 100 * ((distance - wall_threshold) ** 2))
+        
+        # Additional penalty for being too close to walls on multiple sides
+        close_walls = sum(1 for d in [front_distance, left_distance, right_distance] if d > wall_threshold)
+        if close_walls > 1:
+            wall_penalty -= 100 * close_walls
+        
+        # Encourage turning away from walls
+        if front_distance > 0.08 or max(left_distance, right_distance) > 0.08:
+            turn_encouragement = 50 * abs(action[0] - action[1])
+            wall_penalty += turn_encouragement
+        
+        # Reward for moving away from walls
+        if self.previous_front_distance is not None:
+            if front_distance < self.previous_front_distance:
+                wall_penalty += 20
+        
+        self.previous_front_distance = front_distance
+        
+        # Reward for safe navigation in open space
+        if all(sensor < 0.05 for sensor in ir_sensors):
+            reward += 10
+        
+        # Stuck detection with increased penalty
+        if np.all(np.abs(action) < 0.1):  # If the action is very small in all dimensions
+            self.stuck_count += 1
+            if self.stuck_count > 5:  # If stuck for more than 5 timesteps
+                reward -= 100  # Increased penalty for being stuck
+        else:
+            self.stuck_count = 0
+        
+        # Red box detection rewards (scaled down compared to wall penalties)
+        if red_visible:
+            self.consecutive_red_sightings += 1
+            reward += min(12000, self.consecutive_red_sightings**1.1)  # Cap the reward
+            if red_centered:
+                reward += 1
+            if object_close:
+                reward += 3
+        else:
+            self.consecutive_red_sightings = 0
+            reward -= 10
+            
+        # Check for red box inside green area at every stage
+        if self._is_red_inside_green(obs):
+            steps_factor = max(0, 1 - (self.current_step / self.ep_steps))
+            omega_reward = 100000 * (1 + steps_factor)  # Reduced large reward for completing the task
+            reward += omega_reward
+            print(f"Red box inside green area! Omega reward: {omega_reward}")
+            return reward  # End the episode immediately with a large reward
+
+        # Stage-specific rewards
+        if self.curriculum_stage == 0:  # find_red_box stage
+            if red_visible:
+                reward += self.consecutive_red_sightings
+                if red_centered:
+                    reward += 1
+                if object_close:
+                    reward += 2
+            else:
+                reward -= 0.5
+            
+        elif self.curriculum_stage == 1:  # smooth_movement stage
+            # Reward for moving straight
+            if np.abs(action[0] - action[1]) < 0.1:
+                reward += 1
+            
+            # Penalty for frequent direction changes
+            if len(self.action_history) >= 3:
+                last_three_actions = self.action_history[-3:]
+                if all(np.abs(a[0] - a[1]) > 0.2 for a in last_three_actions):
+                    reward -= 2  # Reduced penalty for 3 twists/turns in a row
+            
+            # # Reward for detecting and approaching red box (ARTIFACTR2)
+            # if red_visible:
+            #     reward += 30
+            #     if red_centered:
+            #         reward += 20
+            #     if object_close:
+            #         reward += 50
+            # Keep the existing red box logic, but scale down rewards
+        
+        elif self.curriculum_stage == 2:  # push_red_box stage
+            if red_visible and object_close:
+                reward += 5
+                if red_centered:
+                    reward += 2
+            elif red_visible or object_close:
+                reward += 2
+            else:
+                reward -= 1
+
+        elif self.curriculum_stage == 3:  # find_green_box stage
+            red_reward = self._get_color_reward(obs, 1)
+            green_reward = self._get_color_reward(obs, 2)
+            if red_reward > 0 and green_reward > 0:
+                reward += (red_reward + green_reward) * 1.5
+            elif green_reward > 0:
+                reward += green_reward
+            elif red_reward > 0:
+                reward += red_reward * 0.5
+            else:
+                reward -= 5
+
+        elif self.curriculum_stage == 4:  # push_to_green_box stage
+            red_reward = self._get_color_reward(obs, 1)
+            green_reward = self._get_color_reward(obs, 2)
+            both_colors_reward = 10 if 3 in obs['image_segments'] else 0
+            
+            if red_reward > 0 and green_reward > 0:
+                red_positions = np.where(obs['image_segments'] == 1)[0]
+                green_positions = np.where(obs['image_segments'] == 2)[0]
+                if len(red_positions) > 0 and len(green_positions) > 0:
+                    red_index = red_positions[0]
+                    green_index = green_positions[0]
+                    distance = abs(red_index - green_index)
+                    proximity_reward = 25 / (distance + 1)
+                    reward += red_reward + green_reward + proximity_reward + both_colors_reward
+                else:
+                    reward += red_reward + green_reward + both_colors_reward
+            elif green_reward > 0:
+                reward += green_reward
+            elif red_reward > 0:
+                reward += red_reward * 0.5
+            else:
+                reward -= 5
+        
+        reward += wall_penalty
+        # Encourage forward movement for all stages
+        # if action[0] > 0.1 and action[1] > 0.1: # if action[0] > 0.25 ARTIFACT 5
+        #     reward += 5
+        # Add exploration reward
+        exploration_reward = self._calculate_exploration_reward()
+        reward += exploration_reward
+        
+        # Apply omega penalty if episode is ending due to losing sight of red box
+        if self.lost_red_box:
+            omega_penalty = -100000  # Large negative reward
+            print(f"Lost sight of red box! Omega penalty: {omega_penalty}")
+            reward += omega_penalty
+            self.lost_red_box = False  # Reset the flag
+
+        return reward
+    
+    def _is_red_inside_green(self, obs):
+        image_segments = obs['image_segments']
+        return 3 in image_segments and image_segments[1] == 3 and obs['ir_sensors'][0] > 0.1
+
+    def _check_done(self, obs):
+        if self.current_step >= self.ep_steps:
+            return True
+        
+        if self._is_red_inside_green(obs):
+            return True
+        
+        if self.consecutive_red_sightings == 0 and self.current_step > 100:
+            self.lost_red_box = True
+            return True
+        return False
+
+    def _get_color_reward(self, obs, color, multiplier=1):
+        if color in obs['image_segments']:
+            base_reward = 2 if color == 1 else 1  # Higher base reward for red
+            if obs['image_segments'][1] == color:  # Check if color is in center segment
+                return base_reward * 3 * multiplier  # Triple reward for center
+            else:
+                return base_reward * multiplier
+        return 0
+
+    def _check_curriculum_progress(self):
+        self.recent_episode_rewards.append(self.episode_reward)
+        if len(self.recent_episode_rewards) > 3:
+            self.recent_episode_rewards.pop(0)
+        
+        if len(self.recent_episode_rewards) == 3:
+            avg_reward = sum(self.recent_episode_rewards) / 3
+            if avg_reward >= self.stage_threshold_rewards[self.curriculum_stage]:
+                self.stage_success_count += 1
+            else:
+                self.stage_success_count = 0
+            
+            if self.stage_success_count >= 1:
+                self.increase_difficulty()
+                self.stage_success_count = 0
+                self.recent_episode_rewards = []
+        
+        if self.curriculum_stage == len(self.stages) - 1:
+            if self._is_red_inside_green(self._get_obs()):
+                if self.current_step < self.best_completion_steps:
+                    self.best_completion_steps = self.current_step
+                    print(f"New best completion in {self.best_completion_steps} steps!")
+
+    def increase_difficulty(self):
+        if self.curriculum_stage < len(self.stages) - 1:
+            self.curriculum_stage += 1
+            print(f"Progressing to stage: {self.stages[self.curriculum_stage]}")
+        else:
+            print("Maximum stage reached")
+
+    def get_current_stage(self):
+        return self.stages[self.curriculum_stage]
+
+def run(rob: IRobobo):
+    train = True
+    
+    load_model = True  # Set to True if you want to load a pre-trained model
+    steps = 512  # Reduced number of steps per episode
+    episodes = 2000  # Increased number of episodes
+    total_timesteps = steps * episodes
+
+    # Common setup
+    env = RoboboEnv(rob)
+    env = DummyVecEnv([lambda: env])
+    
+    current_datetime = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    log_dir = os.path.join(FIGRURES_DIR, "models", f"PPO_{current_datetime}")
+    os.makedirs(log_dir, exist_ok=True)
+    if train:
+        # Callbacks setup
+        checkpoint_callback = CheckpointCallback(save_freq=total_timesteps // 1000, save_path=log_dir, name_prefix="ppo_model")
+        event_callback = EveryNTimesteps(n_steps=total_timesteps // 1000, callback=checkpoint_callback)
+        eval_callback = EvalCallback(env, best_model_save_path=log_dir,
+                                    log_path=log_dir, eval_freq=steps * 20,
+                                    deterministic=True, render=False)
+        callbacks = CallbackList([checkpoint_callback, event_callback, eval_callback])
+
+        if not load_model:
+            # Create new model with curriculum learning
+            # model = PPO("MultiInputPolicy", env, verbose=1, tensorboard_log=log_dir,
+            #             learning_rate=lambda f: f * 3e-4,  # Increased initial learning rate
+            #             n_steps=256,
+            #             batch_size=64,
+            #             n_epochs=10,
+            #             gamma=0.99,
+            #             ent_coef=0.01,
+            #             use_sde=True,
+            #             sde_sample_freq=4,
+            #             target_kl=0.02,
+            #             clip_range=lambda f: f * 0.2)  # Clip range annealing
+            # model = PPO("MultiInputPolicy", env, verbose=1, tensorboard_log=log_dir,
+            #     learning_rate=1e-3,  # Fixed, higher learning rate
+            #     n_steps=512,  # Increased steps per update
+            #     batch_size=64,
+            #     n_epochs=10,
+            #     gamma=0.99,
+            #     ent_coef=0.05,  
+            #     use_sde=False,  # Disabled SDE
+            #     target_kl=0.02,
+            #     clip_range=0.3)  
+            # model = PPO("MultiInputPolicy", env, verbose=1, tensorboard_log=log_dir,
+            #     learning_rate=5e-4,  # Reduced learning rate
+            #     n_steps=100,  # Keep increased steps per update
+            #     batch_size=64,
+            #     n_epochs=10,
+            #     gamma=0.99,
+            #     ent_coef=0.02,  
+            #     use_sde=False,  
+            #     target_kl=0.015,  
+            #     clip_range=0.2)  
+            model = PPO("MultiInputPolicy", env, verbose=1, tensorboard_log=log_dir,
+                learning_rate=1e-3,  
+                batch_size=32,  
+                n_epochs=5,  
+                gamma=0.99,
+                ent_coef=0.01,
+                use_sde=False,
+                target_kl=0.02,
+                clip_range=0.2)
+        else:
+            # Load pre-trained model
+            # path = os.path.join(FIGRURES_DIR, "models", "PPO_2024-06-28_11-53-39", "ppo_model_2047_steps") #30718 20479 || ppo_model_10240_stepsGOOD+8185+ppo_model_6139_steps(c3) || back to PPO_2024-06-27_16-47-34 ppo_model_5116_steps(c3) +PPO_2024-06-27_19-43-11ppo_model_4093_steps+PPO_2024-06-27_20-59-07ppo_model_3570_steps+ppo_model_5116_stepsGOODPPO_2024-06-27_21-22-07+ppo_model_71611_steps
+            # model = PPO.load(path, env=env)
+            # model.learning_rate = 2e-2
+            # model.learn(total_timesteps=5000)  # Train for a short while with higher rate
+            # model.learning_rate = 1e-3  # Then reduce it
+            # model.learn(total_timesteps=20000)  # Continue training
+            # # Final lower learning rate phase
+            # model.learning_rate = 5e-4
+            # model.learn(total_timesteps=5000)
+            # model.n_steps = 256
+            # model.batch_size = 64
+            # model.n_epochs = 10
+            # model.ent_coef = 0.05
+            
+            # model.gamma = 0.99
+
+            # model.rollout_buffer = model.rollout_buffer.__class__(
+            #     model.n_steps,
+            #     model.observation_space,
+            #     model.action_space,
+            #     device=model.device,
+            #     gamma=model.gamma,
+            #     gae_lambda=model.gae_lambda,
+            #     n_envs=model.n_envs,
+            # )
+            # model.rollout_buffer.reset()
+            # path = os.path.join(FIGRURES_DIR, "models", "PPO_2024-06-27_23-43-49", "ppo_model_5116_steps")
+            path = os.path.join(FIGRURES_DIR, "models", "PPO_2024-06-28_13-57-22", "ppo_model_10231_stepsWORKWORKWORK.zip")
+
+            model = PPO.load(path, env=env)
+
+            # Define a function for a very high clip range (effectively no clipping)
+            def high_clip(progress_remaining):
+                return 5.0
+
+            # Initial phase with high learning rate and effectively no clipping
+            model.learning_rate = 1e-3
+            model.clip_range = high_clip
+            model.learn(total_timesteps=5000)
+
+            # Second phase with slightly lower learning rate but still high
+            model.learning_rate = 1e-4
+            model.learn(total_timesteps=20000)
+
+            # Define a function for standard clipping
+            def standard_clip(progress_remaining):
+                return 0.2
+
+            # Final phase with lower learning rate and standard clipping
+            model.learning_rate = 5e-4
+            model.clip_range = standard_clip
+            model.learn(total_timesteps=5000)
+
+            # Set other parameters
+            model.n_steps = 512
+            model.batch_size = 64
+            model.n_epochs = 10
+            model.ent_coef = 0.05
+            model.gamma = 0.99
+
+            # Reset the rollout buffer with the final parameters
+            model.rollout_buffer = model.rollout_buffer.__class__(
+                model.n_steps,
+                model.observation_space,
+                model.action_space,
+                device=model.device,
+                gamma=model.gamma,
+                gae_lambda=model.gae_lambda,
+                n_envs=model.n_envs,
+            )
+            model.rollout_buffer.reset()    
+
+        ## Curriculum learning
+        for i in range(5):  # 5 stages of curriculum
+            current_stage = env.env_method("get_current_stage")[0]
+            print(f"Curriculum stage {i+1}/5: {current_stage}")
+            print(f"Threshold reward for this stage: {env.env_method('get_stage_threshold')[0]}")
+            model.learn(total_timesteps=total_timesteps // 5, callback=callbacks)
+            # Progress to the next stage
+            env.env_method("increase_difficulty")
+
+        # Save the final model
+        model.save(os.path.join(log_dir, "final_model"))
+    else:
+        # Evaluation code
+        # Load the specific model you want to evaluate
+        path = os.path.join(FIGRURES_DIR, "models", "PPO_2024-06-28_13-57-22", "ppo_model_10231_stepsWORKWORKWORK.zip")
+        model = PPO.load(path, env=env)
+
+        # Evaluate the model
+        eval_episodes = 10
+        for episode in range(eval_episodes):
+            obs = env.reset()
+            done = False
+            total_reward = 0
+            step = 0
+            
+            while not done:
+                action, _ = model.predict(obs, deterministic=True)
+                obs, reward, done, info = env.step(action)
+                reward = reward[0]  # Extract scalar value from numpy array
+                total_reward += reward
+                step += 1
+                
+                # Print step information
+                print(f"Episode {episode+1}, Step {step}: Reward = {reward:.2f}, Total Reward = {total_reward:.2f}")
+                
+                # Check if the episode is done
+                if done[0]:
+                    break
+
+                # Check if the episode is done due to reaching the maximum steps
+                if step >= env.get_attr('ep_steps')[0]:
+                    print(f"Episode {episode+1} reached maximum steps.")
+                    break
+
+            print(f"Episode {episode+1} finished: Total Reward = {total_reward:.2f}, Steps = {step}")
+            
+            # Print the final stage reached
+            final_stage = env.get_attr('get_current_stage')[0]
+            print(f"Final stage reached: {final_stage}")
+            
+        print("Evaluation complete.")
+
+# Position:  -2.4000000000000026 0.07699995934963236 0.03970504179596901
+# Orientation:  -1.57191039448275 -1.5144899542893442 -1.5719101704888712
+
+
+#PPO_2024-06-27_23-43-49 ppo_model_71611_steps longest training + PPO_2024-06-28_09-09-00 ppo_model_10231_steps
+
+    """
+    +
+    PPO_2024-06-28_10-47-03
+    ppo_model_8185_steps
+    +
+    PPO_2024-06-28_11-53-39
+    ppo_model_2047_steps
+    ------------------------------------------------------------
+    PPO_2024-06-27_23-43-49
+    ppo_model_38875_steps
+    ppo_model_10231_steps
+    ppo_model_5116_steps
+    
+    """