A2C with NARX‑Critic

A2C (Advantage Actor-Critic) uses an actor to select actions and a critic to evaluate states. Our implementation employs a NARX (Nonlinear AutoRegressive with eXogenous inputs) critic, enabling better modeling of dynamics and history by explicitly incorporating past states.

Components

• Actor: Gaussian policy \(\pi_\theta(a|s) = \mathcal{N}(\mu_\theta(s), \sigma_\theta^2)\); implemented in the Actor class (PyTorch)
• Critic (NARX): evaluates \(V(s)\) using an extended input (current state + past signals); see Critic (A2C) and NARX (modular NARX network)
• Experience collection: Runner gathers trajectories and clips actions to the action_space
• Training: A2CLearner.learn updates actor/critic with stabilization (gradient clipping, entropy bonus)

Theory (based on the implementation)

• Discounted returns:

\[ 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 in the code:

\[ A_t = \begin{cases} r_t, & \text{if } discount\_rewards=True \text{ (pure returns)}\\ r_t + \gamma V(s_{t+1}) - V(s_t), & \text{otherwise (TD target)} \end{cases} \]

and then \(A_t = \text{td\_target} - V(s_t)\).

• Losses:

\[ \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 as critic

The NARX network (tensoraerospace/agent/narx/model.py) explicitly feeds previous outputs/states as inputs to predict the next step. In A2C the critic input is the concatenation of the current and previous states (process_memory_narx builds critic_states). This improves \(V(s)\) estimates for systems with significant dynamic memory.

Identification (simplified):

• Training NARX minimizes the MSE between predicted and target outputs over sequences (NARX.train).
• In A2C the critic minimizes the MSE between the target (returns or TD) and the current estimate \(V(s)\).

Training loop

1. Runner.run collects tuples \((s_t, a_t, r_t, s_{t+1}, done_t)\), clipping actions to the action_space.
2. process_memory_narx prepares tensors: actions, rewards (optionally discounted), states, next states, termination flags, and critic_states = [s_t, s_{t-1}].
3. A2CLearner.learn:
4. If discount_rewards=True, critic target td_target = rewards (returns); otherwise r + γ V(s').
5. Advantage: advantage = td_target - V(s).
6. Actor: log-probabilities from Normal(mean, std), entropy; gradient-clipped updates.
7. Critic: MSE; gradient-clipped updates.
8. TensorBoard logging (losses, gradients/parameters, rewards).

Quick start

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

For systems with strong inertia set discount_rewards=False so the critic trains on the TD target with \(V(s')\).

API reference

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.