jamesheald/joint-space-empowerment

Joint-Space Empowerment for Dexterous Coordination in Tendon-Driven Hands

This repository contains the official trained policies for:

Joint-Space Empowerment for Dexterous Coordination in Tendon-Driven Hands
James Heald, Vittorio Caggiano, Vikash Kumar, Maneesh Sahani
ICML 2026 (Spotlight)
📜 Paper | 💻 Code | 🌐 Project Page

---

📂 Repository Structure

The files are organized by task environment and training seed:

joint-space-empowerment
├── Adroit
│   ├── BaodingBalls
│   │   ├── seed_0
│   │   │   ├── rl_model_10000000_steps.zip                    # Trained model
│   │   │   └── rl_model_vecnormalize_10000000_steps.pkl       # Observation normalization statistics
│   │   ├── seed_1  …  seed_4
│   ├── DieReorient
│   │   └── ...
│   ├── KeyTurn
│   │   └── ...
│   └── PenTwirl
│       └── ...
├── MyoHand
│   ├── BaodingBalls
│   │   └── ...
│   ├── DieReorient
│   │   └── ...
│   ├── KeyTurn
│   │   └── ...
│   ├── PenTwirl
│   │   └── ...
│   ├── Reorient100
│   │   ├── Pretraining_Play
│   │   │   ├── seed_0
│   │   │   │   ├── rl_model_replay_buffer_1000000_steps.pkl   # Replay buffer
│   │   │   │   └── rl_model_vecnormalize_1000000_steps.pkl
│   │   │   ├── seed_1  …  seed_4
│   │   └── Training
│   │       └── ...
│   └── Reorient8-sparse
│       └── ...
└── README.md

Loading Example

Policies were trained using Stable-Baselines3 with observation normalization.

from huggingface_hub import hf_hub_download
from myosuite.utils import gym
from src.agents.josepi.sac import SAC as JoSEPi
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize

REPO_ID = "jamesheald/joint-space-empowerment"
SUBFOLDER = "MyoHand/BaodingBalls/seed_0"
MODEL_FILE = "rl_model_10000000_steps.zip"
STATS_FILE = "rl_model_vecnormalize_10000000_steps.pkl"

# 1. Download SB3 zip model and VecNormalize pickle stats
model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILE, subfolder=SUBFOLDER)
stats_path = hf_hub_download(repo_id=REPO_ID, filename=STATS_FILE, subfolder=SUBFOLDER)

# 2. Reconstruct base vectorized environment
raw_env = DummyVecEnv([lambda: gym.make("CustomMyoChallengeBaodingP1-v1", normalize_act=True)])

# 3. Apply normalization stats and freeze updating
env = VecNormalize.load(stats_path, raw_env)
env.training = False     # Freeze running mean/var statistics
env.norm_reward = False  # Return unnormalized environment rewards

# 4. Load SB3 policy
model = JoSEPi.load(model_path, env=env)

# 5. Run inference
obs = env.reset()
while True:
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, done, info = env.step(action)
    if done[0]:
        break