Joey Hejna

The technical report for Composer 2, Cursor's coding model.

Importance Weighted Retrieval (IWR) is a retrieval method that estimates importance weights, or the ratio between the target and prior data distributions, using Gaussian KDEs on top of existing retrieval methods.

We use VLM generated plans to guide low-level RL for dexterous manipulation tasks.

We introduce a new method for robot data curation that uses mutual information to select the most informative demonstrations. We show that this approach outperforms existing methods in both simulation and real-world settings.

This paper introduces a novel approach that leverages priors generated by pre-trained LLMs alongside the precision of preference learning. Our method, termed Language-Guided Preference Learning (LGPL), uses LLMs to generate initial behavior samples, which are then refined through preference-based feedback to learn behaviors that closely align with human expectations.

GVL poses value estimation as a temporal ordering problem over shuffled video frames; this seemingly more challenging task encourages VLMs to more fully exploit their underlying semantic and temporal grounding capabilities to differentiate frames based on their perceived task progress, consequently producing significantly better value predictions.

We introduce DITTO, an algorithm for few-shot adaptation or alignment of large language models using only two or three human demonstrations. We show that DITTO outperforms prior methods by a significant margin on automated evals and a real-world user study.

We introduce ReMix, a recipe for automatically curating data mixtures for training large scale robotic imitation learning policies. We demonstrate improvements over human-curated and uniform dataset mixtures.

We investigate how well success-filtered autonomous imitation can work in real-world robotics settings.

Evaluating robot motions involves more than just the start and end states; it's about how the task is performed. We propose motion instruction fine-tuning (MotIF) and the MotIF-1K dataset to improve VLMs' ability to understand nuanced robotic motions.

Through extensive experimentation and multiple model scales we characterize the over-optimization problem for direct alignment algorithms in large language models.

A large diverse dataset collected across multiple institutions totaling over 75K robot trajectories.

Understanding direct alignment algorithms for LLMs through the perspective of a per-token Markov Decision Process.

Laying the groundwork for a foundation model for robotics.

Reduces Reinforcement Learning from Human Feedback to contrastive learning under the regret model of human preferences, which has recently been shown to be more accurate than the widely accepted reward model. Unlike many RLHF methods, CPL is fully off-policy and works on arbitrary MDPs.

Approaches to preference-based RL typically work in two phases: first a reward function is learned, then it is maximized using a vanilla RL algorithm. We introduce the Inverse Preference Learning framework, where we directly learn a Q-function that models the user's preferences without explicitly learning a reward function.

We introduce DWSL, an algorithm for offline goal-conditioned reinforcement learning that uses only supervised objectives while still learning a constrained optimal policy. DWSL performs particularly well on high-dimensional image domains and seems robust to hyperparameters.

We introduce a novel framework for Q-learning that models the maximal soft-values without needing to sample from a policy and improves performance in online and offline RL settings.

Pretraining preference models greatly reduces query-complexity, enabling humans to teach robots with a reasonable amount of feedback.

We show that predicting instructions along with actions drastically improves performance in combinatorially complex long-horizon imitation settings.

Better robot structures hold the promise of better performance. We propose a new algorithm, TAME, that is able to evolve morphologies without any task specification. This is accomplished using an information theoretic objective that efficiently ranks morphologies based on their ability to explore and control their environment.

We propose transferring RL policies across agents using a hierarchical framework. Then, to remedy poor zero-shot transfer performance we introduce two additional imitation objectives.

A lightweight framework for general deep-learning research in PyTorch.

A repository for training large behavior models in Jax/Flax/tfds.

We examine and compare three methods for explicitly disentangling learned latent representations in VAE models.

Training the models that train Composer, Cursor's coding model.

Working on data collection, labeling, and training more reactive policies.

Working on the machine learning and robotics team.

Developed C++ systems for trading APIs and monitoring systems. Worked on optimizing memory usage of large model training.

Worked on demo systems for Intel's OpenVINO model optimization system in the AWS DeepLens. Explored systems for gradient based explanations of deep networks.

NeurIPS 2023 (Outstanding Reviewer), ICML 2024 (Outstanding Reviewer), ICML 2023, ICLR 2025, CoRL 2024, ICRA 2025, IEEE RA-L, RL-Brew at RLC 2024

Course Assistant, CS 221: Artificial Intelligence, Autumn 2024

Head Course Assistant, CS 221: Artificial Intelligence, Autumn 2023

Teaching Assistant, EECS 127: Optimization Models, Fall 2020

Teaching Assistant, EECS 189: Machine Learning, Spring 2020

Teaching Assistant, CS 70: Discrete Math and Probability Theory, Fall 2019

Introductory ML Notes

Deep Learning Workshop

Reinforcement Learning Workshop