Environments and wrappers¶
Agents take a gymnax, Brax or MuJoCo Playground environment, or its id, and
wrap it themselves. Every continuous agent acts in [-1, 1], and
agent_action_to_env maps its action onto the environment's bounds; an
expert speaks the environment's units, and the agent reads it through
ExpertInAgentUnits. See Environments and
Differentiable simulation.
ajax.environments.create ¶
register_brax_builder ¶
Register a custom brax env builder for env_id.
Mirrors register_playground_builder for brax-stack envs. The builder
takes (n_envs, episode_length) and must return a fully wrapped env
whose _ajax_env_id attribute equals env_id. Subsequent calls to
_build_brax_env(env_id, ...) delegate to this builder. Used when a
safety-experiments-style env wraps a stock brax robot with custom
termination/observation-augmentation that must survive eval's rebuild.
To support action_repeat > 1 the builder must also accept an
action_repeat keyword (passed only when > 1).
register_playground_builder ¶
Register a custom playground env builder for env_id.
The builder takes (n_envs, episode_length) and must return a fully
wrapped env whose _ajax_env_id attribute equals env_id. Subsequent
calls to _build_playground_env(env_id, ...) delegate to this builder.
To support action_repeat > 1 the builder must also accept an
action_repeat keyword (passed only when > 1).
build_env_from_id ¶
build_env_from_id(
env_id,
n_envs=1,
fresh_reset=True,
*,
episode_length=1000,
action_repeat=1,
differentiable_reset=False,
)
Build a wrapped env from its id (gymnax, mujoco_playground or brax).
fresh_reset only concerns mujoco_playground envs (see
_build_playground_env); gymnax and Ajax's brax stack already draw a
new initial state for every episode.
action_repeat (default 1) repeats each agent action for that many
simulator steps on brax / playground envs (see
:func:_build_playground_env); episode_length (a keyword, default
1000) then counts simulator steps. gymnax envs do not support it.
differentiable_reset concerns the auto-resets that compute a fresh
reset inside step: brax envs, and playground envs with
fresh_reset=True (the default). With differentiable_reset=False
(the default) they evaluate the reset only on steps where some env is
done, inside a lax.while_loop, because the reset can cost as much as
many env steps. Gradients through such an env still flow with respect
to the actions, the policy parameters and the env state, and forward
mode (jax.jvp) works with respect to anything.
Only reverse mode through the reset itself -- jax.grad with
respect to something the reset depends on, such as physics parameters
of the env -- fails, loudly, at trace time: "Reverse-mode
differentiation does not work for lax.while_loop". Pass
differentiable_reset=True for that case: the reset is then evaluated
on every step and kept only for done envs, which reverse mode can cross,
at the cost of a reset per step. Transitions are the same either way,
up to float rounding. Gymnax envs (whose auto-reset is always computed)
and playground's cached auto-reset ignore the flag, and evaluation
(evaluate.setup_environment), which never differentiates, rebuilds
the env with the default.
add_ajax_wrappers ¶
add_ajax_wrappers(
env,
*,
normalize_obs=False,
normalize_reward=False,
gamma=None,
apply_obs_normalization=True,
train=True,
norm_info=None,
)
Ajax's own layers over a task env: the observation / reward
normaliser (train=False with norm_info freezes its statistics),
then, on a continuous action space only, the clip to its bounds. The
one place they are composed.
strip_ajax_wrappers ¶
The task env under Ajax's own layers, and the keywords that rebuild
those layers with :func:add_ajax_wrappers. Everything under them is
the task: the user's wrappers, the flattening of
:func:build_env_from_id, a brax stack's time limit.
ajax.wrappers ¶
Wrappers for environment
BraxWrapper ¶
Base class for the wrappers of a brax or MuJoCo Playground stack.
A copy of brax.envs.base.Wrapper (brax 0.14.2, Apache-2.0) without
its Env base class, so that Ajax imports brax only to build a brax
env (brax comes with the brax extra).
FlattenObservationWrapper ¶
Bases: GymnaxWrapper
Flatten the observations of the environment.
get_obs ¶
Recompute the observation from state, flattened to match reset.
ClipAction ¶
Bases: GymnaxWrapper
Clip a continuous action to the env's action space, or to [low,
high] when given: gym's ClipAction. PureJaxRL's, which this one
descends from, clips to [-1, 1] and leaves reading the space as a
TODO; [-1, 1] would make the top of a Box(0, 2) unreachable.
Ajax's own steps map agent actions onto the box first
(:func:ajax.environments.utils.agent_action_to_env), so here it is a
guard for raw stepping (interaction.step).
ClipActionBrax ¶
InitialStateWrapper ¶
Bases: GymnaxWrapper
Override the plant's state at reset (an initial-condition distribution).
init_state_fn(key, state, params) -> state receives the env's own
reset state and returns the state to start from; the observation is
recomputed with env.get_obs. Use it to narrow (or fix) the initial
conditions of a benchmark, e.g. the p(O) of a control meta-dataset
(Busetto et al. 2024 keep it fixed at the nominal steady state).
Every episode starts this way, the auto-reset ones included: gymnax's
step auto-resets with the inner env's reset_env, so step
applies init_state_fn to that reset state too (fields the inner step
carries across its auto-reset stay as it left them).
NormalizeVecObservationBrax ¶
NormalizeVecObservationBrax(
env,
train=True,
norm_info=None,
normalize_obs=True,
normalize_reward=True,
gamma=None,
apply_normalization=True,
)
Bases: _OnlineNormalization, BraxWrapper
Online normalisation of a brax / playground stack's observations and
rewards. The statistics live in state.info["normalization_info"],
and the info keeps the raw observation and reward (RAW_OBS_KEY,
RAW_REWARD_KEY): the state's own are the normalised ones.
NormalizeVecObservationGymnax ¶
NormalizeVecObservationGymnax(
env,
train=True,
norm_info=None,
normalize_obs=True,
normalize_reward=True,
gamma=None,
apply_normalization=True,
)
Bases: _OnlineNormalization, GymnaxWrapper
Online normalisation of a gymnax env's observations and rewards. The
statistics live in the env state, a subclass of the inner env's state
class with a normalization_info field.
AutoResetWrapper ¶
Bases: BraxWrapper
Automatically resets Brax envs that are done, sampling a new random seed for initialization at each reset. This seed is propagated through info["rng"]
By default the reset is evaluated only on steps where some env is done,
inside a lax.while_loop (see :func:_call_if_any), and skipped on
the others. Rollouts stay differentiable with respect to the actions,
the policy parameters and the env state, none of which the reset depends
on, and forward mode (jax.jvp) works with respect to anything. What
reverse mode cannot do is differentiate through the reset itself, with
respect to something it depends on (e.g. physics parameters of the env):
jax.grad then raises "Reverse-mode differentiation does not work for
lax.while_loop" at trace time. differentiable_reset=True evaluates
the reset on every step instead and keeps the result only for done envs,
which reverse mode can cross, at the cost of a reset per step.
BatchRngWrapper ¶
Bases: BraxWrapper
Splits an unbatched PRNG key into n_envs keys on reset.
Playground's BraxAutoResetWrapper.reset (and FreshAutoResetWrapper.reset)
assumes rng is already shape (n_envs, 2) so it can
jax.vmap(jax.random.split)(rng). Ajax callers pass a single unbatched
key, so this adapter bridges the two conventions.
FinalObsWrapper ¶
Bases: BraxWrapper
Stashes state.obs in state.info['final_obs'] on every reset/step.
Place this directly below an auto-reset wrapper in the stack. The auto-reset
wrapper overwrites state.obs with the reset observation on done, so
without this the pre-reset observation is lost — breaking correct value
bootstrapping on truncation (PPO/SAC need V(s_T) at truncation, not
V(s_reset)). With this wrapper, downstream code can read
state.info['final_obs'] to recover the terminal observation.
Backend-agnostic: duck-typed passthrough that works for both brax State
and mujoco_playground State (both expose .obs and .info).
FreshAutoResetWrapper ¶
Bases: BraxWrapper
Auto-resets a batched mujoco_playground env to a fresh initial state.
Drop-in replacement for playground's BraxAutoResetWrapper in Ajax's
playground stack (EpisodeWrapper, VmapWrapper, FinalObsWrapper, this,
BatchRngWrapper; see ajax.environments.create). With its default
full_reset=False, upstream caches each env's first reset and restarts
every later episode from it, so a whole run sees only n_envs initial
conditions, whereas dm_control draws a new one per episode. Upstream's
full_reset=True does draw new ones, but it (a) evaluates a full reset
on every step (on CheetahRun, whose reset runs a 200-step stabilisation,
~180x the cost of a step) and (b) replaces the whole info of done
envs with the reset's, which zeroes truncation, episode_done and
episode_metrics and overwrites final_obs with the reset
observation, breaking the V(s_T) bootstrap at truncation.
On a step where some env is done, this wrapper draws a fresh batched
reset (computed only on such steps, see :func:_call_if_any) and gives
each done env, from it:
dataandobs;- the
infoentries produced by the base environment's ownreset, i.e. its per-episode state (its rng, task targets, ...).
Everything else describes the transition just taken and passes through:
reward, done, metrics, and the bookkeeping that the wrappers
below write into info on every step (EpisodeWrapper's steps,
truncation and episode_*; FinalObsWrapper's final_obs).
differentiable_reset=True evaluates the reset on every step instead,
for reverse-mode gradients through the reset itself; see
:class:AutoResetWrapper.
TerminatedTruncatedWrapper ¶
Bases: GymnaxWrapper
Split a time-limit-inclusive terminal flag back into terminated/truncated.
Stock gymnax >= 1.0 environments already return the two flags apart, so
they do not need this wrapper. It is for third-party environments that
subclass gymnax's Environment but whose step_env still reports the
pre-1.0 done = terminated | truncated, which gymnax's base step
hands back in the terminated slot. Left unwrapped, such an env makes
every time-limit truncation look like a natural termination, so PPO/SAC
drop the V(s_T) bootstrap and silently underestimate the value of
long-running states.
The time limit is the truncated flag gymnax's base step reports:
it is computed from the pre-reset state, whereas the returned state has
already been auto-reset (time == 0) at the limit.
ObservationNoiseWrapper ¶
Bases: GymnaxWrapper
Sensor noise on a gymnax env's observations, obs + sigma * scale *
N(0, I), drawn afresh at every reset and step from a key folded off
the call's own (the env's own draws are unchanged).
The pre-reset observation an episode end leaves in the info gets the
step's draw, so a transition's next observation is the one the agent
then sees. scale is a scalar or one value per dimension (e.g. the
observations' standard deviations). get_obs stays the state's
noiseless observation, which an expert reads
(interaction.get_raw_obs).
ObservationNoiseWrapperBrax ¶
Bases: BraxWrapper
:class:ObservationNoiseWrapper for a batched brax / playground stack
(above its VmapWrapper, as build_env_from_id builds it). Its
step takes no key: the wrapper's own lives in the state's info, tiled
over the envs (row 0 is read), and is split at every reset and step.
init_norm_info ¶
Initialise running stats with leading axis batch_size.
The leading axis is redundant (online_normalize collapses it on
every update, so rows carry identical values) but it is load-bearing
for env-side callers: env-side stats live inside the env state
pytree, which Brax's VmapWrapper vmaps over the n_envs axis.
Shrinking the stats to (1, *) would break that vmap with a
mismatched-axis error. Agent-side callers should pass
batch_size=1 instead (see init_agent_obs_norm).
ajax.environments.obs_norm ¶
Agent-side running observation normalization.
Lives inside CollectorState.obs_norm_info so stats are part of
agent_state and thread through naturally. Normalization is applied
at every actor/critic apply_fn boundary; stats are updated only at
collection time (online interaction and BC dataset prep). The replay
buffer stores RAW augmented observations so changing stats don't poison
old samples; we re-normalize at sample time.
Why agent-side and not env-wrapper-side: when augment_obs_with_expert_state
is on, the expert_state (PID integrators etc.) is appended AFTER the
env step, in the collector. The env wrapper can't see those dims, so
its running stats don't cover them. Agent-side normalization wraps the
full augmented vector.
init_agent_obs_norm ¶
Initialise running stats for an obs vector of size obs_dim.
Stats are stored with leading shape 1 (not n_envs): online_normalize
reduces the per-env axis on every update so all envs would share the
same scalar stats anyway, and apply_obs_norm only needs the
per-feature mean/var. The size-1 leading axis preserves broadcasting
against batched obs (*, obs_dim) while saving an n_envs× memory
factor on every checkpoint and vmap broadcast.
update_obs_norm ¶
Update running stats with obs and return (normalized_obs, new_info).
obs shape: (n_envs, obs_dim). When info is None the call is
a no-op (returns input unchanged) so call sites can stay unconditional.
The leading raw_columns stay raw (see online_normalize).
apply_obs_norm ¶
Normalize obs using existing stats (no update). Used at every
actor/critic apply_fn site that consumes obs sampled from the
buffer or carried by the eval scan, so stats updates stay localised
to the collection sites. Falls back to raw obs while the running
stats are still empty (count == 0), avoiding a divide-by-≈0 at
step 1 before any update.
The statistics cover the agent's own observation, the leading columns;
what is appended after it stays raw: the expert's action
(interaction.augment_obs) and a critic's action, already in action
units, as the env's raw_columns keep a PID's terms raw.
ajax.environments.system_class ¶
System classes: distributions over environment parameters.
A system class is a distribution p(S) over the dynamics of an
environment, represented as a distribution over its EnvParams. It
is the object that turns one environment into a family of environments
for domain randomisation, meta-learning and in-context control (Busetto
et al. 2024, "One controller to rule them all", arXiv:2411.06482).
Sampling returns batched params: every leaf carries a leading axis of
size n, one entry per parallel environment. :func:ajax.environments
.interaction.reset / :func:~ajax.environments.interaction.step detect
that leading axis and vmap over it (gymnax only: brax envs do not take
per-step params). APG, DreamerV3 and TDMPC2 train on them; the agents of
the shared loop compile them as a static argument and refuse them when
built (ActorCritic.supports_batched_env_params).
SystemClass ¶
Base class: a distribution over EnvParams.
Subclasses implement :meth:sample. The nominal params are kept so
callers can always fall back to the reference system (evaluation on
the nominal plant, curriculum stage 1, ...).
FixedSystem
dataclass
¶
UniformPerturbation
dataclass
¶
Bases: SystemClass
Independent uniform relative perturbation of selected parameter fields.
Each listed field p is drawn as p * (1 + U(-scale, scale)),
independently per field and per sampled system (a vector or matrix field
is scaled by one factor per system); every other field
stays at its nominal value. This is the meta-dataset construction of
Busetto et al. 2024 (there: scale=0.05 on the 19 evaporator
constants).
fields must name attributes of the nominal EnvParams; a typo
raises at construction rather than silently sampling nothing.
env_params_is_batched ¶
True when params carries one system per env along a leading axis.
A batched params (from :meth:SystemClass.sample or
:func:broadcast_env_params) has every leaf with the same leading
axis. An unbatched one always has at least one rank-0 leaf (gymnax's
max_steps_in_episode at minimum), even when other fields are
vectors or matrices (e.g. a plant's state-space model), so "any leaf has
rank >= 1" would misclassify those.
n_systems ¶
The number of systems a batched params carries: its leaves'
leading axis.
broadcast_env_params ¶
Replicate an unbatched params n times along a new leading axis.
Used to run a single nominal system through the batched code path (e.g. a fixed system during the first curriculum stage). Every leaf (scalar, vector or matrix) gets the new leading axis; an already batched params is returned unchanged.
ajax.environments.model_reference ¶
Model-reference tracking tasks on top of any gymnax environment.
Model-reference control (Busetto et al. 2024, arXiv:2411.06482, §2)
asks a controller to make a plant's measured output y_k follow the
output y^d_k of a reference model M driven by a user reference
r_k. This module turns any gymnax env into such a task:
- :class:
StepReferenceis the reference distributionp(R): piecewise-constant random steps of random amplitude and duration. - :class:
LinearReferenceModelis a discrete-time linearM(x_{k+1} = A x_k + B r_k,y^d_k = C x_k + D r_k); the paper's first-order model is :meth:LinearReferenceModel.first_order. - :class:
ModelReferenceWrapperexposes the controller's view of the problem: observation[e_k, u_{k-1}](tracking error and previous input, optionally followed by the raw plant observation) and reward-||y^d_k - y_k||^2, so the undiscounted return is the negative closed-loop matching cost of eq. (9) and the paper's M-RMSE issqrt(-return / horizon).
Every operation is plain jnp so the wrapper is differentiable when
the inner env is (see :mod:ajax.environments.differentiable), and it
composes with :mod:ajax.environments.system_class batched params
because the wrapper is per-env (Ajax vmaps reset/step).
StepReference ¶
Random piecewise-constant references r_{[0, N-1]}.
Each segment holds a value drawn uniformly in [min_value,
max_value] for a duration drawn uniformly in [min_duration,
max_duration] steps (both inclusive), independently per output.
The paper's evaporator case uses values in [20, 25] % and
durations in [20, 50] s at 1 s sampling over N = 100.
LinearReferenceModel ¶
Discrete-time linear reference model M.
x_{k+1} = A x_k + B r_k, y^d_k = C x_k + D r_k with
A: (nx, nx), B: (nx, ny), C: (ny, nx), D: (ny, ny).
first_order
classmethod
¶
Diagonal first-order model, one independent channel per output.
Defaults are the paper's eq. (12): x_{k+1} = 0.4286 x_k + 0.7143
r_k, y^d_k = 0.5669 x_k + 0.2914 r_k (unit DC gain).
init_state ¶
State whose output equals the plant's initial output.
The paper initialises x^M_0 at the plant's initial condition
so the desired trajectory starts where the plant starts. Solves
C x = y0 - D r0 in the least-squares sense.
ModelReferenceWrapper ¶
ModelReferenceWrapper(
env,
reference,
model,
output_idx=(0,),
output_fn=None,
include_raw_obs=False,
action_dim=1,
)
Bases: GymnaxWrapper
Turn a gymnax env into a model-reference tracking task.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
env
|
Any
|
gymnax env (the plant). |
required |
reference
|
StepReference
|
reference distribution |
required |
model
|
LinearReferenceModel
|
reference model |
required |
output_idx
|
Sequence[int]
|
observation dims forming the measured output |
(0,)
|
output_fn
|
Optional[Callable[[Array], Array]]
|
|
None
|
include_raw_obs
|
bool
|
append the plant observation to |
False
|
action_dim
|
int
|
dimension of the input |
1
|
The reference sequence is re-sampled whenever the plant episode ends
(the inner env's auto-reset). That step's reward scores the plant's
final, pre-reset output, and info's final observation is the
controller's view of it. Past the reference horizon the last sample is
held; size reference.horizon to the episode length.
ajax.environments.pid_terms ¶
The error terms a PID controller acts on, carried in the observation.
A PID controller computes u_t = Kp e_t + Ki I_t + Kd D_t from the
tracking error e_t = target_t - current_t, its integral I_t and
its difference D_t. :class:PIDTermsWrapper puts (e_t, I_t, D_t)
in front of a gymnax env's observation, so a policy reading them stays a
function of the observation: every agent (replay or on-policy, with or
without memory) trains on the wrapped env unchanged, and training and
evaluation, which step the same wrapped env, see the same terms. The PID
policy (:class:ajax.modules.pid_actor.PIDPolicy) reads them there.
The terms are the discrete-time PID's with sampling period dt
(backward-difference form, as in Astrom & Hagglund, Advanced PID
Control, 2006, cited from memory, unverified)::
I_t = clip(I_{t-1} + dt * e_t, -integral_limit, integral_limit), I_{-1} = 0
D_t = (e_t - e_{t-1}) / dt, D_0 = 0
The clip is clamping anti-windup (integral_limit=None: no clip).
Every episode starts afresh, the auto-reset ones included: the integral
restarts from dt * e_0 and the difference from 0, so no term carries
over from the previous episode.
The wrapper declares the terms raw (raw_obs_columns): Ajax's
observation normalisers leave them in the env's units, so gains acting on
them are a PID's, and normalise the rest of the observation as usual.
PIDTermsWrapper ¶
Bases: GymnaxWrapper
Observation [e, I, D, raw obs] for a gymnax env.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
env
|
Any
|
gymnax env; its observation is flattened, and the indices below index the flat observation. |
required |
current_idx
|
Sequence[int]
|
observation dims of the controlled values. |
required |
target_idx
|
Sequence[int]
|
observation dims of their set-points, one per value. |
required |
dt
|
float
|
sampling period: the integral sums |
1.0
|
integral_limit
|
Optional[float]
|
clamp of the integral (anti-windup); None: no clamp. |
None
|
split_pid_terms ¶
The (e, I, D) a :class:PIDTermsWrapper puts first in obs,
stacked as (..., 3, n_errors).
ajax.environments.differentiable ¶
Differentiable closed-loop rollouts (backpropagation through the simulator).
The collectors in :mod:ajax.environments.interaction wrap every env step
in stop_gradient: RL agents treat the environment as a black box. This
module is the opposite path, for analytic policy gradient methods that
differentiate an objective through the dynamics themselves (Busetto et al.
2024's in-context controller, APG/SHAC-style trainers, ...):
- :func:
with_transition_gradientsconfigures a (possibly wrapped) gymnax env so its transitions keep gradients. gymnax >= 1.0 detaches them by default and only some envs support the opt-in (Pendulum, MountainCarContinuous, PointRobot, Reacher, Swimmer at the time of writing); unsupported envs raise instead of silently training on zero gradients. - :func:
closed_loop_rolloutruns a fixed-horizon closed loop between a stateful policy and the env with nostop_gradientanywhere, sojax.gradof any function of the returned :class:Rolloutflows through the policy, the environment, and the policy's carry across time (BPTT over the whole horizon).
Gymnax only: brax/playground envs take no per-step params and have their own (MJX) gradient story.
Rollout ¶
Time-major closed-loop trajectory, leaves shaped (T, B, ...).
obs[t] is what the policy saw at step t, action[t] what it
applied, reward[t] / done[t] / next_obs[t] the env's
response. resets[t] is the episode-start flag the policy was given
at step t (True at t=0 and after a done). info is the env's
per-step info dict with every leaf stacked over time (e.g. a tracking
wrapper's desired output), for metrics that need more than the reward.
policy_action[t] is the policy's own action, before an
action_map mapped it onto the env (None without one).
with_transition_gradients ¶
Return a copy of env whose transitions preserve gradients.
Walks Ajax's GymnaxWrapper chain (wrappers keep the inner env on
_env) so a wrapped env is re-wrapped around the gradient-enabled
core. Raises ValueError when the core env does not implement
differentiable transitions.
closed_loop_rollout ¶
closed_loop_rollout(
policy_step,
policy_carry,
rng,
env,
env_params,
horizon,
n_envs=None,
action_map=None,
)
Run horizon closed-loop steps with gradients flowing end to end.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
policy_step
|
PolicyStepFn
|
|
required |
policy_carry
|
Any
|
initial policy carry (any pytree, |
required |
rng
|
Array
|
PRNG key; split into per-env reset and step keys. |
required |
env
|
Any
|
gymnax env, already configured with
:func: |
required |
env_params
|
Any
|
unbatched |
required |
n_envs
|
Optional[int]
|
number of parallel envs; required when |
None
|
action_map
|
Optional[Callable[[Array], Array]]
|
maps the policy's action onto what the env receives
(differentiable), e.g. agent units onto the env's bounds; the
rollout then keeps both ( |
None
|
Returns:
| Type | Description |
|---|---|
Tuple[Rollout, Any]
|
|
Action units¶
ajax.environments.utils.agent_action_to_env ¶
Map an agent action in [-1, 1] to the env's action bounds.
Every Ajax env step carrying an agent's or an expert's action calls it:
every continuous agent acts in [-1, 1] (squashed or not), whatever
the env's box. action is [n_envs, *action_shape]. Continuous
actions are clipped to [-1, 1] and, on every dimension of the Box
with finite bounds [low, high], mapped affinely onto them: c + h
* clip(a, -1, 1), with centre c = (high + low) / 2 and half-range
h = (high - low) / 2 (CleanRL's action_bias and
action_scale).
This is DreamerV3's ClipAction + NormalizeAction (a dimension
with an infinite bound keeps the clipped action) and TD-MPC2's action
scaling. For [-1, 1] bounds -- every brax / playground env, gymnax
envs like MountainCarContinuous -- the map is exactly the clip, even
when the bounds are traced; gymnax Pendulum's torque is [-2, 2].
Per-env (batched) gymnax params (ajax.environments.system_class)
give per-env bounds: env e's action is mapped with its own system's
bounds. Discrete actions pass through. Replay stores the agent's raw
action, not this one. :func:env_action_to_agent is the inverse.
ajax.environments.utils.env_action_to_agent ¶
Map an action in the env's bounds to agent units, [-1, 1].
The inverse of :func:agent_action_to_env on the env's box, (a - c)
/ h, with the same action layout and per-env bounds. It does not
clip: an action outside the box stays outside [-1, 1]. It is the
identity on [-1, 1] boxes, on a dimension with an infinite bound
and on discrete actions.
ajax.environments.utils.ExpertInAgentUnits ¶
An expert written in env units, read in agent units ([-1, 1]).
What the expert returns -- expert(obs), a stateful expert(state,
obs)'s action, step_with_gains's action -- goes through
:func:env_action_to_agent's map, unclipped. init_state,
learnable_fields, anchor_gains and step_with_gains exist
exactly when the expert's do (hasattr still tells a stateful
expert). It hashes by identity, for jit's static arguments.
Its bounds are one (A,) box read from concrete params, also inside
jit (static params): a stateless expert is also vmapped over single
observations, where per-env bounds would broadcast its action. Per-env
boxes that differ raise.
wrap
classmethod
¶
expert in agent units; None and a wrapped expert stay as they are.
Running a trained policy¶
ajax.environments.interaction.get_pi ¶
Get the policy distribution for the given observation and actor state.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
done
|
Optional[Array]
|
Done flags for recurrent mode. |
None
|
Returns:
| Type | Description |
|---|---|
Tuple[Distribution, LoadedTrainState]
|
The policy distribution and the updated actor state. |