A2C с NARX‑Critic

A2C (Advantage Actor‑Critic) использует актёра для выбора действий и критика для оценки состояний. В нашей реализации критик — NARX (Nonlinear AutoRegressive with eXogenous inputs), что позволяет лучше моделировать динамику и историю за счёт явного учёта прошедших состояний.

Компоненты

• Актор: гауссовская политика \(\pi_\theta(a|s) = \mathcal{N}(\mu_\theta(s), \sigma_\theta^2)\); параметры — Actor (PyTorch)
• Критик (NARX): оценка \(V(s)\) на основе расширенного входа (текущее состояние + предыдущие сигналы); классы Critic (A2C) и NARX (модульная NARX‑сеть)
• Сбор опыта: Runner собирает траектории, клиппирует действия под action_space
• Обучение: A2CLearner.learn — обновления актёра/критика со стабилизацией (клиппинг градиента, энтропия)

Теория (на базе реализации)

• Дисконтированные возвраты:

\[ G_t = \sum_{k=0}^{\infty} \gamma^k r_{t+k}, \quad G^{\text{episodic}}_t = r_t + \gamma (1-\text{done}) G_{t+1} \]

• Advantage в коде:

\[ A_t = \begin{cases} r_t, & \text{если discount\_rewards=True (чисто по возвратам)}\\ r_t + \gamma V(s_{t+1}) - V(s_t), & \text{иначе (TD‑таргет)} \end{cases} \]

и далее \(A_t = \text{td\_target} - V(s_t)\).

• Потери:

\[ \mathcal{L}_\text{actor} = -\,\mathbb{E}[\log \pi_\theta(a_t|s_t)\, A_t] - \beta\,\mathbb{E}[\mathcal{H}[\pi_\theta(\cdot|s_t)]] \]

\[ \mathcal{L}_\text{critic} = \mathbb{E}\big[(\text{td\_target} - V_\phi(s_t))^2\big] \]

NARX как критик

NARX‑сеть (tensoraerospace/agent/narx/model.py) явно использует прошлый выход/состояние в качестве входа для предсказания следующего. В A2C вход критика формируется как конкатенация текущего состояния и предыдущего состояния (см. process_memory_narx — формирование critic_states). Это повышает качество оценок \(V(s)\) для систем с выраженной динамической памятью.

Идентификация (упрощённо):

• Обучение NARX сводится к минимизации MSE между предсказанным и целевым выходом по последовательности (см. NARX.train).
• В A2C критик обучается по MSE между таргетом (возвраты или TD) и текущей оценкой \(V(s)\).

Обучение (контур)

1. Runner.run собирает пары \((s_t, a_t, r_t, s_{t+1}, done_t)\), клиппирует действия под action_space.
2. process_memory_narx формирует тензоры: действия, вознаграждения (опц. дисконтируются), состояния, следующие состояния, флаги завершения и critic_states = [s_t, s_{t-1}].
3. A2CLearner.learn:
4. Если discount_rewards=True, таргет критика td_target = rewards (возвраты); иначе r + γ V(s').
5. Advantage: advantage = td_target - V(s).
6. Actor: лог‑вероятности из Normal(mean,std), энтропия; обновление с клиппингом градиента.
7. Critic: MSE; обновление с клиппингом градиента.
8. Логирование в TensorBoard (losses, градиенты/параметры, награды).

Быстрый старт

import gymnasium as gym
import torch
from tensoraerospace.agent.a2c.narx import Actor, Critic, A2CLearner, Runner

env = gym.make('LinearLongitudinalF16-v0', number_time_steps=2000)
actor = Actor(state_dim=env.observation_space.shape[0], n_actions=env.action_space.shape[0])
critic = Critic(state_dim=env.observation_space.shape[0])
learner = A2CLearner(actor, critic, gamma=0.99, entropy_beta=0.01)
runner = Runner(env, actor, learner.writer)

memory = runner.run(max_steps=2048)
learner.learn(memory, steps=2048, discount_rewards=True)

Tip

Для систем с сильной инерцией используйте discount_rewards=False, чтобы критик обучался по TD‑таргету с \(V(s')\).

Документация API

A2CLearner(actor, critic, gamma=0.9, entropy_beta=0.01, actor_lr=0.0004, critic_lr=0.004, max_grad_norm=0.5, device=None, log_dir=None, wandb_project=None, wandb_entity=None, wandb_run_name=None, wandb_tags=None, wandb_config=None)

Learner implementing Advantage Actor-Critic (A2C) updates.

Initialize learner with optimizers and hyperparameters.

Parameters:

Name	Type	Description	Default
actor	Module	Policy network.	required
critic	Module	Value network.	required
gamma	float	Discount factor.	0.9
entropy_beta	float	Entropy regularization weight.	0.01
actor_lr	float	Learning rate for actor.	0.0004
critic_lr	float	Learning rate for critic.	0.004
max_grad_norm	float	Gradient clipping norm.	0.5
log_dir	str | None	TensorBoard log directory. If None, the default runs/ directory of SummaryWriter is used.	None

learn(memory, steps, discount_rewards=True)

Update actor/critic using a batch of collected transitions.

Parameters:

Name	Type	Description	Default
memory	list	Collected transitions.	required
steps	int	Global step index used for logging.	required
discount_rewards	bool	If True, uses discounted returns as TD target.	True

get_param_env()

Return serializable configuration needed to reconstruct the learner.

save(path=None, save_gradients=False)

Save A2CLearner model to the specified directory.

Saves actor network, critic network, and configuration. Optionally saves optimizer states for resuming training.

Parameters:

Name	Type	Description	Default
path	str | Path | None	Base save directory. If None, saves to the current working directory.	None
save_gradients	bool	If True, also persist optimizer state dicts.	False

Returns:

Name	Type	Description
Path	Path	The directory where the model was saved.

from_pretrained(repo_name, env=None, access_token=None, version=None, load_gradients=False) classmethod

Load pretrained model from a local directory or Hugging Face Hub.

Parameters:

Name	Type	Description	Default
repo_name	str	Path to a local folder or a Hugging Face repo id (e.g. "namespace/repo_name").	required
env	Optional[Env]	Optional environment instance.	None
access_token	Optional[str]	Hugging Face access token for private repos.	None
version	Optional[str]	Revision / branch / tag on Hugging Face.	None
load_gradients	bool	Restore optimizer states for continued training.	False

Returns:

Name	Type	Description
A2CLearner	A2CLearner	Initialized learner.

publish_to_hub(repo_name, folder_path, access_token=None)

Upload a saved model folder to Hugging Face Hub.

Parameters:

Name	Type	Description	Default
repo_name	str	Repository id on Hugging Face (e.g. "user/my-a2c").	required
folder_path	Union[str, Path]	Local folder produced by :meth:save.	required
access_token	Optional[str]	Hugging Face access token.	None

Runner(env, actor, writer)

Environment interaction loop used to collect training data.

Create runner for data collection.

Parameters:

Name	Type	Description	Default
env	Env	Environment instance.	required
actor	Module	Policy network used to select actions.	required
writer	MetricWriter	TensorBoard writer for logging rewards.	required

reset()

Reset environment and episode state.

run(max_steps, memory=None)

Run the environment for a fixed number of steps and collect transitions.

Parameters:

Name	Type	Description	Default
max_steps	int	Number of environment steps to execute.	required
memory	list	Existing list to append transitions to.	None

Returns:

Name	Type	Description
list	list[tuple[ndarray, float, ndarray, ndarray, bool]]	Collected transitions.