hackathonf23-Stacks/leaky_relu_runner.py

"""
This is the sixth ablation study, using Leaky ReLu instead of ReLu.
"""
import numpy as np
import pandas as pd
import tensorflow as tf
from keras.layers import Dense
from sklearn.preprocessing import MinMaxScaler
import dagshub
import mlflow
import mlflow.keras
import pickle
import matplotlib.pyplot as plt

mlflow.set_tracking_uri('https://dagshub.com/ML-Purdue/hackathonf23-Stacks.mlflow')
dagshub.init(repo_owner='ML-Purdue', repo_name='hackathonf23-Stacks', mlflow=True)

def get_or_create_experiment_id(name):
    exp = mlflow.get_experiment_by_name(name)
    if exp is None:
        exp_id = mlflow.create_experiment(name)
        return exp_id
    return exp.experiment_id
class MaxEntIRL:
    def __init__(self, state_dim):
        self.state_dim = state_dim
        self.model = self._create_irl_model()

    def _create_irl_model(self):
        model = tf.keras.Sequential([
            Dense(self.state_dim, input_shape=(self.state_dim,), activation=tf.nn.leaky_relu),
            Dense(4096, activation=tf.nn.leaky_relu),
            Dense(2048, activation=tf.nn.leaky_relu),
            Dense(self.state_dim, activation='linear')
        ])
        return model

    def generateHumanTrajectories(self, num_trajectories, trajectory_length):
        human_trajectories = []
        for _ in range(num_trajectories):
            trajectory = []
            state = np.zeros(self.state_dim)

            for _ in range(trajectory_length):
                direction_probabilities = self._generate_direction_probabilities()  # Get direction probabilities
                action_coefficients = np.random.choice([-1, 0, 1], p=direction_probabilities)
                action = action_coefficients * 0.1
                new_state = state + action
                trajectory.append((state, action))
                state = new_state

            human_trajectories.append(trajectory)

        return human_trajectories

    def _generate_direction_probabilities(self):
        probabilities = np.random.dirichlet(np.ones(self.state_dim) * 0.1)
        return probabilities

    def loadDataset(self, file_path):
        data = pd.read_csv(file_path)  # Load CSV data
        scaler = MinMaxScaler()
        columns_to_normalize = ['position x [mm]', 'position y [mm]', 'position z (height) [mm]', 'velocity [mm/s]']
        data[columns_to_normalize] = scaler.fit_transform(data[columns_to_normalize])
        return data

    def train_irl_with_dataset(self, data, lr=0.001, epochs=3):
        state_dim = self.state_dim
        optimizer = tf.keras.optimizers.Adam(learning_rate=lr)

        # Extract relevant columns from the loaded dataset
        positions = data[['position x [mm]', 'position y [mm]', 'position z (height) [mm]']].values
        velocities = data['velocity [mm/s]'].values
        mlflow.tensorflow.autolog()
        with mlflow.start_run(experiment_id=get_or_create_experiment_id("Leaky ReLu")):
            for epoch in range(epochs):
                total_loss = 0
                state_frequencies = self._calculate_state_frequencies(positions)

                for idx in range(len(positions)):
                    state = positions[idx]
                    velocity = velocities[idx]

                    with tf.GradientTape() as tape:
                        preferences = self.model(state[np.newaxis, :])
                        prob_human = tf.nn.softmax(preferences)

                        # Define losses
                        max_entropy_loss = -tf.reduce_sum(prob_human * tf.math.log(prob_human + 1e-8), axis=1)
                        alignment_loss = -tf.reduce_sum(state_frequencies * tf.math.log(prob_human + 1e-8), axis=1)
                        maxent_irl_objective = max_entropy_loss + alignment_loss

                        # Compute the gradients
                        grads = tape.gradient(maxent_irl_objective, self.model.trainable_variables)
                        optimizer.apply_gradients(zip(grads, self.model.trainable_variables))
                        total_loss += tf.reduce_sum(maxent_irl_objective)  # Accumulate the total loss

                avg_loss = total_loss / len(positions)
                mlflow.log_metric(f"loss", avg_loss, step=epoch)
                print(f"Epoch {epoch + 1}/{epochs}, MaxEnt IRL Loss: {avg_loss}")

    def train_irl(self, human_trajectories=None, data=None, use_dataset=False, lr=0.001, epochs=3):
        if use_dataset and data is not None:
            # Train using the loaded dataset
            self.train_irl_with_dataset(data, lr=lr, epochs=epochs)
        else:
            # Train using the generative function
            if human_trajectories is None:
                human_trajectories = self.generateHumanTrajectories(num_trajectories, trajectory_length)

            self._train_irl_generative(human_trajectories, lr=lr, epochs=epochs)

    def _train_irl_generative(self, human_trajectories, lr=0.001, epochs=3):
        trajectory_length = len(human_trajectories[0])
        optimizer = tf.keras.optimizers.Adam(learning_rate=lr)

        for epoch in range(epochs):
            total_loss = 0
            state_frequencies = self._calculate_state_frequencies(human_trajectories, trajectory_length)

            for trajectory in human_trajectories:
                for state, _ in trajectory:
                    with tf.GradientTape() as tape:
                        preferences = self.model(state[np.newaxis, :])
                        prob_human = tf.nn.softmax(preferences)

                        # Inside the training loop:
                        max_entropy_loss = -tf.reduce_sum(prob_human * tf.math.log(prob_human + 1e-8), axis=1)
                        alignment_loss = -tf.reduce_sum(state_frequencies * tf.math.log(prob_human + 1e-8), axis=1)
                        maxent_irl_objective = max_entropy_loss + alignment_loss

                    grads = tape.gradient(maxent_irl_objective, self.model.trainable_variables)
                    optimizer.apply_gradients(zip(grads, self.model.trainable_variables))
                    total_loss += maxent_irl_objective

            avg_loss = total_loss / (len(human_trajectories) * trajectory_length)
            print(f"Epoch {epoch + 1}/{epochs}, MaxEnt IRL Loss: {avg_loss}")

    def _calculate_state_frequencies(self, positions):
        state_counts = np.sum(positions, axis=0)
        state_frequencies = state_counts / (len(positions) * self.state_dim)
        return state_frequencies
    def save_model(self, file_path):
        model_config = self.model.get_config()
        with open(file_path, 'wb') as f:
            pickle.dump(model_config, f)

    @classmethod
    def load_model(cls, file_path, state_dim):
        with open(file_path, 'rb') as f:
            model_config = pickle.load(f)
            irl_instance = cls(state_dim)
            irl_instance.model = tf.keras.Sequential.from_config(model_config)
        return irl_instance
    # Indicate test completion status
    
state_dim = 3  # Dimension of the state space
irl = MaxEntIRL(state_dim)
num_trajectories = 100
trajectory_length = 20

model_path = '/Users/vinay/Desktop/Computer_Science_Projects/ReScience/hackathonf23-Stacks/models/leaky_relu_model.pkl'

# Define the state dimension for the model
state_dim = 3  # Assuming state_dim is 3

# Load the model using your custom method
def load_model(file_path, state_dim):
    with open(file_path, 'rb') as file:
        model_config = pickle.load(file)
        irl_instance = MaxEntIRL(state_dim)
        irl_instance.model = tf.keras.Sequential.from_config(model_config)
    return irl_instance

# Load the model using the custom load_model function
loaded_model = load_model(model_path, state_dim)

# Load the test data from the "test.csv" file
test_data_path = '/Users/vinay/Desktop/Computer_Science_Projects/ReScience/hackathonf23-Stacks/data/test.csv'
data = pd.read_csv(test_data_path)

# Data preprocessing
scaler = MinMaxScaler()
columns_to_normalize = ['position x [mm]', 'position y [mm]', 'position z (height) [mm]']
columns_to_normalize = scaler.fit_transform(data[columns_to_normalize])

# Make predictions using the loaded model
predictions = loaded_model.model.predict(columns_to_normalize)
model_loss = ((predictions - columns_to_normalize))
magnitudes = np.linalg.norm(model_loss, axis=1)
magnitude_columns_to_normalize = np.linalg.norm(columns_to_normalize, axis=1)
percent_error_magnitudes = np.abs(magnitudes - magnitude_columns_to_normalize) / np.abs(magnitude_columns_to_normalize) * 100
avg_displacement_error = np.mean(magnitudes)
print(f"Average Displacement Error with Leaky ReLu Model: {avg_displacement_error:.2f}")


plt.figure(figsize=(8, 6))
plt.boxplot(magnitudes, vert=False)
plt.title('Error of Trajectory Magnitudes with Leaky ReLu Ablation:')
plt.xlabel('Magnitude (meters)')
plt.yticks([])
plt.show()