结合最大熵原理的逆向强化学习算法及其在AI Agent中的应用

import numpy as np
import gym
from gym.envs.toy_text import discrete

class GridWorld(discrete.DiscreteEnv):
    def __init__(self, shape=(5,5), goal=(4,4)):
        nS = np.prod(shape)
        nA = 4  # 上下左右
        P = {s: {a: [] for a in range(nA)} for s in range(nS)}

        def to_s(row, col): return row * shape[1] + col
        def inc(row, col, a):
            if a == 0: row = max(row - 1, 0)
            elif a == 1: row = min(row + 1, shape[0]-1)
            elif a == 2: col = max(col - 1, 0)
            elif a == 3: col = min(col + 1, shape[1]-1)
            return row, col

        for row in range(shape[0]):
            for col in range(shape[1]):
                s = to_s(row, col)
                for a in range(nA):
                    newrow, newcol = inc(row, col, a)
                    s_ = to_s(newrow, newcol)
                    done = (newrow, newcol) == goal
                    rew = 1.0 if done else 0.0
                    P[s][a].append((1.0, s_, rew, done))

        isd = np.zeros(nS)
        isd[0] = 1.0
        super(GridWorld, self).__init__(nS, nA, P, isd)

def maxent_irl(feature_matrix, P_a, gamma, trajs, lr, n_iters):
    n_states, d_states = feature_matrix.shape
    theta = np.random.uniform(size=(d_states,))

    def get_policy(theta):
        rewards = feature_matrix.dot(theta)
        V = np.zeros(n_states)
        for _ in range(20):
            V = np.log(np.sum(np.exp(rewards + gamma * P_a.dot(V).reshape(n_states, -1)), axis=1))
        return np.exp(rewards + gamma * P_a.dot(V).reshape(n_states, -1) - V[:, None])

    def expected_feature_counts(policy):
        mu = np.zeros((n_states,))
        mu[0] = 1.0
        for _ in range(100):
            mu = mu.dot(policy)
        return mu.dot(feature_matrix)

    empirical_fc = np.zeros(d_states)
    for traj in trajs:
        for state, _ in traj:
            empirical_fc += feature_matrix[state]
    empirical_fc /= len(trajs)

    for _ in range(n_iters):
        policy = get_policy(theta)
        exp_fc = expected_feature_counts(policy)
        grad = empirical_fc - exp_fc
        theta += lr * grad

    return feature_matrix.dot(theta)

# 特征矩阵 (one-hot 状态特征)
n_states = 25
feature_matrix = np.eye(n_states)

# 随机生成专家轨迹 (假设专家能到达终点)
trajs = [[(0,0), (1,0), (2,0), (3,0), (4,0), (4,1), (4,2), (4,3), (4,4)]]

# 状态转移概率矩阵
P_a = np.zeros((n_states, 4, n_states))
# ... 构建P_a (略，参考环境定义)

rewards = maxent_irl(feature_matrix, P_a, 0.9, trajs, lr=0.01, n_iters=100)
print("Recovered rewards:", rewards.reshape((5,5)))