direct and residual movement inconsistency fixed

ml/env.py

@@ -184,7 +184,19 @@ class JackBotEnv(gym.Env):
         return self._get_obs(), {}
 
     def _get_obs(self) -> np.ndarray:
-        return self.robot.get_observation(command=self.command)
+        # Read raw joint angles from backend
+        raw_angles = np.asarray(self.robot.backend.get_joint_angles(), dtype=np.float32).flatten()
+
+        min_lim = self.min_joint_limits.flatten()
+        max_lim = self.max_joint_limits.flatten()
+
+        # Map raw joint radians [min, max] -> normalized [-1, 1]
+        normalized_joints = 2.0 * (raw_angles - min_lim) / (max_lim - min_lim) - 1.0
+        normalized_joints = np.clip(normalized_joints, -1.0, 1.0)
+
+        # Concatenate normalized joints with active command vector
+        obs = np.concatenate([normalized_joints, self.command]).astype(np.float32)
+        return obs
 
     def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
         previous_action = self.last_action.copy()
@@ -194,25 +206,36 @@ class JackBotEnv(gym.Env):
         self.last_last_action = self.last_action.copy()
         self.last_action = action.copy()
 
-        # Command resampling
-        if self.random_command and (self.step_count >= self.next_cmd_resample_step or self._curriculum_advanced):
+        # Command resampling ONLY if random_command is True
+        if self.random_command and (
+            self.step_count >= self.next_cmd_resample_step or self._curriculum_advanced):
             self.command = self.sample_command()
             random_interval = np.random.randint(self.min_cmd_hold_steps, self.max_cmd_hold_steps + 1)
             self.next_cmd_resample_step = self.step_count + random_interval
 
-        # Extract [vx, vy, omega]
+        # Extract active [vx, vy, omega]
         cmd_vx, cmd_vy, cmd_omega = self.command
 
-        # Mirror main.py input resolution logic: update robot_state and vector_dirmov directly
+        # Mirror input resolution logic to keep robot state synchronized
         self.robot.robot_state = "walking" if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01) else "idle"
         self.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
-        
-        joint_range = np.minimum(
-            self.default_joint_angles - self.min_joint_limits,
-            self.max_joint_limits - self.default_joint_angles
-        )
-        target_angles = self.default_joint_angles + action * joint_range
-        self.robot.tick(action=target_angles)
+
+        # Direct Mode: Target joint scaling
+        action_flat = np.asarray(action, dtype=np.float32).flatten()
+        action_clipped = np.clip(action_flat, -1.0, 1.0)
+
+        if self.robot_mode == "direct":
+            # Map [-1, 1] linearly to physical joint limits [min, max]
+            min_lim = self.min_joint_limits.flatten()
+            max_lim = self.max_joint_limits.flatten()
+            
+            target_angles = min_lim + (action_clipped + 1.0) * 0.5 * (max_lim - min_lim)
+        else:
+            # Residual mode mapping logic
+            target_angles = self.default_joint_angles.flatten() + action_clipped * 0.20
+
+        # Apply target joint angles to physics engine
+        self.robot.tick(action=target_angles, physics_substeps=4)
 
         if self.robot_mode != "kinematics" and self.step_count % 60 == 0:
             random_force = np.random.uniform(-2.0, 2.0, size=2)
@@ -222,8 +245,9 @@ class JackBotEnv(gym.Env):
         self._update_distance_metrics()
         self._update_curriculum()
 
+        # Build next observation preserving active command
         obs = self._get_obs()
-        reward = self._compute_reward(action, previous_action)
+        reward = self._compute_reward(action_flat, previous_action)
 
         self.cumulative_reward += reward
         self.robot_reward += reward
@@ -235,6 +259,9 @@ class JackBotEnv(gym.Env):
         if self.step_count % 120 == 0 and self.use_gui:
             self._update_hud()
 
+        if self.use_gui:
+            time.sleep(1.0 / self.control_freq)
+
         return obs, reward, terminated, truncated, info
 
     def _update_distance_metrics(self):