CartPole环境运行PPO算法时遇AttributeError及解包错误的解决求助

我正在跑Keras官网那个CartPole的PPO示例代码，结果训练环节接连踩了两个坑，折腾半天没搞定，来请教下大家怎么解决比较稳妥：

• 第一个坑：AttributeError
  代码里有一行 observation = observation.reshape(1,-1)，运行直接报错'tuple' object has no attribute 'reshape'。查了下才发现，新版本Gym里env.reset()返回的不是单纯的观察数组，而是**(初始观察数组, 空info字典)**这样的元组，自然没法直接调用reshape方法。

• 第二个坑：解包错误
  我试着改成只取元组里的数组部分，比如用observation[0].reshape(1,-1)或者直接把env.reset()[0]赋值给observation，但运行到observation_new, reward, done, _ = env.step(action[0].numpy())这一行时又炸了，报错too many values to unpack (expected 4)。

我不想大改原有代码的逻辑，请问有没有办法既能解决这两个问题，又尽量保持原代码的结构？

下面是我整理的最小可复现代码：

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import gym
import scipy.signal

env = gym.make("CartPole-v0")

steps_per_epoch = 4000
epochs = 30
hidden_sizes = (64, 64)
observation_dimensions = env.observation_space.shape[0]
num_actions = env.action_space.n

observation_input = keras.Input(shape=(observation_dimensions,), dtype=tf.float32)
logits = mlp(observation_input, list(hidden_sizes) + [num_actions], tf.tanh, None)
actor = keras.Model(inputs=observation_input, outputs=logits)

# 这里的env.reset()返回元组，是第一个问题的源头
observation, episode_return, episode_length = env.reset(), 0, 0

for epoch in range(epochs):
    # Initialize the sum of the returns, lengths and number of episodes
    sum_return = 0
    sum_length = 0
    num_episodes = 0

    for t in range(steps_per_epoch):
        # if render:
        #     env.render()

        # 这里调用reshape报错，因为observation是元组
        observation = observation.reshape(1, -1)
        logits, action = sample_action(observation)
        # 这里解包报错，因为env.step返回的是(obs, reward, terminated, truncated, info)
        observation_new, reward, done, _ = env.step(action[0].numpy())

        episode_return += reward
        episode_length += 1

        # Get the value and log-probability of the action
        value_t = critic(observation)
        logprobability_t = logprobabilities(logits, action)

        # Store obs, act, rew, v_t, logp_pi_t
        buffer.store(observation, action, reward, value_t, logprobability_t)

        # Update the observation
        observation = observation_new

        # Finish trajectory if reached to a terminal state
        terminal = done
        if terminal or (t == steps_per_epoch - 1):
            last_value = 0 if done else critic(observation.reshape(1, -1))
            buffer.finish_trajectory(last_value)
            sum_return += episode_return
            sum_length += episode_length
            num_episodes += 1
            # 这里同样是env.reset()返回元组
            observation, episode_return, episode_length = env.reset(), 0, 0

# 依赖的辅助函数
def mlp(x, sizes, activation=tf.tanh, output_activation=None):
    # Build a feedforward neural network
    for size in sizes[:-1]:
        x = layers.Dense(units=size, activation=activation)(x)
    return layers.Dense(units=sizes[-1], activation=output_activation)(x)

@tf.function
def sample_action(observation):
    logits = actor(observation)
    action = tf.squeeze(tf.random.categorical(logits, 1), axis=1)
    return logits, action

备注：内容来源于stack exchange，提问作者Max