TSSM Agent for Agent Arena

This is an ML research project exploring latent dynamics for deformable manipulation. As is common in research code, I focus on reproducible scripts, metrics, and logged runs rather than unit tests and production-grade infrastructure.

The transDreamer folder contains the transformer based model, which will be trained using the garment dataset. To use this, transdreamer_adapter.py implements all modifications necessary.

The scripts folder contains all executable python scripts used for each milestone.

Setup

To set up the project, clone halid1020/softgym and halid1020/agent-arena-v0 into the root of this project.

The resulting structure should look as follows:

project/
├── agent-arena-v0/
├── softgym/
├── tssm_agent/
│   ├── transdreamer/
│   ├── configs/ 
│   ├── (...)         
│   └── scripts/
├── (...) 
├── setup.sh   
└── README.md

Follow the setup instructions for softgym and agent-arena-v0, after this, the scripts in tssm_agent/ should be able to access the agent-arena API. To test this, run test_api.py.

When inside the agent-arena conda environment, run test_dependencies.py to check if pytorch, cuda, openGL and PyBullet are working.

Install skikit-image for evaluation functions:

conda install -c conda-forge scikit-image

Load the datasets into /data, create the folder if it doesn't exist.

Softgym

Follow the instructions in halid1020/softgym.

Agent-Arena

Follow the instructions in halid1020/agent-arena-v0.

TSSM Agent

Run [./setup]

Scripts

The scripts are used to separate the work into components, each for a specific functionality. See scripts subfolder.

viewing training logs with tensorboard

To view training logs and metrics using TensorBoard:

tensorboard --logdir tssm_agent/models/transdreamer/transdreamer_garment_rgb --load_fast=false 

Then open http://localhost:6006 in your browser to view the training progress, loss curves, and other metrics.

Overview (objectives and executables)

• Reconstruct trajectories (prior/posterior/reward): Use tssm_agent/scripts/visualise_offline.py (PlaNet/RSSM) and tssm_agent/scripts/visualise_td_offline.py (Transformer-based) to generate reconstructions on the garment test set.
• Quantitative metrics for predictions/latents: Use tssm_agent/scripts/evaluate_offline.py (baseline) and tssm_agent/scripts/evaluate_td_offline.py (Transformer). Dataset stats via tssm_agent/scripts/dataset_info.py.
• Transformer in place of RNN (PlaNet prior compatibility): See implementation in tssm_agent/transdreamer/ with integration via tssm_agent/transdreamer/transdreamer_adapter.py.
• Integration and training: Train/evaluate baseline with tssm_agent/scripts/train_offline.py and the Transformer model with tssm_agent/scripts/train_td_offline.py.

Secondary goals:

• Train/evaluate on garment dataset: train_offline.py, train_td_offline.py, and corresponding evaluate_* scripts.
• Compare latent representation quality: Compare metrics/logs from evaluate_offline.py vs evaluate_td_offline.py (see TensorBoard section above).
• MPC policy testing in simulation: tssm_agent/scripts/evaluate_rssm_planning.py (PlaNet) vs tssm_agent/scripts/evaluate_td_planning.py (Transformer).

For executable examples and exact run commands, see tssm_agent/scripts/readme.md.