Tensorflowのサンプル試し（倒立振子問題）

■ChatGPTでサンプルを生成して動作確認する。
前回Tensorflow2.13の環境を整えたので、いくつかのテーマでChatGPTからサンプル生成して動作確認していきたい。なお、tf_agentsは強化学習のライブラリとのこと。

深層強化学習の簡単なサンプルを求めたところ、下のコードが提案された。
倒立振子のコードとのこと（ほうきを手にのせて前後でバランスを取るようなやつ）。

このコードの意味を学んだりパラメータを変えて動作確認したりを考えているけど、ChatGPTを使うと簡単にサンプルも出すことができるので、勉強の仕方や方向性について何か考えないといけない気がする。

import tensorflow as tf
from tf_agents.environments import suite_gym
from tf_agents.environments import tf_py_environment
from tf_agents.networks import q_network
from tf_agents.agents.dqn import dqn_agent
from tf_agents.replay_buffers import tf_uniform_replay_buffer
from tf_agents.utils import common
from tf_agents.trajectories import trajectory
import numpy as np

# 環境セットアップ
env_name = 'CartPole-v1'
train_py_env = suite_gym.load(env_name)
eval_py_env = suite_gym.load(env_name)
train_env = tf_py_environment.TFPyEnvironment(train_py_env)
eval_env = tf_py_environment.TFPyEnvironment(eval_py_env)

# Qネットワーク構築
fc_layer_params = (100,)
q_net = q_network.QNetwork(
    train_env.observation_spec(),
    train_env.action_spec(),
    fc_layer_params=fc_layer_params)

# DQNエージェント初期化
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
train_step_counter = tf.Variable(0)

agent = dqn_agent.DqnAgent(
    train_env.time_step_spec(),
    train_env.action_spec(),
    q_network=q_net,
    optimizer=optimizer,
    td_errors_loss_fn=common.element_wise_squared_loss,
    train_step_counter=train_step_counter)

agent.initialize()

# リプレイバッファ準備
replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
    data_spec=agent.collect_data_spec,
    batch_size=train_env.batch_size,
    max_length=10000)

# データ収集関数
def collect_step(environment, policy, buffer):
    time_step = environment.current_time_step()
    action_step = policy.action(time_step)
    next_time_step = environment.step(action_step.action)
    traj = trajectory.from_transition(time_step, action_step, next_time_step)
    buffer.add_batch(traj)

# 初期データ収集（100ステップ）
random_policy = agent.collect_policy
print("初期データ収集開始...")
for _ in range(100):
    collect_step(train_env, random_policy, replay_buffer)
print("初期データ収集完了！")

# データセット作成
dataset = replay_buffer.as_dataset(
    num_parallel_calls=3,
    sample_batch_size=64,
    num_steps=2).prefetch(3)
iterator = iter(dataset)

# トレーニングループ
num_iterations = 10000
print("トレーニング開始...")
for _ in range(num_iterations):
    experience, _ = next(iterator)
    train_loss = agent.train(experience).loss
    step = train_step_counter.numpy()

    if step % 100 == 0:
        print(f"Step {step}: 損失 = {train_loss:.4f}")

print("トレーニング完了！")

# 学習済みポリシーのテスト
num_episodes = 5
print("\n評価フェーズ開始...")
for episode in range(1, num_episodes + 1):
    time_step = eval_env.reset()
    episode_reward = 0
    step_count = 0
    while not time_step.is_last():
        action_step = agent.policy.action(time_step)
        time_step = eval_env.step(action_step.action)
        episode_reward += time_step.reward.numpy().item()
        step_count += 1
    print(f"エピソード {episode}：報酬 = {episode_reward:.2f}, ステップ数 = {step_count}")
print("評価フェーズ終了！")