Soft Geometric Inductive Bias for Object Centric Dynamics

This repository contains the official code for

"Soft Geometric Inductive Bias for Object Centric Dynamics"

by Hampus Linander, Conor Heins, Marco Perin, Alexander Tschantz, Christopher L. Buckley

Dependencies

Install uv.

Run entry points with uv run, e.g.

uv run pytest

Create an explicit venv

uv venv
uv sync
source .venv/bin/activate

CUDA

Install the kinetixcuda group

Quick usage

# One‑off command in a throw‑away env
uv run --group kinetixcuda --no-default-groups -m pytest

Or, in a dedicated virtual environment:

# create a dedicated venv once
uv venv .venv_cuda

# activate it and sync just the jaxcuda stack + your project
source .venv_cuda/bin/activate
uv sync --active --group kinetixcuda --no-default-groups

Run experiments

To see all available parameters use uv run python src/experiments/train.py --help

Environments (specified with --level-path) located in src/experiments/envs.

Generating plots

Prerequisite: all sweeps have been run For each folder in plots/ there is a script to download the raw metric data from WandB, start by running

uv run python download_data.py

Then create a merged database using duckdb -f create_db.sql

After this, the notebooks can be executed to generate the plots.