隠れマルコフモデルのパラメータ推定が上手くいかない件について

推定された遷移確率:
[[1. 0.]
 [0. 0.]]
推定された出力確率:
[[0.18333333 0.13333333 0.17333333 0.16333333 0.15       0.19666667]
 [0.         0.         0.         0.         0.         0.        ]]

import numpy as np
import matplotlib.pyplot as plt

class HMM:
    def __init__(self, states, observations):
        self.states = states
        self.observations = observations
        self.initial_prob = None
        self.transition_prob = None
        self.emission_prob = None

    def _init_params(self):
        n_states = len(self.states)
        n_observations = len(self.observations)
        self.initial_prob = np.array([1.0, 0])  # 初期状態を指定
        self.transition_prob = np.array([[0.99, 0.01], [0.05, 0.95]])  # 遷移確率を指定
        self.emission_prob = np.array([[1/6, 1/6, 1/6, 1/6, 1/6, 1/6], [0, 0, 0, 1/3, 1/3, 1/3]])

    def _estimate_params(self, hidden_states, observations):
        n_states = len(self.states)
        n_observations = len(self.observations)
        n_timesteps = len(hidden_states)

        # パラメータを初期化
        estimated_initial_prob = np.zeros(n_states)
        estimated_transition_prob = np.zeros((n_states, n_states))
        estimated_emission_prob = np.zeros((n_states, n_observations))

        # 初期確率の推定
        estimated_initial_prob[self.states.index(hidden_states[0])] = 1

        # 遷移確率の推定
        for t in range(n_timesteps - 1):
            current_state = self.states.index(hidden_states[t])
            next_state = self.states.index(hidden_states[t + 1])
            estimated_transition_prob[current_state, next_state] += 1

        # 出力確率の推定
        for t in range(n_timesteps):
            state = self.states.index(hidden_states[t])
            observation = self.observations.index(observations[t])
            estimated_emission_prob[state, observation] += 1

        estimated_initial_prob /= np.sum(estimated_initial_prob) + 1e-10
        estimated_transition_prob /= np.sum(estimated_transition_prob, axis=1, keepdims=True) + 1e-10
        estimated_emission_prob /= np.sum(estimated_emission_prob, axis=1, keepdims=True) + 1e-10

        return estimated_initial_prob, estimated_transition_prob, estimated_emission_prob

    def baum_welch(self, hidden_states, observations, n_iter=300):
        self._init_params()

        for _ in range(n_iter):
            # フォワードステップ
            forward_prob = self.forward(hidden_states, observations)

            # バックワードステップ
            backward_prob = self.backward(hidden_states, observations)

            # 推定されたパラメータの更新
            estimated_initial_prob, estimated_transition_prob, estimated_emission_prob = self._estimate_params(hidden_states, observations)

            # パラメータの更新
            self.initial_prob = estimated_initial_prob
            self.transition_prob = estimated_transition_prob
            self.emission_prob = estimated_emission_prob

        return self.initial_prob, self.transition_prob, self.emission_prob

    def forward(self, hidden_states, observations):
        n_timesteps = len(observations)
        n_states = len(self.states)
        forward_prob = np.zeros((n_timesteps, n_states))

        forward_prob[0, :] = self.initial_prob * self.emission_prob[:, self.observations.index(observations[0])]

        for t in range(1, n_timesteps):
            for j in range(n_states):
                forward_prob[t, j] = np.sum(forward_prob[t - 1, :] * self.transition_prob[:, j]) * self.emission_prob[j, self.observations.index(observations[t])]

        return forward_prob

    def backward(self, hidden_states, observations):
        n_timesteps = len(observations)
        n_states = len(self.states)
        backward_prob = np.zeros((n_timesteps, n_states))

        backward_prob[n_timesteps - 1, :] = 1

        for t in reversed(range(n_timesteps - 1)):
            for i in range(n_states):
                backward_prob[t, i] = np.sum(backward_prob[t + 1, :] * self.transition_prob[i, :] * self.emission_prob[:, self.observations.index(observations[t + 1])])

        return backward_prob

# パラメータの真の値
states = ['Fair', 'Biased']
observations = [1, 2, 3, 4, 5, 6]
initial_prob = np.array([1.0, 0.0])
transition_prob = np.array([[0.99, 0.01], [0.05, 0.95]])
emission_prob = np.array([[1/6, 1/6, 1/6, 1/6, 1/6, 1/6], [0, 0, 0, 1/3, 1/3, 1/3]])

# HMMのインスタンス化
hmm = HMM(states, observations)

# 隠れ状態の生成
hidden_states = []
observed_states = []

for _ in range(300):
    state = np.random.choice(states, p=initial_prob)
    # 観測結果の生成（修正後）
    observation = np.random.choice(observations, p=emission_prob[states.index(state)].tolist())
    hidden_states.append(state)
    observed_states.append(observation)

# 推定されたパラメータの取得
estimated_initial_prob, estimated_transition_prob, estimated_emission_prob = hmm.baum_welch(hidden_states, observed_states, n_iter=100)

print("推定された遷移確率:")
print(estimated_transition_prob)
print("推定された出力確率:")
print(estimated_emission_prob)

# 隠れ状態の推移の可視化
fig, ax = plt.subplots(figsize=(8, 6))
x = range(len(hidden_states))
y = [states.index(state) for state in hidden_states]
ax.plot(x, y, label='Hidden States')
ax.set_xticks(x)
ax.set_xticklabels(hidden_states)
ax.set_xlabel('Time Step')
ax.set_ylabel('Hidden State')
ax.legend()
plt.show()

# 観測結果の推移の可視化
fig, ax = plt.subplots(figsize=(8, 6))
x = range(len(observed_states))
ax.plot(x, observed_states, label='Observed States')
ax.set_xticks(x)
ax.set_xticklabels(hidden_states)
ax.set_xlabel('Time Step')
ax.set_ylabel('Observed State')
ax.legend()
plt.show()