env cleanup

rewards adjustment
pybullet logic contained in SimManager

Robot.py

@@ -172,9 +172,7 @@ class Robot:
     # --- RL METHODS ---
 
     def apply_rl_action(self, action: np.ndarray) -> None:
-        """
-        Applies continuous RL action deltas [-1, 1] to current joint angles.
-        """
+        """Applies continuous RL action deltas [-1, 1] to current joint angles."""
         action = np.asarray(action, dtype=np.float32)
         scaled_action = np.clip(action, -1.0, 1.0) * self.action_scale
 
@@ -191,20 +189,19 @@ class Robot:
     def get_observation(self, command: Optional[np.ndarray] = None) -> np.ndarray:
         """
         Returns observation vector [18 joint angles] + [optional 4 command dimensions].
-        Queries PyBullet if backend is PyBulletBackend; otherwise falls back to internal state.
+        Queries SimManager helper if PyBulletBackend is used; falls back to internal state otherwise.
         """
-        if isinstance(self.backend, PyBulletBackend) and self.backend.sim and hasattr(self.backend.sim, 'physics_client'):
-            physics_client = self.backend.sim.physics_client
+        if isinstance(self.backend, PyBulletBackend) and self.backend.sim:
             body_id = self.backend.body_id if self.backend.body_id is not None else 0
-            
-            # Retrieve joint mapping from SimManager/Simulation if available
-            if hasattr(self.backend.sim, 'robot_joints') and body_id in self.backend.sim.robot_joints:
-                joint_indices = self.backend.sim.robot_joints[body_id]
+            if hasattr(self.backend.sim, 'get_robot_joint_angles'):
+                joint_angles = self.backend.sim.get_robot_joint_angles(body_id)
+            elif hasattr(self.backend.sim, 'physics_client'):
+                physics_client = self.backend.sim.physics_client
+                joint_indices = self.backend.sim.robot_joints.get(body_id, list(range(18))) if hasattr(self.backend.sim, 'robot_joints') else list(range(18))
+                joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=physics_client)
+                joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32)
             else:
-                joint_indices = list(range(18))
-
-            joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=physics_client)
-            joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32)
+                joint_angles = self.current_rad.data.flatten().astype(np.float32)
         else:
             joint_angles = self.current_rad.data.flatten().astype(np.float32)
 

ml/SimManager.py

@@ -1,9 +1,10 @@
 """
 ml/SimManager.py - PyBullet Simulation & Multi-Body Manager
 """
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Optional
 import pybullet as p
 import pybullet_data
+import numpy as np
 import DataTypes as dt
 
 
@@ -84,4 +85,82 @@ class SimManager:
     def disconnect(self):
         if self.physics_client is not None and p.isConnected(self.physics_client):
             p.disconnect(self.physics_client)
-            self.physics_client = None
+            self.physics_client = None
+
+    # --- ROBOT GETTERS AND SETTERS ---
+
+    def reset_robot_base(
+        self,
+        body_id: int,
+        position: List[float],
+        orientation: Optional[List[float]] = None,
+        linear_velocity: Optional[List[float]] = None,
+        angular_velocity: Optional[List[float]] = None
+    ) -> None:
+        """Resets a robot body's base position, orientation, and velocities."""
+        if orientation is None:
+            orientation = [0.0, 0.0, 0.0, 1.0]
+        if linear_velocity is None:
+            linear_velocity = [0.0, 0.0, 0.0]
+        if angular_velocity is None:
+            angular_velocity = [0.0, 0.0, 0.0]
+
+        p.resetBasePositionAndOrientation(
+            body_id, position, orientation, physicsClientId=self.physics_client
+        )
+        p.resetBaseVelocity(
+            body_id, linearVelocity=linear_velocity, angularVelocity=angular_velocity,
+            physicsClientId=self.physics_client
+        )
+
+    def get_robot_pose(self, body_id: int) -> Tuple[List[float], List[float]]:
+        """Returns base position (x, y, z) and orientation quaternion (x, y, z, w)."""
+        pos, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
+        return list(pos), list(orn)
+
+    def get_robot_rpy(self, body_id: int) -> Tuple[float, float, float]:
+        """Returns roll, pitch, yaw angles in radians for the given robot body."""
+        _, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
+        roll, pitch, yaw = p.getEulerFromQuaternion(orn)
+        return float(roll), float(pitch), float(yaw)
+
+    def get_robot_pose_and_rpy(self, body_id: int) -> Tuple[List[float], Tuple[float, float, float]]:
+        """Returns base position and (roll, pitch, yaw) tuple in radians."""
+        pos, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
+        roll, pitch, yaw = p.getEulerFromQuaternion(orn)
+        return list(pos), (float(roll), float(pitch), float(yaw))
+
+    def get_robot_velocity(self, body_id: int) -> Tuple[List[float], List[float]]:
+        """Returns linear velocity (vx, vy, vz) and angular velocity (wx, wy, wz)."""
+        lin_v, ang_v = p.getBaseVelocity(body_id, physicsClientId=self.physics_client)
+        return list(lin_v), list(ang_v)
+
+    def get_robot_joint_angles(self, body_id: int, joint_indices: Optional[List[int]] = None) -> np.ndarray:
+        """Returns joint angles as a 1D numpy array float32 for specified or registered joint indices."""
+        if joint_indices is None:
+            joint_indices = self.robot_joints.get(body_id, list(range(18)))
+        joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=self.physics_client)
+        return np.array([state[0] for state in joint_states], dtype=np.float32)
+
+    def set_robot_color(self, body_id: int, rgba: List[float]) -> None:
+        """Changes visual color RGBA of base link and all joints of the specified robot body."""
+        num_joints = p.getNumJoints(body_id, physicsClientId=self.physics_client)
+        p.changeVisualShape(body_id, -1, rgbaColor=rgba, physicsClientId=self.physics_client)
+        for j in range(num_joints):
+            p.changeVisualShape(body_id, j, rgbaColor=rgba, physicsClientId=self.physics_client)
+
+    def measure_robot_heights(self, robot_ids: List[int]) -> List[float]:
+        """Gets current Z height for all specified robot body IDs."""
+        heights = []
+        for body_id in robot_ids:
+            pos, _ = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
+            heights.append(pos[2])
+        return heights
+
+    def settle_and_measure_height(self, robot_ids: List[int], steps: int = 200, fallback_height: float = 0.14) -> float:
+        """Steps simulation for designated steps so robot settles, then calculates target standing height."""
+        for _ in range(steps):
+            self.step()
+        heights = self.measure_robot_heights(robot_ids)
+        mean_height = float(np.mean(heights)) if heights else fallback_height
+        return mean_height if mean_height > 0.0 else fallback_height

ml/env.py

@@ -8,7 +8,6 @@ from typing import Optional, Tuple, Dict, Any, List
 import gymnasium as gym
 from gymnasium import spaces
 import numpy as np
-import pybullet as p
 
 from config import cfg
 from Robot import Robot, PyBulletBackend
@@ -79,9 +78,7 @@ class JackBotEnv(gym.Env):
         self.collapse_height_fraction = 0.55
         self.tilt_failure_rad = 0.9
 
-        # Small random exploration pulse to break local optima.
-        # The bonus is only granted on a rare step and only when the robot
-        # is still stable enough that a bold move is likely to remain safe.
+        # Small random exploration pulse settings
         self.exploration_bonus_prob = 0.03
         self.exploration_bonus_interval = 120
         self.exploration_bonus_scale = 0.08
@@ -93,7 +90,7 @@ class JackBotEnv(gym.Env):
         self.curriculum_phase = 0
         self.curriculum_episode_limit = 150
         self.curriculum_stage_requirements = {
-            1: {"survival_steps": 200, "distance": 0.00, "stability_roll_pitch": 0.35},
+            1: {"survival_steps": 1000, "distance": 0.00, "stability_roll_pitch": 0.35},
             2: {"survival_steps": 350, "distance": 10.0, "stability_roll_pitch": 0.30},
             3: {"survival_steps": 550, "distance": 15.00, "stability_roll_pitch": 0.25},
         }
@@ -103,13 +100,6 @@ class JackBotEnv(gym.Env):
         self.leader_crown = LeaderCrown(physics_client_id=self.sim_manager.physics_client)
         self.last_time = time.time()
 
-    def _set_robot_color(self, pb_id: int, rgba: List[float]):
-        """Helper to change the visual color of a robot body and all its links."""
-        num_joints = p.getNumJoints(pb_id, physicsClientId=self.sim_manager.physics_client)
-        p.changeVisualShape(pb_id, -1, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client)
-        for j in range(num_joints):
-            p.changeVisualShape(pb_id, j, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client)
-
     def _robot_base_position(self, robot_id: int, spacing: float = 0.5) -> list[float]:
         return [0.0, 0.0, 0.14]
 
@@ -148,20 +138,13 @@ class JackBotEnv(gym.Env):
 
         for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
             spawn_pos = self._robot_base_position(idx, self.robot_spacing)
-            spawn_orn = [0, 0, 0, 1]
-
-            p.resetBasePositionAndOrientation(
-                pb_id, spawn_pos, spawn_orn, physicsClientId=self.sim_manager.physics_client
-            )
-            p.resetBaseVelocity(
-                pb_id, linearVelocity=[0, 0, 0], angularVelocity=[0, 0, 0],
-                physicsClientId=self.sim_manager.physics_client
-            )
+            spawn_orn = [0.0, 0.0, 0.0, 1.0]
 
+            self.sim_manager.reset_robot_base(pb_id, spawn_pos, spawn_orn)
             robot_obj.reset_to_init()
 
             if self.use_gui:
-                self._set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0])
+                self.sim_manager.set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0])
 
         self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32)
         self.max_distance_from_start = [0.0]
@@ -178,18 +161,14 @@ class JackBotEnv(gym.Env):
         else:
             self.commands = np.zeros((1, 4), dtype=np.float32)
 
-        for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
-            pos, _ = p.getBasePositionAndOrientation(
-                pb_id, physicsClientId=self.sim_manager.physics_client
-            )
+        for idx, pb_id in enumerate(self.pb_robots):
+            pos, _ = self.sim_manager.get_robot_pose(pb_id)
             self.start_positions[idx] = [float(pos[0]), float(pos[1]), float(pos[2])]
 
-        for _ in range(200):
-            self.sim_manager.step()
-
-        self.target_height = self._measure_settled_height()
-        if self.target_height <= 0.0:
-            self.target_height = 0.14
+        # Settle for 200 steps after dropping in a standing position, then measure target height
+        self.target_height = self.sim_manager.settle_and_measure_height(
+            self.pb_robots, steps=200, fallback_height=0.14
+        )
 
         if self.use_gui:
             self.hud.reset()
@@ -199,83 +178,63 @@ class JackBotEnv(gym.Env):
         return self._get_obs(), {}
 
     def get_robot_velocities(self) -> list:
-            """Exposes velocities for the SB3 metrics callback."""
-            vels = []
-            for pb_id in self.pb_robots:
-                lin_v, ang_v = p.getBaseVelocity(pb_id, physicsClientId=self.sim_manager.physics_client)
-                vels.append((lin_v, ang_v))
-            return vels
+        """Exposes velocities for the SB3 metrics callback."""
+        vels = []
+        for pb_id in self.pb_robots:
+            lin_v, ang_v = self.sim_manager.get_robot_velocity(pb_id)
+            vels.append((lin_v, ang_v))
+        return vels
 
     def get_robot_distance_metrics(self) -> list:
-            """Exposes distance-from-start metrics for logging without affecting reward."""
-            metrics = []
-            for idx, pb_id in enumerate(self.pb_robots):
-                pos, _ = p.getBasePositionAndOrientation(
-                    pb_id, physicsClientId=self.sim_manager.physics_client
-                )
-                start_x, start_y, _ = self.start_positions[idx]
-                dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
-                self.max_distance_from_start[idx] = max(self.max_distance_from_start[idx], dist)
-                metrics.append((dist, self.max_distance_from_start[idx]))
+        """Exposes distance-from-start metrics for logging without affecting reward."""
+        metrics = []
+        for idx, pb_id in enumerate(self.pb_robots):
+            pos, _ = self.sim_manager.get_robot_pose(pb_id)
+            start_x, start_y, _ = self.start_positions[idx]
+            dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
+            self.max_distance_from_start[idx] = max(self.max_distance_from_start[idx], dist)
+            metrics.append((dist, self.max_distance_from_start[idx]))
 
-            self.max_survival_steps = max(self.max_survival_steps, self.step_count)
-            return metrics
+        self.max_survival_steps = max(self.max_survival_steps, self.step_count)
+        return metrics
 
     def get_current_robot_metrics(self) -> list:
-            """Returns current per-robot reward, distance, and velocity summaries for alive robots only."""
-            metrics = []
-            for idx, pb_id in enumerate(self.pb_robots):
-                if self.failed_robots_mask[idx]:
-                    continue
+        """Returns current per-robot reward, distance, and velocity summaries for alive robots only."""
+        metrics = []
+        for idx, pb_id in enumerate(self.pb_robots):
+            if self.failed_robots_mask[idx]:
+                continue
 
-                pos, _ = p.getBasePositionAndOrientation(
-                    pb_id, physicsClientId=self.sim_manager.physics_client
-                )
-                linear_vel, angular_vel = p.getBaseVelocity(
-                    pb_id, physicsClientId=self.sim_manager.physics_client
-                )
-                start_x, start_y, _ = self.start_positions[idx]
-                dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
-                speed = float(np.linalg.norm(np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)))
-                yaw_rate = float(abs(angular_vel[2]))
-                metrics.append({
-                    "reward": float(self.robot_rewards[idx]),
-                    "distance_from_start": dist,
-                    "speed": speed,
-                    "yaw_rate": yaw_rate,
-                    "alive": True,
-                    "survival_steps": int(self.step_count),
-                })
+            pos, _ = self.sim_manager.get_robot_pose(pb_id)
+            linear_vel, angular_vel = self.sim_manager.get_robot_velocity(pb_id)
+            start_x, start_y, _ = self.start_positions[idx]
+            dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
+            speed = float(np.linalg.norm(np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)))
+            yaw_rate = float(abs(angular_vel[2]))
+            metrics.append({
+                "reward": float(self.robot_rewards[idx]),
+                "distance_from_start": dist,
+                "speed": speed,
+                "yaw_rate": yaw_rate,
+                "alive": True,
+                "survival_steps": int(self.step_count),
+            })
 
-            return metrics
+        return metrics
 
     def get_survival_steps(self) -> int:
-            """Returns the current survival length for the environment's current episode."""
-            return int(self.step_count)
+        """Returns the current survival length for the environment's current episode."""
+        return int(self.step_count)
 
     def _get_obs(self) -> np.ndarray:
         obs_list = []
         for idx, (pb_id, joint_indices) in enumerate(zip(self.pb_robots, self.robot_joint_indices)):
-            joint_states = p.getJointStates(
-                pb_id, 
-                joint_indices, 
-                physicsClientId=self.sim_manager.physics_client
-            )
-            joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32)
+            joint_angles = self.sim_manager.get_robot_joint_angles(pb_id, joint_indices)
             robot_obs = np.concatenate([joint_angles, self.commands[idx]])
             obs_list.append(robot_obs)
 
         return np.concatenate(obs_list).astype(np.float32)
 
-    def _measure_settled_height(self) -> float:
-        heights = []
-        for pb_id in self.pb_robots:
-            pos, _ = p.getBasePositionAndOrientation(
-                pb_id, physicsClientId=self.sim_manager.physics_client
-            )
-            heights.append(pos[2])
-        return float(np.mean(heights)) if heights else 0.14
-
     def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
         self.step_count += 1
         self.total_steps += 1
@@ -322,10 +281,7 @@ class JackBotEnv(gym.Env):
         survival_ok = self.max_survival_steps >= req["survival_steps"]
         distance_ok = self.max_distance_from_start[0] >= req["distance"]
 
-        position, orientation = p.getBasePositionAndOrientation(
-            self.pb_robots[0], physicsClientId=self.sim_manager.physics_client
-        )
-        roll, pitch, _ = p.getEulerFromQuaternion(orientation)
+        position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robots[0])
         stability_ok = abs(roll) <= req["stability_roll_pitch"] and abs(pitch) <= req["stability_roll_pitch"]
         height_ok = position[2] >= max(0.09, self.target_height * 0.85)
 
@@ -363,78 +319,55 @@ class JackBotEnv(gym.Env):
         previous_actions = previous_action.reshape(1, 18)
 
         for idx, pb_id in enumerate(self.pb_robots):
-            pos, orientation = p.getBasePositionAndOrientation(
-                pb_id, physicsClientId=self.sim_manager.physics_client
-            )
-            linear_vel, angular_vel = p.getBaseVelocity(
-                pb_id, physicsClientId=self.sim_manager.physics_client
-            )
-            roll, pitch, _ = p.getEulerFromQuaternion(orientation)
+            pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id)
+            linear_vel, angular_vel = self.sim_manager.get_robot_velocity(pb_id)
             command = self.commands[idx]
 
             cmd_vx = command[0]
             cmd_vy = command[1]
             cmd_yaw = command[3]
 
-            # -------------------------------------------------------------
-            # 1. LINEAR VECTOR SPEED MAXIMIZATION (Magnitude + Direction)
-            # -------------------------------------------------------------
+            # 1. LINEAR VECTOR SPEED MAXIMIZATION
             cmd_dir = np.array([cmd_vx, cmd_vy], dtype=np.float32)
             cmd_norm = np.linalg.norm(cmd_dir)
 
             if cmd_norm > 0.05:
-                # Normalize target direction vector
                 unit_cmd_dir = cmd_dir / cmd_norm
                 actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)
 
-                # Speed aligned with target direction (m/s)
                 aligned_speed = float(np.dot(actual_vel_2d, unit_cmd_dir))
-
-                # Moderate movement reward to encourage directional motion
                 linear_speed_reward = 2.5 * aligned_speed
 
-                # Penalize sideways drift (perpendicular velocity to commanded direction)
                 perp_vel = actual_vel_2d - aligned_speed * unit_cmd_dir
                 drift_penalty = 0.35 * float(np.dot(perp_vel, perp_vel))
             else:
-                # If no linear command given, penalize all horizontal movement
                 linear_speed_reward = 0.0
                 drift_penalty = 1.0 * (linear_vel[0]**2 + linear_vel[1]**2)
 
-            # -------------------------------------------------------------
             # 2. TURNING SPEED MAXIMIZATION
-            # -------------------------------------------------------------
-            actual_yaw_rate = angular_vel[2]  # rad/s in PyBullet Z-axis
+            actual_yaw_rate = angular_vel[2]
 
             if abs(cmd_yaw) > 0.05:
-                # Reward turning in the commanded direction, but less aggressively
                 turning_reward = 1.5 * (actual_yaw_rate * cmd_yaw)
             else:
-                # Penalize unwanted rotation when joystick turn is centered
                 turning_reward = -0.6 * (actual_yaw_rate ** 2)
 
-            # -------------------------------------------------------------
             # 3. SMALL ALIVE BONUS + AGE-RAMPED STILLNESS PENALTY
-            # -------------------------------------------------------------
-            alive_reward = 0.02
+            alive_reward = 0.5
 
             age_ratio = min(1.0, self.step_count / max(1, self.curriculum_episode_limit))
             still_penalty = 0.0
             if (cmd_norm > 0.1 or abs(cmd_yaw) > 0.1) and (abs(linear_vel[0]) < 0.02 and abs(actual_yaw_rate) < 0.05):
                 still_penalty = 0.35 + 0.85 * age_ratio
 
-            # -------------------------------------------------------------
             # 4. POSTURE & STABILITY PENALTIES
-            # -------------------------------------------------------------
-            height_penalty = 6.0 * ((self.target_height - pos[2]) ** 2) if pos[2] < self.target_height else 0.0
+            height_penalty = 10.0 * ((self.target_height - pos[2]) ** 2) if pos[2] < self.target_height else 0.0
             stability_penalty = 1.5 * (roll**2 + pitch**2)
-
-            # Penalize only large control jumps; allow continuous low-amplitude motion
+            print(-height_penalty)
             control_delta = np.abs(current_actions[idx] - previous_actions[idx])
             large_delta_mask = control_delta > 0.12
             large_delta_penalty = 0.002 * float(np.sum(np.square(control_delta[large_delta_mask]))) if np.any(large_delta_mask) else 0.0
 
-            # Combined Step Reward
             r_step = (
                 alive_reward
                 + linear_speed_reward
@@ -465,8 +398,7 @@ class JackBotEnv(gym.Env):
         rolls = []
         pitches = []
         for pb_id in self.pb_robots:
-            pos, orient = p.getBasePositionAndOrientation(pb_id, physicsClientId=self.sim_manager.physics_client)
-            roll, pitch, _ = p.getEulerFromQuaternion(orient)
+            pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id)
             heights.append(pos[2])
             rolls.append(math.degrees(roll))
             pitches.append(math.degrees(pitch))
@@ -490,9 +422,7 @@ class JackBotEnv(gym.Env):
 
         best_idx = int(np.argmax(self.robot_rewards))
         leader_pb_id = self.pb_robots[best_idx]
-        leader_pos, _ = p.getBasePositionAndOrientation(
-            leader_pb_id, physicsClientId=self.sim_manager.physics_client
-        )
+        leader_pos, _ = self.sim_manager.get_robot_pose(leader_pb_id)
         self.leader_crown.update(leader_pos)
 
     def _update_robot_failures(self):
@@ -501,10 +431,7 @@ class JackBotEnv(gym.Env):
             if self.failed_robots_mask[idx]:
                 continue
 
-            position, orientation = p.getBasePositionAndOrientation(
-                pb_id, physicsClientId=self.sim_manager.physics_client
-            )
-            roll, pitch, _ = p.getEulerFromQuaternion(orientation)
+            position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id)
 
             collapse_threshold = max(0.06, self.collapse_height_fraction * self.target_height)
             is_tilted = abs(roll) > self.tilt_failure_rad or abs(pitch) > self.tilt_failure_rad
@@ -513,7 +440,7 @@ class JackBotEnv(gym.Env):
             if is_tilted or is_collapsed:
                 self.failed_robots_mask[idx] = True
                 if self.use_gui:
-                    self._set_robot_color(pb_id, COLOR_FAILED)
+                    self.sim_manager.set_robot_color(pb_id, COLOR_FAILED)
 
     def close(self):
         self.sim_manager.disconnect()

ml/train.py

@@ -178,13 +178,18 @@ def train(
 
     device = resolve_device(device)
 
+    policy_kwargs = dict(
+    log_std_init=-1.5,  # Sets initial std ~ 0.22 instead of 1.0
+    net_arch=dict(pi=[256, 256], vf=[256, 256])
+)
+
     model = PPO(
         "MlpPolicy",
         env,
         verbose=1,
         seed=seed,
-        learning_rate=3.5e-4,      # Cut LR in half (from 3e-4) to smooth out updates
-        n_steps=2048,              # Larger rollout buffer per env for stable gradients
+        learning_rate=1.5e-4,      # Cut LR in half (from 3e-4) to smooth out updates
+        n_steps=1024,              # Larger rollout buffer per env for stable gradients
         batch_size=128,            # Larger minibatches reduce noise
         n_epochs=10,               # Number of epoch updates per rollout
         gamma=0.99,                # Discount factor
@@ -195,6 +200,7 @@ def train(
         vf_coef=0.5,
         max_grad_norm=0.5,
         device=device,
+        policy_kwargs=policy_kwargs,
         tensorboard_log=str(Path(__file__).resolve().parent / "tensorboard"),
     )