I am a Ph.D. candidate in Electrical and Computer Engineering at The Ohio State University, advised by Prof. Ness B. Shroff.

My research interests are broadly in reinforcement learning, with a focus on safe exploration, offline-to-online learning, and multi-objective reinforcement learning. I'm particularly interested in developing learning algorithms that can explore efficiently while respecting safety constraints. I also work on post-training for large language models, where I've explored how reinforcement learning for tool use changes the reasoning behavior of language models, as well as methods for distributed (federated) alignment.

In summer 2026, I was a Crash Avoidance Engineering Intern at the Transportation Research Center, where I worked on autonomous vehicle safety. Before starting my Ph.D. at Ohio State, I received my B.S. and M.S. in Electrical and Computer Engineering from the University of Tehran.

View my CV (PDF)

News & Highlights

Publications

Selective Forgetting in Tool-Augmented Language-Model Agents
A. Roknilamouki, et al.
NeurIPS 2026 Workshop on Towards Test-Time Continual Learning Agents. Accepted.
Paper
Escaping Offline Pessimism: Vector-Field Reward Shaping for Safe Frontier Exploration
A. Roknilamouki, A. Ghosh, E. Ekici, and N. Shroff
Reinforcement Learning Conference (RLC), 2026.
Paper
Provably Efficient RL for Linear MDPs under Instantaneous Safety Constraints in Non-Convex Feature Spaces
A. Roknilamouki, A. Ghosh, M. Shi, F. Nourzad, E. Ekici, and N. Shroff
International Conference on Machine Learning (ICML), 2025.
Paper
Safe and Reliable Deep Reinforcement Learning for Covert Routing
A. Roknilamouki, F. T. Dagefu, N. B. Shroff, et al.
IEEE International Conference on Mobile Ad-Hoc and Smart Systems (MASS), 2025. Invited paper

Manuscripts Under Review

Thinking Beyond Tools: How Tool-Use RL Reshapes Deliberation
A. Roknilamouki, et al.
Submitted to a top-tier AI conference, 2026.
FIRM: Federated In-client Regularized Multi-objective Alignment for Large Language Models
F. Nourzad, A. Roknilamouki, E. Ekici, J. Liu, and N. Shroff
Under review at a top-tier AI conference, 2026.
arXiv

Professional Experience

Crash Avoidance Engineering Intern
Transportation Research Center (TRC) · East Liberty, OH
  • Researched autonomous-driving safety using Autoware.
  • Investigated vision-language-action (VLA) models and world models as learning-based approaches to safe and reliable autonomous driving.
Jun.–Aug. 2026

Research Experience

Graduate Research Assistant
The Ohio State University · Columbus, OH
  • Tool-augmented LLM agents: Studied how tool-use RL reshapes reasoning in Qwen and Gemma thinking models through behavioral and representation-level analyses. Developed reasoning-replay and training-time interventions to examine deliberation styles and trade-offs with tool-use reliability.
  • Offline-to-online RL and safe exploration: Developed vector-field reward shaping that combines gradient-based attraction and rotational flow to sustain exploration along uncertainty boundaries, with theoretical analysis and validation using SAC and normalizing-flow policies.
  • Provably safe RL: Developed NCS-LSVI for linear MDPs with instantaneous hard constraints in non-star-convex feature spaces. Established sublinear regret with zero safety violations with high probability, and introduced Objective–Constraint Decomposition to address a gap in prior analysis.
  • Federated LLM alignment: Developed FIRM, a communication-efficient federated RLHF framework for multiple alignment objectives, using in-client regularization to mitigate client drift and establishing finite-time convergence guarantees.
  • Deep RL for covert routing: Developed a two-phase framework combining backpressure pre-training and conservative fine-tuning to address training instability, support throughput optimization, and stabilize queue dynamics.

Education

The Ohio State University
Ph.D. in Electrical and Computer Engineering
  • Expected graduation: May 2028.
2022–present
University of Tehran
M.S. in Electrical and Computer Engineering
  • Thesis: Distributed control of autonomous surface vessels using Model Predictive Control (MPC).
  • Ranked 1st in the Electrical Engineering graduate class.
2019–2022
University of Tehran
B.S. in Electrical and Computer Engineering
2015–2019

Skills

LanguagesPython, C++, C, R, MATLAB
ML / LLM StackJAX, PyTorch, Flax, Optax, TRL, Hugging Face, scikit-learn, NumPy, pandas
MethodsSafe and deep RL; RLHF and LLM alignment; federated learning; imitation learning; pre-training
Models & ArchitecturesDiffusion transformers, sequence models, action chunking, Llama
Systems & ToolsLinux, SLURM, CUDA, Docker, Git, GitHub, Jupyter, Google Colab

Academic Activities & Service

Honors & Awards