MaxEnt IRL için aşağıdaki kodu hazırladım. Her adımı detaylı olarak kontrol edebilir misin?

from scipy.optimize import minimize
from scipy.linalg import solve
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import time

def maxentdefaultparams(algorithm_params):
    """
    Fill in default parameters for the MaxEnt algorithm.
    """
    # Create default parameters.
    default_params = {
        'seed': 0,
        'laplace_prior': 0,
        'all_features': 1,
        'true_features': 0
    }
    # Set parameters.
    algorithm_params = filldefaultparams(algorithm_params, default_params)
    return algorithm_params

def filldefaultparams(algorithm_params, default_params):
    """
    Fill in default parameters for any missing parameters in algorithm_params.
    """
    if algorithm_params is None:
        algorithm_params = {}
    for key in default_params:
        if key not in algorithm_params:
            algorithm_params[key] = default_params[key]
    return algorithm_params

def maxentsoftmax(q):
    """
    Compute row-wise softmax of a matrix q.
    Returns v, where v(s) = log(sum_a exp(q(s,a)))
    """
    max_q = np.max(q, axis=1, keepdims=True)
    v = max_q + np.log(np.sum(np.exp(q - max_q), axis=1, keepdims=True))
    return v.ravel()  # Return as 1D array of length states

def linearvalueiteration(mdp_data, r_sa, max_iterations=1000, tolerance=1e-4):
    """
    Perform linear value iteration for MaxEnt IRL.

    mdp_data: dict containing MDP information
    r_sa: reward for each state-action pair (S x A)
    """
    states = mdp_data['states']
    actions = mdp_data['actions']
    gamma = mdp_data['discount']

    # Initialize V(s)
    V = np.zeros(states)

    # Precompute transition probabilities
    P = mdp_data['sa_p']    # shape (S, A, S)
    policy = np.full((states, actions), 1.0 / actions)  # Uniform policy initialization

    for iteration in range(max_iterations):
        V_old = V.copy()
        Q = np.zeros((states, actions))
        for a in range(actions):
            # For each action, compute Q(s,a)
            Q[:, a] = r_sa[:, a] + gamma * P[:, a, :].dot(V)
        # Update V(s)
        V = maxentsoftmax(Q)
        # Check for convergence
        if np.max(np.abs(V - V_old)) < tolerance:
            break

    # Compute policy
    log_policy = Q - V[:, np.newaxis]
    policy = np.exp(log_policy)
    # Normalize policy
    policy = policy / np.sum(policy, axis=1, keepdims=True)
    return policy, log_policy

def linearmdpfrequency(mdp_data, policy_data, initD):
    """
    Compute state visitation frequencies D for MaxEnt IRL.

    mdp_data: dict with MDP information
    policy_data: dict containing 'p', the policy (S x A)
    initD: initial state distribution (S x 1)
    Returns D: state visitation frequencies (S x 1)
    """
    states = mdp_data['states']
    gamma = mdp_data['discount']

    P = mdp_transition_matrix(mdp_data, policy_data['p'])

    # Compute D
    # D = (I - gamma * P.T)^{-1} * initD
    I = np.eye(states)
    A = I - gamma * P.T
    D = solve(A, initD)
    return D

def mdp_transition_matrix(mdp_data, policy):
    """
    Compute the transition matrix P_pi under a given policy.

    mdp_data: dict with MDP information
    policy: policy matrix (S x A)
    Returns P_pi: transition matrix (S x S)
    """
    states = mdp_data['states']
    actions = mdp_data['actions']

    P_pi = np.zeros((states, states))

    for s in range(states):
        for a in range(actions):
            P_pi[s, :] += policy[s, a] * mdp_data['sa_p'][s, a, :]
    return P_pi

def maxentdiscounted(r, F, muE, mu_sa, mdp_data, initD, laplace_prior=0):
    """
    Compute MaxEnt objective and gradient. Discounted infinite-horizon version.
    """
    actions = mdp_data['actions']

    # Convert to full reward
    wts = r.copy()
    r_full = F @ r

    # Compute policy and log policy
    r_sa = np.tile(r_full.reshape(-1,1), (1, actions))
    policy, log_policy = linearvalueiteration(mdp_data, r_sa)

    # Compute value by adding up log example probabilities
    val = np.sum(log_policy * mu_sa)

    # Add Laplace prior for regularization
    if laplace_prior:
        val -= laplace_prior * np.sum(np.abs(wts))

    # Invert for descent (since we are minimizing)
    val = -val

    # Compute gradient
    D = linearmdpfrequency(mdp_data, {'p': policy}, initD)

    # Compute gradient
    dr = muE - F.T @ D

    # Laplace prior
    if laplace_prior:
        dr -= laplace_prior * np.sign(wts)

    # Invert for descent
    dr = -dr

    return val, dr

def maxentrun(algorithm_params, mdp_data, mdp_model,
              feature_data, example_samples, true_features, verbosity):
    """
    Ziebart's Maximum Entropy IRL, with optional prior on reward values.

    Args:
        algorithm_params: parameters of the MaxEnt algorithm
        mdp_data: MDP definition
        mdp_model: MDP model (not used in this implementation)
        feature_data: feature data
        example_samples: list of example trajectories, each a list of (state, action) tuples
        true_features: true feature matrix (if available)
        verbosity: verbosity level

    Returns:
        irl_result: result of the IRL algorithm
    """
    # Fill in default parameters
    algorithm_params = maxentdefaultparams(algorithm_params)

    np.random.seed(algorithm_params['seed'])

    # Initialize variables
    sa_p = mdp_data['sa_p']  # Transition probabilities (S x A x S)
    states, actions, _ = sa_p.shape

    N = len(example_samples)
    T = len(example_samples[0])

    # Build feature membership matrix
    if algorithm_params['all_features']:
        F = feature_data['splittable']  # Assuming splittable is a numpy array
        # Add dummy feature
        F = np.hstack([F, np.ones((states, 1))])
    elif algorithm_params['true_features']:
        F = true_features
    else:
        F = np.eye(states)
    # Count features
    features = F.shape[1]

    # Compute feature expectations
    muE = np.zeros(features)
    mu_sa = np.zeros((states, actions))
    ex_s = np.zeros((N, T), dtype=int)
    ex_a = np.zeros((N, T), dtype=int)
    for i in range(N):
        for t in range(T):
            s = example_samples[i][t][0]
            a = example_samples[i][t][1]
            ex_s[i,t] = s
            ex_a[i,t] = a
            mu_sa[s,a] += 1
            muE += F[s]

    # Normalize mu_sa
    mu_sa /= N * T

    # Generate initial state distribution for infinite horizon
    initD = np.bincount(ex_s[:,0], minlength=states).astype(float)
    initD /= initD.sum()

    # Objective function
    def fun(r):
        val, grad = maxentdiscounted(r, F, muE, mu_sa, mdp_data, initD, algorithm_params['laplace_prior'])
        return val, grad

    # Optimization options
    if verbosity == 0:
        opt_options = {'disp': False, 'maxiter': 1000}
    else:
        opt_options = {'disp': True, 'maxiter': 1000}

    start_time = time.time()

    # Initialize reward parameters
    r0 = np.random.rand(features)

    # Run unconstrained minimization
    res = minimize(fun, r0, method='L-BFGS-B', jac=True, options=opt_options)

    # Get timing
    time_elapsed = time.time() - start_time
    if verbosity != 0:
        print('Optimization completed in {:.2f} seconds.'.format(time_elapsed))

    # Get final reward weights
    r = res.x

    # Convert to full tabulated reward
    wts = r.copy()
    r_full = F @ r

    # Return corresponding reward function
    r_sa = np.tile(r_full.reshape(-1,1), (1, actions))
    soln = mdp_solve(mdp_data, r_sa)
    v = soln['v']
    q = soln['q']
    p = soln['p']

    # Construct returned structure
    irl_result = {
        'r': r_sa,
        'v': v,
        'p': p,
        'q': q,
        'r_itr': [r_sa],
        'model_itr': [wts],
        'model_r_itr': [r_sa],
        'p_itr': [p],
        'model_p_itr': [p],
        'time': time_elapsed
    }

    return irl_result

def mdp_solve(mdp_data, r_sa, max_iterations=1000, tolerance=1e-4):
    """
    Solve the MDP given the reward r_sa (S x A).
    Returns solution containing 'v', 'q', 'p'.
    """
    states = mdp_data['states']
    actions = mdp_data['actions']
    gamma = mdp_data['discount']

    # Use standard value iteration
    V = np.zeros(states)
    P = mdp_data['sa_p']  # Transition probabilities (S x A x S)
    for iteration in range(max_iterations):
        V_old = V.copy()
        Q = np.zeros((states, actions))
        for a in range(actions):
            Q[:, a] = r_sa[:, a] + gamma * P[:, a, :].dot(V)
        V = np.max(Q, axis=1)
        if np.max(np.abs(V - V_old)) < tolerance:
            break

    # Compute policy
    policy = np.zeros((states, actions))
    for s in range(states):
        best_actions = np.argwhere(Q[s, :] == np.max(Q[s, :])).flatten()
        policy[s, best_actions] = 1.0 / len(best_actions)

    soln = {'v': V, 'q': Q, 'p': policy}
    return soln

def run_maxent():
    # Run the MaxEnt IRL algorithm
    irl_result = maxentrun(algorithm_params, mdp_data, None, feature_data,
                           example_samples, true_features, verbosity)

    # Extract reward per state
    reward_function = irl_result['r']  # Shape (S, A)
    # Assuming rewards are the same across actions
    reward_per_state = reward_function[:, 0]
    # Extract policy
    policy = irl_result['p']

    # Plot the inferred reward function as a bar chart
    plt.figure(figsize=(10, 6))
    plt.bar(range(len(reward_per_state)), reward_per_state, color='skyblue', edgecolor='black')
    plt.title('Inferred Reward Function')
    plt.xlabel('State Index')
    plt.ylabel('Reward Value')
    plt.grid(axis='y')
    plt.tight_layout()
    plt.show()

    # Plot the optimal policy heatmap
    plt.figure(figsize=(10, 8))
    ax = sns.heatmap(policy, annot=True, cmap='coolwarm', cbar_kws={'label': 'Probability'},
                     xticklabels=['Action {}'.format(i) for i in range(policy.shape[1])],
                     yticklabels=['State {}'.format(i) for i in range(policy.shape[0])])
    ax.invert_yaxis()  # To have state index 0 at the bottom
    plt.title('Optimal Policy Heatmap Derived from MaxEnt IRL')
    plt.xlabel('Action')
    plt.ylabel('State')
    plt.tight_layout()
    plt.show()

    # Print the inferred reward function and policy
    print("Inferred reward function (r):")
    print(reward_function)
    print("\nValue function (v):")
    print(irl_result['v'])
    print("\nPolicy (p):")
    print(irl_result['p'])