if __name__ == "__main__":
    import numpy as np
    from ch25_deep_rl import SimpleCartPole  # reuse Chapter 25's env

    torch.manual_seed(0)
    env = SimpleCartPole()
    policy = PolicyNetwork(state_dim=4, n_actions=2)
    optimizer = torch.optim.Adam(policy.parameters(), lr=1e-2)

    returns_history = []
    for episode in range(300):
        ep_return = reinforce_episode(env, policy, optimizer)
        returns_history.append(ep_return)

    print("Avg return, first 20 episodes:",
          np.mean(returns_history[:20]))
    print("Avg return, last 20 episodes: ",
          np.mean(returns_history[-20:]))
