Articulated

Learning Neural Manifold Representations for Articulated Roboticss

Note: The ECE594 project is currently limited in scope to "Aim 1" above.

Overview

This project investigates whether state-estimation objectives for articulated bodies induce structured neural representations analogous to grid codes in spatial navigation. We train recurrent networks on path integration for a robotic arm with configuration space Q = SO(3) × SO(3), analyze the learned representations, and evaluate their utility for downstream reinforcement learning.

Project Stages:

1. Body-state estimation (Team Estimation): Train RNN to perform path integration on joint angular velocities
2. Representation analysis (Team Interpretation): Visualize and analyze latent structure
3. RL evaluation (Team RL): Use learned representations as embeddings for Reacher-v5

Getting Started

1. Install Poetry

# Linux/macOS
curl -sSL https://install.python-poetry.org | python3 -

# Verify installation
poetry --version

2. Clone and Setup

git clone <repository-url>
cd articulated

# Install dependencies
poetry install

# Activate the virtual environment
poetry shell

3. Verify Installation

pytest
black --check .
ruff check .

Project Structure

articulated/
├── articulated/
│   ├── estimation/          # Team Estimation
│   │   ├── datamodule.py    # Trajectory data generation
│   │   ├── model.py         # RNN architectures
│   │   └── train.py         # Training script
│   │
│   ├── analysis/            # Team Interpretation
│   │   ├── dimensionality.py  # PCA
│   │   ├── tuning.py        # Tuning curve analysis
│   │   └── visualization.py # Plotting utilities
│   │
│   ├── rl/                  # Team RL
│   │   ├── environment.py   # Reacher-v5 wrappers
│   │   ├── agent.py         # RL agents
│   │   └── train.py         # Training script
│   │
│   ├── shared/              # Shared utilities
│   │   └── robot_arm.py     # Kinematics on SO(3) × SO(3)
│   │
│   ├── configs/             # Configuration files
│   │   ├── estimation/      # Team Estimation configs
│   │   └── rl/              # Team RL configs
│   │
│   └── notebooks/           # Analysis notebooks
│
├── tests/
├── docs/
└── pyproject.toml

---

Team Guides

Team Estimation

Goal: Train RNN to perform path integration on SO(3) × SO(3).

Key files:

• articulated/estimation/datamodule.py: Implement trajectory generation (inputs,targets)
• articulated/estimation/model.py: Define RNN architecture
• articulated/estimation/train.py: Training script

Key TODOs:

1. Implement proper SO(3) × SO(3) trajectory generation in _generate_single_trajectory()
2. Implement proper "place cell" targets on SO(3) × SO(3)
3. Experiment with RNN vs LSTM vs GRU architectures

Run training:

python -m articulated.estimation.train --config articulated/configs/estimation/rnn.yaml

Interface with other teams:

• The StateEstimationModel.get_embedding() will be used by Team RL
• The StateEstimationModel.get_hidden_states() will be used by Team Interpretation

Resources:

• Banino et al. 2018 - Grid cells in RNNs
• Path integration in neural networks

---

Team Interpretation

Goal: Analyze how the RNN encodes joint position using dimensionality reduction and tuning curves.

Key files:

• articulated/analysis/dimensionality.py: PCA
• articulated/analysis/tuning.py: Tuning curve analysis
• articulated/analysis/visualization.py: Plotting

Key TODOs:

1. Apply PCA to hidden states from trained models
2. Compute tuning curves w.r.t. configuration variables (joint angles)
3. Check if representations have structure (e.g. grid-like, etc.).

Getting hidden states:

from articulated.estimation.model import StateEstimationModel

model = StateEstimationModel.load_for_embedding("checkpoints/estimation/best.ckpt")
hidden_states = model.get_hidden_states(velocity_trajectory)

Resources:

• Banino et al. 2018
• Mechanistic interpretability

---

Team RL

Goal: Compare RL performance using raw observations vs learned embeddings on Reacher-v5.

Key files:

• articulated/rl/environment.py - Reacher wrappers
• articulated/rl/agent.py - RL agents (PPO, SAC)
• articulated/rl/train.py - Training script

Key TODOs:

1. Understand Reacher-v5 observation space
2. Implement _extract_angular_velocities() in environment wrapper
3. Run baseline (raw obs) vs embedded experiments

Run training:

# Baseline (raw observations)
python -m articulated.rl.train --config articulated/configs/rl/baseline.yaml

# With embeddings (requires trained estimation model)
python -m articulated.rl.train --config articulated/configs/rl/embedded.yaml

Resources:

• Gymnasium Reacher-v5
• Stable-Baselines3 docs
• RL tutorial video

---

Development

Code Style

black .           # Format
ruff check .      # Lint
ruff check --fix  # Auto-fix
mypy articulated/ # Type check

Running Tests

pytest
pytest -v
pytest tests/test_models.py

Contributors

Project Leads:

• Mathilde Papillon (papillon@ucsb.edu)
• Francisco Acosta (facosta@ucsb.edu) Geometric Intelligence Lab @ UCSB

ECE594 Project Team:

• Team Estimation: Awsaf Rahman (lead), Siheng Wang
• Team Interpretation: Rohan Koshy, Siheng Wang
• Team RL: Pushpita Joardar, Hun Tae Kim