TD-MPC2

TD-MPC2 learns a model of the environment in a latent space and plans with it at every step. Its world model has no decoder: an encoder, latent dynamics, a reward predictor and an ensemble of Q-functions, trained together on short sub-trajectories from replay. To act, a sampling planner (MPPI) tries action sequences in the latent model, guided by a learned policy prior, and executes the first action of the best ones. TDMPC2MultiTask trains one such model offline on the data of several tasks.

Paper: Hansen, Su and Wang, TD-MPC2: Scalable, Robust World Models for Continuous Control, ICLR 2024 (arXiv:2310.16828). Official code: nicklashansen/tdmpc2.

Ajax reproduces the paper-era code (b67b21c, whose agent is identical to 5f6fade). The world-model design notes give the full specification, and deviations.md lists every version choice and every departure from the reference.

Can I use?

Feature	TDMPC2	TDMPC2MultiTask
Discrete actions	no	no
Continuous actions	yes	yes
Memory (memory=)	no	no
Environments	gymnax, Brax, MuJoCo Playground	trains on a dataset; evaluates on any of them
Observations	flat vectors, raw (no normalisation)	flat vectors of different sizes per task
Terminating episodes	refused, unless episodic=True	no
Several seeds in one call	yes	yes
Extension phases	pretrain, post_update, eval_metrics	the same
Resume and checkpoints	yes; a checkpoint restored by a new agent equals the uninterrupted run (test)	yes, the same

Example

# doc: skip (about two minutes on a CPU; tests/agents/TDMPC2/test_TDMPC2.py covers this path)
from ajax import TDMPC2
from ajax.logging.wandb_logging import LoggingConfig

# model_size=1 and a smaller planner than the paper's, for a quick CPU run.
agent = TDMPC2("Pendulum-v1", model_size=1, num_samples=64, num_elites=8,
               num_pi_trajs=4, iterations=2, seed_steps=400)
logging = LoggingConfig(config={}, log_frequency=600, use_wandb=False)
state, evaluations = agent.train(seed=0, n_timesteps=1_200, num_episode_test=2,
                                 logging_config=logging)
print(agent.gamma)  # 0.975, derived from Pendulum's 200-step episodes
print(evaluations["Eval/episodic mean reward"])

The paper's DMC protocol: action repeat 2 over 1,000 simulator steps (500 agent steps per episode), the 5M model, one environment. The paper's "1M environment steps" are 500,000 agent steps:

# doc: skip (the paper's full budget; benchmarks/world_models/paper_protocol.py runs it)
from ajax import TDMPC2

agent = TDMPC2("CheetahRun", model_size=5, action_repeat=2, episode_length=1000)
state, _ = agent.train(seed=0, n_timesteps=500_000)

Key hyperparameters

Every default is the paper's (config.yaml at 5f6fade).

Argument	Default	What it does
model_size	5	preset of 1, 5, 19, 48 or 317 M parameters; fills enc_dim, mlp_dim, latent_dim, num_enc_layers, num_q
n_envs	1	environments stepped in lockstep, one update per environment step
horizon	3	length of the training sub-trajectories and of the plans
batch_size	256	sub-trajectories per update
buffer_size	1,000,000	environment steps kept in replay
seed_steps	None	random steps before training; None is max(1000, 5 T)
updates_per_step	1.0	updates per environment step after the seed phase
learning_rate	3e-4	a float or a schedule; the encoder trains at learning_rate * enc_lr_scale (0.3)
gamma	None	None derives it from the episode length T: (T/5 - 1) / (T/5), clipped to [0.95, 0.995]
iterations, num_samples, num_elites	6, 512, 64	MPPI iterations, sampled plans and kept plans
num_pi_trajs	24	plans proposed by the policy prior
consistency_coef, reward_coef, value_coef	20, 0.1, 0.1	world-model loss weights
rho	0.5	per-step decay of the loss weights along the horizon
entropy_coef	1e-4	policy prior's entropy weight
tau	0.01	target Q update rate
episodic	False	accept terminating tasks (the latest code's episodic mode)
action_repeat, episode_length	1, 1000	simulator steps per agent step, and per episode (Brax, Playground)

Notes

• What n_timesteps counts. Agent steps (environment steps after action repeat); logs also carry env_frames. With episodic=True it counts rows: environment steps plus one per finished episode.
• Terminations. The paper-era agent bootstraps through every episode end, so a terminating task would train on wrong targets. Ajax counts terminations and raises a ValueError after train if there were any (deviation T10). episodic=True grafts the latest official code's episodic mode (commit eece801) onto the paper-era agent instead: a termination head and variable-length episodes in replay (T29).
• Evaluation (with a logging_config): every log_frequency environment steps, num_episode_test episodes of the planner in evaluation mode, on a rebuilt environment, each evaluation from fresh initial states (the same ones when a seed is rerun; T9). The paper uses log_frequency=50_000 with 10 episodes. LoggingConfig's default of 1,000 makes a 500-step-episode run spend about twice as long evaluating as training.
• Replay memory is sized at each train call for the steps the run will have taken: min(buffer_size, total) (T26), in the episodic mode as well (T29). At full capacity it takes 124 MB per seed on DMC Walker. A run trained in chunks, or restored from a checkpoint by a new agent, replays exactly what the uninterrupted run does.
• Actions. The planner's actions lie in [-1, 1]; Ajax maps them onto the environment's bounds, as the reference scales its actions.
• Several environments. n_envs > 1 plans for each environment separately and takes n_envs updates per step (T7); the reference uses one.
• The DMC tasks run on MuJoCo Playground's MJX ports, whose physics differ from dm_control's (E18), so comparisons with the published curves are approximate.

Multi-task: TDMPC2MultiTask

The paper's multi-task models are trained offline, on the pooled replay of single-task agents. TDMPC2MultiTask does the same: a learnable task embedding conditions the model, observations are zero-padded to the largest task's size, the unused action dimensions are masked, and each task keeps the discount of its own episode length. Training does no planning, and n_timesteps counts updates. Build the dataset from single-task runs with export_episodes (one run's replay), concatenate_episodes (several seeds of a task) and pool_tasks; save_dataset and load_dataset store it as an .npz file.

# doc: skip (about five minutes on a CPU; tests/agents/TDMPC2/test_TDMPC2MultiTask.py covers this path)
from ajax import TDMPC2, TDMPC2MultiTask
from ajax.agents.TDMPC2.dataset import export_episodes, pool_tasks
from ajax.logging.wandb_logging import LoggingConfig

small = dict(model_size=1, num_samples=64, num_elites=8, num_pi_trajs=4, iterations=2)
tasks = ["Reacher-misc", "PointRobot-misc"]  # 8 and 6 observations, 2 actions, T = 100

episodes = []
for task in tasks:  # one single-task run per task; its replay is the task's data
    agent = TDMPC2(task, seed_steps=500, **small)
    state, _ = agent.train(seed=0, n_timesteps=1_000)
    episodes.append(export_episodes(agent, state, seed_index=0, name=task))
dataset = pool_tasks(episodes)  # every task's episodes must have the same length

multi = TDMPC2MultiTask(dataset, eval_envs=tasks, batch_size=256, **small)
logging = LoggingConfig(config={}, log_frequency=500, use_wandb=False)
state, history = multi.train(seed=0, n_timesteps=1_000, num_episode_test=2,
                             logging_config=logging)
print(history["Eval/Reacher-misc/episodic mean reward"])
print(history["Eval/episodic mean reward"])  # the mean over tasks

It also logs Eval/normalized score, the mean over tasks of return / 10 (the paper's normalisation of DMC returns). multi.evaluate(state, num_episodes=10) runs the same evaluation outside training, on a state train returned, and gives the same keys, one value per seed; without eval_envs there is nothing to evaluate on, and it raises. Defaults are the single-task agent's, except batch_size=1024 and task_dim=96, the paper's (T14). The evaluation follows the offline trainer's protocol (T16). Ajax's datasets are generated by its own single-task runs, not the paper's mt30 and mt80 data, so they validate the mechanisms rather than the paper's numbers (T17).

Results

Recorded on CPU in PERFORMANCE_REPORT.md, with model_size=1 and every other setting at the paper's default, two seeds: Pendulum reaches -143 and -124 after 6,000 steps (a random policy scores about -1,200), and Playground's CartpoleBalance (action repeat 2) reaches 810 and 854 after 10,000 steps, still improving. The GPU validation against the published curves is specified in VALIDATION.md and has not been run yet. The probes (tests/probing) train TD-MPC2 briefly on tiny environments whose answers need its model (a payoff past its planning horizon; twin tasks for the multi-task agent), and its own tests (tests/agents/TDMPC2) check that a run split by a checkpoint equals the run left whole, for both agents. On the CPU speed benchmark (benchmarks/agent_bench.py: Pendulum, 4 environments, 8,000 steps, a smaller model than the paper's) TD-MPC2 ran at 292 steps per second of execution (245 with compilation) at main a8a13c2.

API

ajax.agents.TDMPC2.TDMPC2.TDMPC2

TDMPC2(
    env_id,
    n_envs=1,
    model_size=5,
    enc_dim=None,
    mlp_dim=None,
    latent_dim=None,
    num_enc_layers=None,
    num_q=None,
    simnorm_dim=8,
    num_bins=101,
    vmax=10.0,
    dropout=0.01,
    horizon=3,
    rho=0.5,
    consistency_coef=20.0,
    reward_coef=0.1,
    value_coef=0.1,
    learning_rate=0.0003,
    enc_lr_scale=0.3,
    grad_clip_norm=20.0,
    pi_eps=1e-05,
    weight_decay=0.0,
    tau=0.01,
    entropy_coef=0.0001,
    log_std_min=-10.0,
    log_std_max=2.0,
    batch_size=256,
    buffer_size=1000000,
    seed_steps=None,
    updates_per_step=1.0,
    gamma=None,
    discount_denom=5,
    discount_min=0.95,
    discount_max=0.995,
    iterations=6,
    num_samples=512,
    num_elites=64,
    num_pi_trajs=24,
    min_std=0.05,
    max_std=2.0,
    temperature=0.5,
    episodic=False,
    termination_coef=1.0,
    episode_length=1000,
    action_repeat=1,
    env_params=None,
    extensions=(),
)

Bases: ActorCritic

TD-MPC2 for continuous control, single task, online.

Paper DMC protocol: action_repeat=2, episode_length=1000 (simulator steps, so T = 500 agent steps, gamma = 0.99, seed_steps = 2500), model_size=5, n_envs=1; n_timesteps counts agent steps, so the paper's "1M env steps" is n_timesteps=500_000 (logs carry env_frames, tdmpc2_spec 4.22).

Static hyperparameters (they fix shapes or trace-time structure): the architecture (model_size preset of {1, 5, 19, 48, 317} M parameters filling the None widths enc_dim, mlp_dim, latent_dim, num_enc_layers, num_q; an explicit width overrides it; tdmpc2_spec 1.20), simnorm_dim, num_bins, vmax (two-hot bins over [-vmax, vmax] in symlog space; the reference's vmin = -vmax, T24), horizon, batch_size, buffer_size, seed_steps and the planner sizes iterations (+2 when the action dimension is >= 20), num_samples, num_elites, num_pi_trajs. Every other hyperparameter is a static Python float as well, except learning_rate, which accepts a float or a schedule Callable[[int], float] (the world model's encoder runs at learning_rate * enc_lr_scale, 5f6fade:tdmpc2/tdmpc2.py:21-28).

Derived from the episode length T in agent steps (:func:ajax.environments.utils.agent_episode_length: brax / playground episode_length // action_repeat, gymnax the env's max_steps_in_episode): gamma=None is clip((T/denom - 1) / (T/denom), discount_min, discount_max) (tdmpc2_spec 2.20, tdmpc2.py:36-49) and seed_steps=None is max(1000, 5 T) (4.3, envs/__init__.py:81).

Replay (:mod:ajax.agents.TDMPC2.buffer, the b67b21c episode buffer): buffer_size env steps, stored in a ring sized at every :meth:train call for the env steps the run has taken when the call ends, min(buffer_size, total) (the reference's clamp to the run length, which saves memory without changing what is replayed; deviation T26). A resumed run -- the same agent continuing, or a new one restoring a checkpoint into its n_timesteps=0 skeleton (ajax.checkpoint) -- moves the carried episodes into the ring of its own call (:meth:~ajax.agents.TDMPC2.buffer.EpisodeBuffer.adopt), so a run trained in chunks replays exactly what the uninterrupted run does. The last call's capacity and bytes are :attr:replay_capacity and :attr:replay_bytes_per_seed (also in :attr:config, hence the run config of the loggers). Memory per seed: R * n_envs * (T + 1) * (obs_dim + A + 1) * 4 bytes with R = ceil(capacity / (T n_envs)) + 1: 124 MB for DMC walker (obs 24, A 6, T 500) at the full 1,000,000-step capacity; the seed vmap multiplies it.

Evaluation (with a logging_config): every log_frequency env steps, num_episode_test episodes of the planner in eval_mode on a rebuilt env with the training action repeat and episode length, each evaluation from fresh initial states (online_trainer.py:27-48; deviation T9). One evaluation makes num_episode_test * T planner decisions (batched over the episodes; on CPU each costs roughly half a training env step, which plans and updates). The paper protocol is log_frequency = 50_000 with num_episode_test = 10 (tdmpc2_spec 4.21); :class:~ajax.logging.wandb_logging.LoggingConfig's default log_frequency of 1000 makes a T = 500 run spend about twice as long evaluating as training.

Continuous actions only. The env must be fixed-length and non-terminating (all DMC tasks): a termination, off the T-step schedule or on its last step, raises ValueError after :meth:train (deviation T10). Observations and rewards are the env's raw values (no normalisation, as in the reference).

episodic=True takes terminating tasks instead (the latest code's episodic mode, nicklashansen/tdmpc2@eece801; deviation T29): a termination head, 1 - terminated in the targets and the planner, the variable-length replay of ceil(min(buffer_size, total) / n_envs) rows per env (total the rows of the run when the call ends, sized and adopted on a resume as the episode ring is), and n_timesteps counting rows (env steps plus one per finished episode). T is then the time limit, still bootstrapped. A run that never updated because every episode terminated within horizon steps raises ValueError after :meth:train.

Extensions: pretrain, post_update (folded after every update, step = env steps) and eval_metrics (with init_state); the other phases are rejected at construction.

Parameters:

Name	Type	Description	Default
env_id	str | EnvType	env id (gymnax, brax or mujoco_playground) or a prebuilt gymnax env.	required
n_envs	int	lockstep envs; each stepping tick runs n_envs updates (UTD 1 per env step; deviation T7 for > 1).	1
updates_per_step	float	updates per env step after the seed-phase burst, max(1, round(n_envs * updates_per_step)) per stepping tick (static); the default 1 is the reference's UTD, which config.yaml does not expose (deviation T28). The burst stays seed_steps updates. Episodic mode: max(1, round(stepped * updates_per_step)) on a tick that stepped stepped > 0 envs.	1.0
episode_length	int	simulator steps per episode on brax / playground (ignored on gymnax, which uses its params' max_steps_in_episode).	1000
action_repeat	int	simulator steps per agent step (brax / playground; the paper's DMC protocol uses 2).	1
env_params	Optional[EnvParams]	gymnax env params (one system).	None
weight_decay	float	decoupled weight decay of the world model's and the policy's Adam (OptimizerConfig.weight_decay); Ajax's, not a config.yaml hyperparameter: 0 keeps the reference's plain Adam.	0.0
extensions	Sequence[Extension]	composable research features (see :mod:ajax.extensions).	()

Every other argument is the hyperparameter of the same name in nicklashansen/tdmpc2@5f6fade:tdmpc2/config.yaml (lr is learning_rate), with its paper default; see the class docstring for the static ones and the derived defaults.

resume_iteration_offset

resume_iteration_offset(initial_state)

The absolute tick a resumed run starts at, equal across seeds (:func:~ajax.environments.row_collector.resume_tick).

train

train(
    seed=42,
    n_timesteps=int(1000000.0),
    num_episode_test=10,
    logging_config=None,
    on_ids_ready=None,
    initial_state=None,
    total_timesteps=None,
    **kwargs,
)

:meth:ActorCritic.train; n_timesteps counts env (agent) steps.

Sizes the replay ring for the env steps the run will have taken when this call ends and, on a resume, moves the carried episodes into it (class docstring). Returns (state, evaluations), as every agent on the shared loop (:mod:ajax.agents.loop): every logged key's values at the evaluations with a logging_config, else None (nothing is evaluated or logged). Raises ValueError after the run if any env terminated (deviation T10): the paper-era TD-MPC2 bootstraps through every episode end, so a terminating task would train on wrong targets. The episodic mode takes terminations.

ajax.agents.TDMPC2.TDMPC2MultiTask.TDMPC2MultiTask

TDMPC2MultiTask(
    dataset,
    eval_envs=None,
    model_size=5,
    enc_dim=None,
    mlp_dim=None,
    latent_dim=None,
    num_enc_layers=None,
    num_q=None,
    task_dim=PAPER_TASK_DIM,
    simnorm_dim=8,
    num_bins=101,
    vmax=10.0,
    dropout=0.01,
    horizon=3,
    rho=0.5,
    consistency_coef=20.0,
    reward_coef=0.1,
    value_coef=0.1,
    learning_rate=0.0003,
    enc_lr_scale=0.3,
    grad_clip_norm=20.0,
    pi_eps=1e-05,
    weight_decay=0.0,
    tau=0.01,
    entropy_coef=0.0001,
    log_std_min=-10.0,
    log_std_max=2.0,
    batch_size=1024,
    discount_denom=5,
    discount_min=0.95,
    discount_max=0.995,
    iterations=6,
    num_samples=512,
    num_elites=64,
    num_pi_trajs=24,
    min_std=0.05,
    max_std=2.0,
    temperature=0.5,
    episode_length=1000,
    action_repeat=1,
    extensions=(),
)

Multi-task TD-MPC2, trained offline on a pooled dataset (paper era).

The reference's multi-task path is offline only (train.py:52; tdmpc2_spec 3.25, 4.1): the updates of the single-task agent, conditioned on each slice's task (:func:ajax.agents.TDMPC2.multitask.update: a learnable task embedding with max-norm 1, zero-padded observations, prefix action masks, per-task discounts), on batches of slices drawn uniformly over the dataset's pooled episodes, no planning during training. n_timesteps counts updates (the reference's iterations; paper: 10M at batch 1024, spec 4.B). The paper protocol is batch_size=1024 (the default here) and task_dim=96 for every model size (deviation T14).

The dataset (:class:~ajax.agents.TDMPC2.dataset.MultiTaskDataset, e.g. :func:~ajax.agents.TDMPC2.dataset.pool_tasks of single-task runs' exports) fixes the tasks: task i has the dataset's obs_dims[i], action_dims[i] and episode length T_i, whose discount is the reference's heuristic of T_i (discount_denom, discount_min, discount_max; spec 2.20). The model sees the padded observation and action spaces; the planner's iteration rule uses the padded action dim (spec 3.1). The dataset is one array set on the device, shared by every seed: the constructor places it there once (whatever built it: host NumPy arrays are transferred then, not at every chunk) and :meth:train passes it to its jitted, seed-vmapped program as an unbatched argument, so the seeds read the same copy (dataset.nbytes) and the program does not embed it.

Logging (with a logging_config): every log_frequency updates, the last update's quantities (Train/..., timestep = updates done), the extensions' eval_metrics and, with eval_envs, an evaluation: a host loop over the tasks runs num_episode_test episodes of each in its own env (eval_envs[i], built with episode_length simulator steps and action_repeat; its dims and agent-step episode length must be task i's), the planner in eval_mode with t0 on each episode's first step: the offline trainer's protocol (offline_trainer.py:22-39; deviation T16: evaluate.py's exploration noise is not reproduced). Logged per task Eval/<name>/episodic mean reward and length, their mean over tasks Eval/episodic mean reward and the normalised score Eval/normalized score, the mean over tasks of return / 10 (the paper's DMC normalisation, meaningful for DMC returns in [0, 1000]; spec 4.20). The first evaluation follows the first log_frequency updates (the reference's follows its first update). Each evaluation starts its episodes from fresh initial states. The embedding renorm of a look-up at evaluation is not persisted (deviation T15).

Static hyperparameters, as :class:~ajax.agents.TDMPC2.TDMPC2.TDMPC2's: the architecture (model_size preset and explicit widths), task_dim, simnorm_dim, num_bins, vmax, horizon, batch_size and the planner sizes. learning_rate accepts a float or a schedule (the task embedding is trained by the world-model optimizer at learning_rate, the encoder at learning_rate * enc_lr_scale).

Extensions: pretrain, post_update (folded after every update, step = updates done) and eval_metrics (with init_state), the phases of the single-task agent; the others are rejected at construction.

Parameters:

Name	Type	Description	Default
dataset	MultiTaskDataset	the pooled multi-task dataset (it fixes the tasks).	required
eval_envs	Optional[Sequence[str | EnvType]]	one env id or prebuilt gymnax env per task, in task-id order, for evaluation; None trains without evaluating.	None
task_dim	int	task-embedding width (paper: 96; deviation T14).	PAPER_TASK_DIM
episode_length	int	simulator steps per episode of the brax / playground evaluation envs.	1000
action_repeat	int	simulator steps per agent step of the evaluation envs (the paper's DMC protocol: 2).	1
extensions	Sequence[Extension]	composable research features (see :mod:ajax.extensions).	()

Every other argument is the hyperparameter of the same name of :class:~ajax.agents.TDMPC2.TDMPC2.TDMPC2 (5f6fade's config.yaml), with the paper's multi-task batch size 1024.

resume_update_offset

resume_update_offset(initial_state)

The updates a resumed run has done, equal across seeds.

evaluate

evaluate(state, num_episodes=10)

One evaluation of every task (host loop), per seed.

state has a leading seed axis (as :meth:train returns it). Returns {key: [n_seeds]}: per task Eval/<name>/episodic mean reward and Eval/<name>/mean episodic length, their means over tasks and Eval/normalized score (class docstring). It does not advance state.n_logs: :meth:train does after logging, so the next evaluation draws fresh keys (:func:ajax.log.evaluation_keys).

train

train(
    seed=42,
    n_timesteps=int(1000000.0),
    num_episode_test=10,
    logging_config=None,
    on_ids_ready=None,
    initial_state=None,
    total_timesteps=None,
)

Run n_timesteps more updates for every seed (vmapped).

A fresh run initialises the networks (and folds the extensions' init_state and pretrain); initial_state (a state this agent's train returned, or a checkpoint restored into the n_timesteps=0 skeleton) continues a run; its buffers are donated to the first chunk (as ActorCritic.train donates them). The updates run as jitted scans; with a logging_config whose log_frequency (updates) is set, the scans stop every log_frequency updates of the run (an absolute cadence: a resumed run keeps it) to log the training metrics and, with eval_envs, to evaluate every task (class docstring). total_timesteps is the whole run's updates, the extensions' total_steps (None: this call's end), as :meth:ActorCritic.train's.

Returns (state, history): the state with a leading seed axis and, when it logged, {key: [n_seeds, n_logs]} (training and evaluation metrics), else None.

ajax.agents.TDMPC2.dataset.export_episodes

export_episodes(
    agent,
    state,
    seed_index=None,
    *,
    name="task",
    allow_evicted=False,
)

The committed episodes of a single-task TD-MPC2 run, oldest first.

Reads state.buffer_state (the episode ring of agent's run, :mod:ajax.agents.TDMPC2.buffer) after the ticks its collector has run: every complete episode the ring holds, in the order collected (rounds in order, env order within a round; the episode in progress is not exported). Its rows are the ring's obs-aligned rows, unchanged.

Parameters:

Name	Type	Description	Default
agent	TDMPC2	the :class:~ajax.agents.TDMPC2.TDMPC2.TDMPC2 that trained state (its n_envs and episode length).	required
state	Any	the run's state, as agent.train returned it; with a leading seed axis, seed_index picks the seed.	required
seed_index	Optional[int]	the seed to export, required for a seed-batched state.	None
allow_evicted	bool	export what the ring still holds when it wrapped and evicted older episodes. By default this raises: a dataset of the full training history (DESIGN.md §7) needs buffer_size at least the run length.	False

Raises:

Type	Description
ValueError	on a missing or extraneous seed_index, when no episode is complete, or when episodes were evicted (unless allow_evicted).

ajax.agents.TDMPC2.dataset.concatenate_episodes

concatenate_episodes(episodes)

The episodes of several runs of one task (e.g. seeds), in order.

They must agree on the rows, dims, episode length and name.

ajax.agents.TDMPC2.dataset.pool_tasks

pool_tasks(tasks)

The multi-task dataset of tasks, task i being tasks[i].

Observations and actions are zero-padded at the end to the largest dims (MultitaskWrapper, spec 1.23, 4.16); episodes keep their order, task after task. Every task's episodes must have the same number of rows (cutting longer episodes into segments, as the mt80 data probably did, spec 4.18, is not done here). Names must be unique. The arrays are placed on the default device once.

ajax.agents.TDMPC2.dataset.save_dataset

save_dataset(dataset, path)

Write dataset to path (an uncompressed NumPy .npz).

The file holds the four arrays under their field names and meta, a JSON string with the format name and version and the per-task metadata (obs dims, action dims, episode lengths, names): plain arrays and text, readable without pickle (:func:load_dataset). The dataset is checked first (:meth:MultiTaskDataset.check) and the file is written to a temporary name and renamed, so an interrupted save leaves no partial file under path.

ajax.agents.TDMPC2.dataset.load_dataset

load_dataset(path)

The dataset :func:save_dataset wrote to path, checked.

The arrays stay host NumPy arrays (TDMPC2MultiTask places the dataset on the device once). Raises ValueError on another format or version, or when the data do not follow the schema.