Skip to content

Latest commit

 

History

History
101 lines (64 loc) · 13.8 KB

File metadata and controls

101 lines (64 loc) · 13.8 KB

Implementation and reproduction notes

Scope and source selection

This release was organized from the authors' DMControl research source. The original training factory routed --algorithm soda to algorithms/soda_conv.py, a T2RD implementation retained under an experimental name. Its active backpropagation used three objectives:

Source operation Release operation Weight
compute_kernel_loss compute_consistency_loss 1
compute_full_reconstruction_loss_1 (original-latent direction) compute_reconstruction_loss 100
compute_pred_InverseDynamics_loss (next-content term) compute_dynamics_loss 0.1

The source's soda_conv.py had the 50/50 selector commented out and used only random convolution. The neighboring soda.py had mixed augmentation enabled but consistency backpropagation disabled. File names alone did not identify the final method. Historical log folders may contain earlier versions or ablations; the authors confirmed that those runs are references for configuration, while the paper's three loss objectives define the full method. This release follows the authors' clarification and the paper: mixed augmentation plus all three objectives. It does not select one of those incomplete experimental variants as the full method.

A second, local archive contains both the final latent-reconstruction method and structural ablations. Its soda.py selects convolution plus pixel reconstruction, and soda_T2RD.py also uses pixel reconstruction despite its name. The paper identifies pixel reconstruction as a variant. The full method is present in soda_RDSR.py and soda copy.py: their executable syntax trees are identical, with 50/50 augmentation and all three latent-representation losses active. The release was directly compared with that unmodified full-method code, without enabling or replacing any reference loss or augmentation branch.

For a comparison against the untouched active source, --augmentation conv preserves its convolution-only path. The default is the intended mixed augmentation. No source files or checkpoints were modified while preparing this release.

Numerical conventions retained

  • Observations are rendered as 9×100×100; replay sampling crops to 9×84×84. Action selection uses a center crop. This follows Section 5.1 and the source. Table 9's 84×84 rendering / random-shift description is not substituted for the actual crop pipeline.
  • The 11-layer CNN produces a flattened latent of dimension 14,112. A 200-dimensional intermediate representation feeds two 100-dimensional branches. SAC uses only the content branch, with separate actor/critic projections.
  • Consistency and latent reconstruction compare L2-normalized vectors with mean MSE. The dynamics MSE is unnormalized. We do not reinterpret the equations as summed Euclidean losses.
  • The source's reconstruction MLP retains its final ReLU; the dynamics network has three linear layers with two intervening ReLUs.
  • SAC, critic-target update, and auxiliary optimization keep the original order. Auxiliary learning uses the same transition batch as SAC. The source's separate auxiliary batch-size option was unused and has been removed.
  • The three auxiliary backward passes and their Adam steps retain their original order. EMA updates parameters only; BatchNorm buffers continue to update through forwards.
  • The policy/content encoder is shared by actor, critic, and auxiliary learning. Detachment during actor updates is retained.
  • The source recursively reinitialized online and target auxiliary networks after copying them. The release keeps independently initialized branches, not a silently changed identical initialization. Removing unrelated network constructors changes the random stream before those initializations.
  • The source's style-side auxiliary predictor is retained as part of the dual-output module and state mapping. In the active three-loss graph, some style-side parameters receive no direct gradient: the original-style reconstruction target is stop-gradient and reverse reconstruction is disabled. No new style loss is introduced in this release.
  • The training loop keeps the 1,000 initial random actions, the initial burst of 1,000 optimization updates, time-limit bootstrapping, and inclusive final step. 500k names the final step index and entails 500,001 environment decisions in this inherited loop. Action repeat multiplies physics steps. This convention should be considered when comparing frame budgets.

Removing unused losses without changing state updates

The source computed inverse-dynamics, reward-prediction, and reverse-reconstruction losses without backpropagating them. It also contained unused classifier, attention, mutual-information, distance, pixel-reconstruction, and alternative baseline code.

Those losses, heads, optimizers, classifier replay storage, visualization scripts, and duplicate algorithm files were removed. However, some discarded loss branches first called the shared encoder and changed BatchNorm running statistics. The release preserves those forwards explicitly under torch.no_grad() in compute_dynamics_loss, in the same order. Deleting these calls would alter later policy outputs even though the removed loss had no backward call.

The auxiliary predictor and target retain the source's train-mode behavior; policy action selection switches the shared actor/critic modules to evaluation mode. The release does not silently change this BatchNorm convention.

Configuration differences made explicit

The authors confirmed that experiments were launched with python3 src/train.py --algorithm soda --seed 1, and that soda was the historical name of their modified T2RD agent. The release documents --algorithm t2rd, retaining soda as an alias. Both select the same T2RD implementation. The actual command's defaults, including auxiliary Adam at 1e-3, are authoritative. Table 9 and an older shell script list 3e-4; these are not used to override the author-confirmed command.

The server train.py snapshot contained a commented periodic-evaluation block. The subsequently supplied local archive has an active block guarded by step % eval_freq == 0 and step > 0, and it creates three separate test environments. The release evaluates during training at the first interval and subsequent episode boundaries. Every round evaluates the training domain and then each persistent color-hard, video-easy, and video-hard environment, with 30 episodes each. These three settings are all evaluated in the same round, sequentially, as in the original loop.

Older logs contain both 10k and 20k intervals, including variations within a task, and both schedules with and without a step-0 baseline. The release defaults to 20k (10k for Ball-in-Cup), with explicit --eval_freq and --eval_at_start options to reproduce a chosen run's schedule. --eval_at_start is disabled by default to match the local active loop; enabling it reproduces the step-0 evaluation seen in earlier logs. The prior release unconditionally evaluated at step 0, which is now an explicit choice. The previously inspected Ball-in-Cup runs can be represented by --eval_freq 10k --eval_at_start.

The local audit parsed 33 configurations across the five supported tasks, plus three out-of-scope Cheetah configurations. With each record's interval supplied explicitly, every comparable retained hyperparameter agreed, including action repeats 2/8/4/4/4 for Finger/Cartpole/Walker Walk/Walker Stand/Ball-in-Cup, auxiliary learning rate 0.001, and batch size 128. Logs available during download contained all three shifted-domain reward columns; observed intervals agreed with their saved configurations. Unavailable or partial downloaded results were not treated as failed training runs.

Evaluation uses the same training environment between episodes, persistent shifted environments, and the current in-memory learner; it restores the learner's training mode afterward. Frequencies not aligned with an episode boundary run at the next boundary. Evaluation is recorded in eval.log and per-setting JSON files. A standalone checkpoint evaluator remains available as an optional tool.

Replay placeholder preallocation has been removed; transition storage and random sampling are preserved. A final state-dict save was added after training without inserting additional optimizer updates. Completed-episode returns are logged with the correct episode; a final partial return is labeled separately, and its log step counts all decisions in the inclusive loop.

The environment factory now honors each new instance's seed, image size, action repeat and horizon, instead of reusing the first registered instance's arguments. This corrects the requested shifted-environment seed (seed + 42) during training. The previous registration bug could instead reuse the training seed; this is an explicit environment-seeding correction, not a claim of historical evaluation-trajectory identity. Frame-stack bounds now correctly describe uint8 pixels without changing pixel values. EGL accepts a comma-separated list of numeric CUDA device indices and selects the first visible device, as the single-GPU learner does.

The release checkpoints contain network and optimizer states plus configuration, rather than pickled experimental agent classes. They support evaluation and restoring optimizer state in code; replay contents and simulator/RNG state are not included, so the command-line interface does not promise exact interrupted-training resume. Old pickled checkpoints are not directly accepted by the new evaluator.

Numerical regression

Run against an external copy of the research source:

python tests/compare_data_path.py --reference-root /path/to/research-source
python tests/compare_reference.py --reference-root /path/to/research-source --mode conv
python tests/compare_reference.py --reference-root /path/to/research-source --mode mixed

The full mixed method in the local archive can be compared without altering its code:

python tests/compare_reference.py --reference-root /path/to/local-archive \
  --reference-file src/algorithms/soda_RDSR.py --mode native

This native comparison passed updates 1000, 1001, and 1002, including convolution, SAC-only, and overlay cases. Losses, gradients, retained weights, BatchNorm buffers, Adam states, temperature, and actions agreed; the maximum parameter/buffer difference was 0.0.

The comparator loads the original class definitions while omitting unused optional imports. For mixed, it enables only the two augmentation choices inside the source's active soda_conv.py; all three original loss expressions remain unchanged. It initializes both implementations with identical retained weights and uses matched inputs and random state.

Validation on a CUDA GPU with PyTorch 1.12.1 compared updates at steps 1000, 1001, and 1002, covering an auxiliary convolution update, a SAC-only update, and an auxiliary overlay update. All compared loss values, gradients, parameters, BatchNorm buffers, Adam states, entropy temperature, and deterministic actions agreed; the maximum parameter/buffer absolute difference was 0.0. The test tolerance is rtol=1e-5, atol=1e-7.

The comparator uses controlled overlay images to isolate numerical comparison from dataset I/O. Crop coordinates, convolution outputs, overlay arithmetic, and replay samples were also compared directly with the original functions and matched exactly.

A fresh Python 3.7.6 Conda environment was created from conda-forge, without inheriting the existing research environment's packages. The documented install script and pip check passed. This caught and fixed a previously omitted xmltodict dependency. Numerical comparisons were rerun successfully in this fresh environment.

A short real-simulator run used a 64-image subset of actual Places365 images, all three losses at aux_lr=1e-3, and a shortened episode horizon. It exercised four automatic evaluation boundaries (0, 4, 8, 12), periodic and final checkpoints, evaluation JSON, and MP4 output, then continued optimization after evaluation. All five tasks passed reset/step and observation-space checks in all five supported settings (25 combinations). Repeated environment creation passed distinct-seed, same-seed reproducibility, image-size, horizon, and uint8-bound tests. A comma-separated numeric CUDA device list was exercised with EGL.

A further Walker Stand run with actual Places365 input and a shortened horizon passed automatic evaluations at steps 4, 8, and 12, with no step-0 evaluation, and continued optimization afterward. The training-flow tests cover both startup schedules.

The training-flow regression verifies exact evaluation order, reuse of the training environment, persistent shifted environments, collected transition count, warm-up update cadence, checkpoint timing, and completed/partial episode returns. The CUDA evaluation-state regression confirms that evaluation preserves parameters, buffers, Adam states, temperature, and module modes, and that subsequent training updates succeed. Checkpoint states and actions passed a round-trip check; matched uninterrupted and reloaded branches also agreed exactly over three subsequent updates on PyTorch 1.12.1. These are execution and local numerical checks, not newly reproduced research scores.

Available regression commands after installation:

python tests/test_training_flow.py
python tests/check_evaluation_state.py
python tests/check_env_instances.py
python tests/smoke_envs.py
python tests/check_checkpoint.py /path/to/release-checkpoint.pt

These checks establish the tested local computation equivalence, not bit-for-bit equivalence of independently seeded training runs, and not reproduction of the paper's five-seed learning curves. Full 500k-step benchmark training has not been rerun as part of this cleanup. The robotic manipulation code was not present in the supplied source directory and is outside this release.