Motor car simulator

8 Novaxfx

class CarSimulator:
    """
    2D car physics simulator. The agent learns to control:
    - throttle: acceleration/braking
    - steering: turning angle
    
    State includes position, velocity, heading, and track sensors.
    """
    
    def __init__(self, track_type='oval', max_steps=1000):
        self.max_steps = max_steps
        self.dt = 0.05  # 20 Hz simulation
        
        # Car physics parameters
        self.length = 2.5       # wheelbase (m)
        self.max_speed = 30.0   # m/s (~108 km/h)
        self.max_accel = 5.0    # m/s^2
        self.max_steering = np.radians(30)  # max steering angle
        self.friction = 0.98    # rolling friction
        
        # Track definition
        self.track = self._generate_track(track_type)
        self.track_width = 8.0
        
        # State
        self.position = np.zeros(2)   # x, y
        self.velocity = 0.0            # forward speed
        self.heading = 0.0             # radians
        self.steering_angle = 0.0
        self.throttle = 0.0
        self.steps = 0
        self.total_distance = 0.0
        self.lap_progress = 0.0
        
        # Sensor configuration
        self.sensor_angles = np.linspace(-np.pi/2, np.pi/2, 7)  # 7 ray sensors
        
    def _generate_track(self, track_type):
        """Generate track waypoints."""
        if track_type == 'oval':
            # Simple oval track
            t = np.linspace(0, 2*np.pi, 100)
            center_x = 50 * np.cos(t)
            center_y = 30 * np.sin(t)
            # Scale to make it more oval
            center_x *= 1.5
        elif track_type == 'figure8':
            t = np.linspace(0, 2*np.pi, 200)
            center_x = 40 * np.sin(t)
            center_y = 30 * np.sin(2*t)
        else:
            # Random winding track
            t = np.linspace(0, 2*np.pi, 150)
            center_x = 50 * np.cos(t) + 15 * np.cos(3*t)
            center_y = 40 * np.sin(t) + 10 * np.sin(3*t)
            
        return np.column_stack([center_x, center_y])
    
    def reset(self):
        """Reset car to starting position."""
        start_idx = 0
        self.position = self.track[start_idx].copy()
        
        # Face along track direction
        next_point = self.track[(start_idx + 1) % len(self.track)]
        self.heading = np.arctan2(next_point[1] - self.position[1],
                                   next_point[0] - self.position[0])
        self.velocity = 0.0
        self.steering_angle = 0.0
        self.throttle = 0.0
        self.steps = 0
        self.total_distance = 0.0
        self.lap_progress = 0.0
        
        return self._get_observation()
    
    def _get_track_distance(self, point, segment_start, segment_end):
        """Calculate perpendicular distance from point to track segment."""
        # Vector from segment start to end
        segment = segment_end - segment_start
        segment_len = np.linalg.norm(segment)
        if segment_len < 1e-6:
            return np.linalg.norm(point - segment_start)
        
        # Project point onto segment
        t = np.clip(np.dot(point - segment_start, segment) / (segment_len**2), 0, 1)
        projection = segment_start + t * segment
        
        # Perpendicular distance from center line
        center_dist = np.linalg.norm(point - projection)
        
        # Signed distance (positive = right of track direction, negative = left)
        cross = np.cross(segment, point - segment_start)
        signed_dist = np.sign(cross) * center_dist
        
        return signed_dist
    
    def _get_observation(self):
        """
        Get sensor observations for the agent.
        Returns: [speed, track_center_offset, track_angle, sensor_readings...]
        """
        # Find nearest track segment
        distances_to_track = np.linalg.norm(self.track - self.position, axis=1)
        nearest_idx = np.argmin(distances_to_track)
        
        # Track direction at nearest point
        next_idx = (nearest_idx + 1) % len(self.track)
        prev_idx = (nearest_idx - 1) % len(self.track)
        
        track_dir = self.track[next_idx] - self.track[nearest_idx]
        track_angle = np.arctan2(track_dir[1], track_dir[0])
        
        # Position relative to track center
        signed_dist = self._get_track_distance(
            self.position, 
            self.track[nearest_idx], 
            self.track[next_idx]
        )
        center_offset = signed_dist / (self.track_width / 2)  # normalized [-1, 1] roughly
        
        # Heading relative to track direction
        heading_error = np.angle(np.exp(1j * (self.heading - track_angle)))
        
        # Ray sensors - distance to track boundaries
        sensor_readings = self._cast_sensors()
        
        # Construct observation
        obs = np.array([
            self.velocity / self.max_speed,           # normalized speed
            np.clip(center_offset, -2.0, 2.0),        # position on track
            heading_error / np.pi,                     # heading error normalized
            self.steering_angle / self.max_steering,   # current steering
        ])
        
        obs = np.concatenate([obs, sensor_readings])
        return obs.astype(np.float32)
    
    def _cast_sensors(self):
        """Cast ray sensors to detect track boundaries."""
        readings = []
        
        for sensor_angle in self.sensor_angles:
            absolute_angle = self.heading + sensor_angle
            direction = np.array([np.cos(absolute_angle), np.sin(absolute_angle)])
            
            # Simple ray marching to find track edge
            max_range = 30.0
            step = 0.5
            
            for dist in np.arange(step, max_range, step):
                probe_point = self.position + direction * dist
                
                # Check if off track
                nearest_idx = np.argmin(np.linalg.norm(self.track - probe_point, axis=1))
                next_idx = (nearest_idx + 1) % len(self.track)
                
                signed_dist = self._get_track_distance(
                    probe_point,
                    self.track[nearest_idx],
                    self.track[next_idx]
                )
                
                if abs(signed_dist) > self.track_width / 2:
                    readings.append(max_range / dist)  # normalized inverse distance
                    break
            else:
                readings.append(1.0)  # clear path
                
        return np.array(readings, dtype=np.float32)
    
    def step(self, action):
        """
        Execute one simulation step.
        Action: [throttle, steering] each in [-1, 1]
        """
        self.throttle = np.clip(action[0], -1.0, 1.0)
        steer_input = np.clip(action[1], -1.0, 1.0)
        
        # Physics update
        # Steering rate limit for realism
        target_steering = steer_input * self.max_steering
        steer_rate = 2.0 * self.dt  # steering speed
        self.steering_angle += np.clip(
            target_steering - self.steering_angle, 
            -steer_rate, steer_rate
        )
        
        # Acceleration
        accel = self.throttle * self.max_accel
        
        # Bicycle model kinematics
        # Simple but effective car model
        if abs(self.steering_angle) > 0.01 and self.velocity > 0.1:
            # Turning radius
            turning_radius = self.length / np.tan(self.steering_angle)
            angular_velocity = self.velocity / turning_radius
        else:
            angular_velocity = 0.0
            
        # Update state
        self.velocity += accel * self.dt
        self.velocity *= self.friction  # rolling resistance
        self.velocity = np.clip(self.velocity, -self.max_speed/2, self.max_speed)
        
        self.heading += angular_velocity * self.dt
        self.position[0] += self.velocity * np.cos(self.heading) * self.dt
        self.position[1] += self.velocity * np.sin(self.heading) * self.dt
        
        self.total_distance += abs(self.velocity) * self.dt
        self.steps += 1
        
        # Calculate reward
        obs = self._get_observation()
        reward, done = self._calculate_reward(obs)
        
        return obs, reward, done, {}
    
    def _calculate_reward(self, obs):
        """Calculate step reward and check termination."""
        speed = obs[0] * self.max_speed
        center_offset = obs[1]  # normalized
        
        # Speed reward - encourage going fast but controlled
        speed_reward = speed / self.max_speed
        
        # Centering reward - stay near center line
        center_penalty = -abs(center_offset) ** 2
        
        # Combined reward
        reward = speed_reward * 0.5 + center_penalty * 2.0 + 0.1  # survival bonus
        
        # Check termination
        done = False
        
        # Off track
        if abs(center_offset) > 1.5:  # margin beyond track width
            reward = -10.0
            done = True
            
        # Too slow for too long (spinning out)
        if abs(speed) < 1.0 and self.steps > 100:
            reward = -5.0
            done = True
            
        # Max steps reached
        if self.steps >= self.max_steps:
            done = True
            
        return reward, done

# ============================================================================
# PPO NEURAL NETWORK (ACTOR-CRITIC)
# ============================================================================

class ActorCritic(nn.Module):
    """
    Combined actor-critic network for PPO.
    Actor outputs mean and std for action distribution.
    Critic outputs state value estimate.
    """
    
    def __init__(self, state_dim, action_dim, hidden_dim=256):
        super().__init__()
        
        # Shared feature extractor
        self.shared = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
        )
        
        # Actor head (policy)
        self.actor_mean = nn.Linear(hidden_dim, action_dim)
        self.actor_log_std = nn.Parameter(torch.zeros(action_dim))
        
        # Critic head (value)
        self.critic = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim//2),
            nn.ReLU(),
            nn.Linear(hidden_dim//2, 1)
        )
        
        # Initialize weights
        self.actor_mean.weight.data.mul_(0.1)
        self.actor_mean.bias.data.mul_(0.0)
        
    def forward(self, state):
        features = self.shared(state)
        
        # Actor output
        action_mean = self.actor_mean(features)
        action_mean = torch.tanh(action_mean)  # bounded actions
        
        action_std = torch.exp(self.actor_log_std).expand_as(action_mean)
        
        # Critic output
        value = self.critic(features)
        
        return action_mean, action_std, value
    
    def get_action(self, state, deterministic=False):
        """Sample or return deterministic action."""
        mean, std, value = self.forward(state)
        
        if deterministic:
            action = mean
        else:
            dist = Normal(mean, std)
            action = dist.sample()
            
        # Log probability for training
        dist = Normal(mean, std)
        log_prob = dist.log_prob(action).sum(dim=-1)
        
        return action, log_prob, value
    
    def evaluate(self, state, action):
        """Evaluate actions for PPO update."""
        mean, std, value = self.forward(state)
        dist = Normal(mean, std)
        
        log_prob = dist.log_prob(action).sum(dim=-1)
        entropy = dist.entropy().sum(dim=-1)
        
        return log_prob, value.squeeze(-1), entropy

# ============================================================================
# PPO AGENT
# ============================================================================

class PPOAgent:
    """
    Proximal Policy Optimization agent for continuous motor control.
    """
    
    def __init__(self, state_dim, action_dim, device='cpu'):
        self.device = device
        self.state_dim = state_dim
        self.action_dim = action_dim
        
        # Hyperparameters
        self.lr = 3e-4
        self.gamma = 0.99           # discount factor
        self.gae_lambda = 0.95      # GAE parameter
        self.ppo_epochs = 10
        self.clip_epsilon = 0.2
        self.value_coef = 0.5
        self.entropy_coef = 0.01
        self.batch_size = 64
        
        # Network
        self.policy = ActorCritic(state_dim, action_dim).to(device)
        self.optimizer = optim.Adam(self.policy.parameters(), lr=self.lr)
        
        # Memory for rollout collection
        self.reset_buffers()
        
    def reset_buffers(self):
        self.states = []
        self.actions = []
        self.log_probs = []
        self.rewards = []
        self.values = []
        self.dones = []
        
    def select_action(self, state, training=True):
        """Select action from current policy."""
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
        
        with torch.no_grad() if not training else torch.enable_grad():
            if training:
                action, log_prob, value = self.policy.get_action(state_tensor)
            else:
                action, _, _ = self.policy.get_action(state_tensor, deterministic=True)
                return action.cpu().numpy()[0]
                
        return (action.cpu().numpy()[0], 
                log_prob.cpu().numpy()[0], 
                value.cpu().numpy()[0])
    
    def store_transition(self, state, action, log_prob, reward, value, done):
        """Store transition in rollout buffer."""
        self.states.append(state)
        self.actions.append(action)
        self.log_probs.append(log_prob)
        self.rewards.append(reward)
        self.values.append(value)
        self.dones.append(done)
        
    def compute_gae(self, next_value):
        """Compute Generalized Advantage Estimation."""
        advantages = []
        gae = 0
        
        for t in reversed(range(len(self.rewards))):
            if t == len(self.rewards) - 1:
                next_val = next_value
            else:
                next_val = self.values[t + 1]
                
            delta = (self.rewards[t] + 
                    self.gamma * next_val * (1 - self.dones[t]) - 
                    self.values[t])
            
            gae = delta + self.gamma * self.gae_lambda * (1 - self.dones[t]) * gae
            advantages.insert(0, gae)
            
        return advantages
    
    def update(self, next_state, done):
        """Perform PPO update on collected rollout."""
        # Get final value estimate
        with torch.no_grad():
            next_state_tensor = torch.FloatTensor(next_state).unsqueeze(0).to(self.device)
            _, _, next_value = self.policy.get_action(next_state_tensor)
            next_value = next_value.cpu().numpy()[0]
            
        # Compute returns and advantages
        advantages = self.compute_gae(next_value)
        returns = [adv + val for adv, val in zip(advantages, self.values)]
        
        # Convert to tensors
        states = torch.FloatTensor(np.array(self.states)).to(self.device)
        actions = torch.FloatTensor(np.array(self.actions)).to(self.device)
        old_log_probs = torch.FloatTensor(np.array(self.log_probs)).to(self.device)
        returns = torch.FloatTensor(returns).to(self.device)
        advantages = torch.FloatTensor(advantages).to(self.device)
        
        # Normalize advantages
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        
        # PPO training epochs
        dataset_size = len(self.states)
        indices = np.arange(dataset_size)
        
        for _ in range(self.ppo_epochs):
            np.random.shuffle(indices)
            
            for start in range(0, dataset_size, self.batch_size):
                end = start + self.batch_size
                batch_idx = indices[start:end]
                
                batch_states = states[batch_idx]
                batch_actions = actions[batch_idx]
                batch_old_log_probs = old_log_probs[batch_idx]
                batch_returns = returns[batch_idx]
                batch_advantages = advantages[batch_idx]
                
                # Evaluate current policy
                log_probs, values, entropy = self.policy.evaluate(
                    batch_states, batch_actions
                )
                
                # PPO clipped objective
                ratio = torch.exp(log_probs - batch_old_log_probs)
                surr1 = ratio * batch_advantages
                surr2 = torch.clamp(ratio, 1 - self.clip_epsilon, 
                                   1 + self.clip_epsilon) * batch_advantages
                
                actor_loss = -torch.min(surr1, surr2).mean()
                
                # Value loss
                value_loss = nn.MSELoss()(values, batch_returns)
                
                # Entropy bonus for exploration
                entropy_loss = -entropy.mean()
                
                # Total loss
                loss = (actor_loss + 
                       self.value_coef * value_loss + 
                       self.entropy_coef * entropy_loss)
                
                # Optimize
                self.optimizer.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(self.policy.parameters(), 0.5)
                self.optimizer.step()
                
        # Clear buffers
        self.reset_buffers()
        
        return {
            'actor_loss': actor_loss.item(),
            'value_loss': value_loss.item(),
            'entropy': entropy.mean().item()
        }

# ============================================================================
# TRAINING LOOP
# ============================================================================

def train_car_simulator(num_episodes=1000, render_interval=100):
    """
    Main training function.
    
    Architecture:
    1. Environment generates states from car physics
    2. Agent outputs throttle and steering commands
    3. Car responds with new states
    4. Reward based on speed and track position
    5. PPO updates policy to maximize reward
    """
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Training on {device}")
    
    # Create environment
    env = CarSimulator(track_type='oval', max_steps=1000)
    
    # Get dimensions from environment
    obs = env.reset()
    state_dim = len(obs)
    action_dim = 2  # [throttle, steering]
    
    print(f"State dim: {state_dim}, Action dim: {action_dim}")
    
    # Create agent
    agent = PPOAgent(state_dim, action_dim, device=device)
    
    # Training stats
    episode_rewards = []
    episode_lengths = []
    
    for episode in range(num_episodes):
        state = env.reset()
        episode_reward = 0
        episode_length = 0
        done = False
        
        # Collect rollout
        while not done:
            # Agent selects action
            action, log_prob, value = agent.select_action(state, training=True)
            
            # Environment step
            next_state, reward, done, _ = env.step(action)
            
            # Store transition
            agent.store_transition(state, action, log_prob, reward, value, done)
            
            episode_reward += reward
            episode_length += 1
            state = next_state
            
            # Update policy after collecting enough data or episode end
            if len(agent.states) >= 2048 or done:
                if len(agent.states) > 0:
                    agent.update(next_state, done)
        
        episode_rewards.append(episode_reward)
        episode_lengths.append(episode_length)
        
        # Logging
        if (episode + 1) % 10 == 0:
            avg_reward = np.mean(episode_rewards[-10:])
            avg_length = np.mean(episode_lengths[-10:])
            print(f"Episode {episode+1}: Reward={avg_reward:.1f}, "
                  f"Steps={avg_length:.0f}, Distance={env.total_distance:.1f}m")
            
        # Save checkpoint
        if (episode + 1) % 500 == 0:
            torch.save(agent.policy.state_dict(), f'car_policy_ep{episode+1}.pth')
    
    # Final save
    torch.save(agent.policy.state_dict(), 'car_policy_final.pth')
    
    return agent, episode_rewards

# ============================================================================
# EVALUATION / INFERENCE
# ============================================================================

def evaluate_policy(policy_path='car_policy_final.pth', num_episodes=3):
    """Run trained policy."""
    
    device = torch.device('cpu')
    env = CarSimulator(track_type='oval', max_steps=2000)
    obs = env.reset()
    
    state_dim = len(obs)
    action_dim = 2
    
    # Load policy
    policy = ActorCritic(state_dim, action_dim).to(device)
    policy.load_state_dict(torch.load(policy_path, map_location=device))
    policy.eval()
    
    for ep in range(num_episodes):
        state = env.reset()
        done = False
        total_reward = 0
        step = 0
        
        print(f"\n--- Episode {ep+1} ---")
        
        while not done:
            state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
            with torch.no_grad():
                action, _, _ = policy.get_action(state_tensor, deterministic=True)
            
            state, reward, done, _ = env.step(action.cpu().numpy()[0])
            total_reward += reward
            step += 1
            
            if step % 100 == 0:
                print(f"  Step {step}: pos=({env.position[0]:.1f}, {env.position[1]:.1f}), "
                      f"vel={env.velocity:.1f}m/s")
        
        print(f"Total reward: {total_reward:.1f}, Steps: {step}")

# ============================================================================
# MAIN
# ============================================================================

if __name__ == "__main__":
    # Training
    print("=" * 50)
    print("MOTOR LEARNING: CAR SIMULATOR")
    print("Algorithm: PPO (Proximal Policy Optimization)")
    print("=" * 50)
    
    # Uncomment to train:
    # agent, rewards = train_car_simulator(num_episodes=1000)
    
    # Uncomment to evaluate:
    # evaluate_policy('car_policy_final.pth')
    
    # Quick demo with untrained policy
    print("\nRunning demo with random policy...")
    env = CarSimulator(track_type='oval')
    state = env.reset()
    
    for step in range(200):
        # Random action
        action = np.random.uniform(-0.5, 1.0, size=2)  # slight forward bias
        state, reward, done, _ = env.step(action)
        
        if step % 50 == 0:
            print(f"Step {step}: speed={env.velocity:.1f}m/s, "
                  f"pos=({env.position[0]:.1f}, {env.position[1]:.1f})")
        if done:
            print(f"Episode ended at step {step}")
            break

CSS

Motor car simulator

Recommended