Offline Reinforcement Learning Imitation Learning
# Offline Reinforcement Learning & Imitation Learning
## Introduction & Motivation
Standard reinforcement learning algorithms assume the agent can interact with its environment freely, collecting new experience through trial and error and using that fresh experience to improve its policy. This assumption is reasonable in simulated environments (video games, physics simulators) but becomes impractical or outright dangerous in many real-world settings: a robot exploring randomly in a factory can damage equipment or injure people, a recommendation system exploring randomly can seriously degrade user experience, and a medical treatment policy cannot ethically be trained through live trial and error on patients. In all of these settings, what is often available instead is a large, fixed dataset of previously logged interactions, collected under some other policy (a human operator, a prior rule-based system, or logged historical behavior), and the goal becomes learning the best possible policy from this static dataset alone, without any further interaction with the environment.
This problem, offline reinforcement learning (also called batch reinforcement learning), differs in a fundamental way from both standard online RL and standard supervised learning. Unlike online RL, the agent cannot correct a poor policy by gathering more targeted data; the dataset is fixed. Unlike standard supervised learning, the goal is not simply to imitate the behavior present in the dataset, but potentially to find a policy that outperforms the behavior that generated the data, by recombining and generalizing across the logged trajectories using the reward signal, a property known as "stitching."
Imitation learning, a closely related but distinct problem family, similarly learns from a fixed dataset of prior behavior, but without necessarily assuming access to a reward signal at all; the goal is simply to learn a policy that reproduces the behavior demonstrated in the dataset, typically collected from an expert demonstrator (a human operator or an existing high-performing controller), as closely as possible. Where offline RL asks "how good can we do, potentially better than the data-generating policy, given this fixed dataset and its rewards," imitation learning typically asks "how closely can we reproduce this expert's demonstrated behavior," and the two problem settings share substantial methodological overlap, particularly at their intersection (imitation learning augmented with reward information, or offline RL applied to a narrow expert dataset where matching, not exceeding, the demonstrator may be the practical goal).
The practical importance of both problem settings is considerable: robotics, autonomous driving, healthcare treatment policy design, industrial control, and recommendation systems all generate enormous volumes of logged interaction data as a natural byproduct of operation, and offline RL and imitation learning offer a path to leveraging this already-collected data to train capable policies without the cost, risk, or simple infeasibility of extensive live online exploration in these settings.
## Core Concepts & Theory
The central technical challenge specific to offline reinforcement learning is distributional shift between the fixed dataset's behavior policy (whatever generated the logged data) and the policy currently being learned or evaluated. Standard off-policy RL algorithms, such as Q-learning variants, involve a bootstrapping step in which the value of a state-action pair is estimated partly using the learned Q-function's own predictions at the next state, evaluated under actions selected by the current (possibly still-changing) policy. When this evaluation is restricted entirely to a fixed offline dataset, the Q-function receives no opportunity to correct systematic overestimation errors for state-action pairs poorly represented in the dataset, and the learned policy, in seeking to maximize this Q-function, tends to gravitate toward exactly these erroneously overestimated, poorly-supported actions, a failure mode often called "extrapolation error" or, when systematic, "Q-value overestimation." Left unaddressed, this compounding overestimation error can cause naive off-policy algorithms to perform catastrophically when applied to a purely offline dataset, despite performing reasonably in the online setting where the agent can correct such errors through further interaction.
Addressing this distributional shift is the organizing theme behind most offline RL algorithms, which generally fall into a few broad families: policy-constraint methods, which explicitly restrict the learned policy to stay close to the behavior policy that generated the dataset (limiting how far the policy can stray into poorly-supported, extrapolation-error-prone regions of action space); value-regularization methods, which instead directly penalize the learned Q-function for assigning high value to actions unlikely under the data-generating behavior policy, without necessarily constraining the policy's action distribution directly; and model-based methods, which learn an explicit dynamics model of the environment from the offline data and then apply planning or model-based RL techniques, typically combined with an explicit penalty for operating in regions of state-action space where the learned model is likely to be inaccurate due to limited data coverage.
Imitation learning's core theoretical concern is somewhat different: behavior cloning, the simplest imitation learning approach (directly framing the problem as supervised learning, predicting the expert's action given the observed state via a standard supervised loss), suffers from a compounding-error problem during deployment, because small prediction errors cause the learned policy's visited states to gradually drift away from the distribution of states seen in the expert demonstration data, and since the policy was never trained on these off-distribution states, its errors tend to compound over time, a phenomenon formalized in the DAgger (Dataset Aggregation) framework's analysis and addressed by DAgger's core technique of iteratively querying the expert for corrective labels on states actually visited by the current learned policy, rather than relying solely on the original, fixed expert demonstration data.
## Mathematical Formulation
A conservative approach to mitigating Q-value overestimation, exemplified by Conservative Q-Learning (CQL), modifies the standard Q-learning objective by adding a regularization term that explicitly pushes down the Q-values assigned to actions not well-supported by the offline dataset, while pushing up the Q-values of actions that are actually observed in the dataset, formulated as an additional loss term added to the standard temporal-difference Q-learning loss:
$$ L_{CQL} = \alpha \left( \mathbb{E}_{s \sim D, a \sim \pi(a|s)}[Q(s, a)] - \mathbb{E}_{s, a \sim D}[Q(s, a)] ight) + L_{TD} $$
where the first expectation is taken over actions sampled from the current learned policy (or, in a common practical variant, from a broad action distribution such as uniform, used to probe for overestimated out-of-distribution actions), the second expectation is taken over the state-action pairs actually observed in the offline dataset D, and L_TD is the standard Q-learning temporal-difference loss; the net effect of this additional term is that the learned Q-function is explicitly discouraged from assigning inflated values to actions the current policy might select but which are not well-represented in the training data, directly counteracting the extrapolation error described above.
Behavior cloning, the foundational supervised-learning approach to imitation learning, simply minimizes the negative log-likelihood of the expert's demonstrated actions under the learned policy, treating the problem as ordinary supervised classification or regression over the collected demonstration dataset D of state-action pairs from the expert:
$$ L_{BC} = -\mathbb{E}_{(s, a) \sim D_{expert}} \left[ \log \pi_ heta(a \mid s) ight] $$
which is straightforward to optimize and requires no interaction with the environment or knowledge of a reward function at all, but as discussed above, provides no mechanism for correcting the compounding-error problem that arises once the learned policy's own actions cause it to visit states outside the expert demonstration data's distribution, a fundamental limitation of the pure behavior-cloning objective that all more sophisticated imitation learning techniques attempt to address in one way or another.
## Advanced Theory & Extensions
The Decision Transformer reframes offline RL entirely as a sequence modeling problem, borrowing directly from the autoregressive Transformer architecture used in large language models: rather than learning a Q-function or an explicit policy-improvement procedure, it trains a Transformer to autoregressively predict the next action in a trajectory, conditioned on the sequence of past states, actions, and a specified target return-to-go (the desired total future reward from the current point in the trajectory onward). At deployment time, the agent is conditioned on a high desired target return, and the model, having learned the correlation between high target returns and the higher-quality trajectory segments present in the training data, generates actions that plausibly achieve returns close to that target, sidestepping the explicit dynamic-programming-based value estimation (and its associated overestimation issues) that underlies most other offline RL approaches, at the cost of being fundamentally limited by the quality of trajectories actually present in the training data (it cannot easily extrapolate to returns substantially beyond what was demonstrated in training).
Inverse reinforcement learning (IRL) takes a different approach to leveraging expert demonstrations than direct behavior cloning: rather than directly learning a policy that mimics the expert's actions, IRL attempts to infer the underlying reward function that best explains why the expert behaved as it did, under the assumption that the expert is behaving (at least approximately) optimally or near-optimally with respect to some unknown reward function. Once a plausible reward function is recovered, standard reinforcement learning techniques can then be used to train a policy that optimizes this inferred reward, potentially producing a policy that generalizes better to novel situations not directly covered by the original demonstrations than a purely imitative behavior-cloned policy would, since the recovered reward function ideally captures the underlying intent behind the expert's behavior rather than merely its surface-level action patterns.
Generative Adversarial Imitation Learning (GAIL) reframes imitation learning using an adversarial training setup directly analogous to generative adversarial networks: a discriminator is trained to distinguish state-action pairs generated by the current learned policy from those present in the expert demonstration dataset, while the policy is trained (via standard reinforcement learning, treating the discriminator's output as a reward signal) to fool this discriminator, producing behavior that is statistically indistinguishable from the expert's demonstrated behavior without requiring an explicitly recovered reward function as an intermediate step, unlike classical inverse reinforcement learning approaches.
Offline-to-online fine-tuning, an increasingly common practical strategy, trains an initial policy using offline RL on a fixed logged dataset, then continues training that policy with a limited, carefully managed amount of further online interaction, aiming to combine offline RL's ability to bootstrap from a large existing dataset with online RL's ability to correct the offline policy's remaining errors and further improve performance beyond what pure offline training alone could achieve, though this transition itself introduces its own technical challenges, since a policy and value function trained to be conservative for offline safety can be poorly calibrated for effective, sample-efficient online exploration and improvement.
## Computational Considerations
Offline RL algorithms are generally not more computationally expensive per training step than standard off-policy RL algorithms (both typically train a Q-function and a policy using minibatches sampled from a replay buffer or dataset), but they do introduce additional considerations: policy-constraint and value-regularization terms often require additional forward passes (for instance, evaluating the Q-function at multiple sampled actions to estimate the regularization term in CQL-style objectives), somewhat increasing per-step compute relative to a comparable online off-policy algorithm without these additional terms.
Model-based offline RL methods incur the additional cost of training an accurate dynamics model of the environment from the offline dataset before or alongside policy learning, and of quantifying that model's uncertainty (commonly via model ensembles) to appropriately penalize the policy for operating in poorly-modeled regions of state-action space, adding meaningful additional training compute and implementation complexity relative to purely model-free offline RL approaches, in exchange for the potential benefit of being able to generate additional synthetic rollouts from the learned model to augment the fixed, finite real offline dataset.
Decision-Transformer-style sequence-modeling approaches to offline RL inherit the computational profile of Transformer training generally, including the quadratic attention cost in trajectory sequence length, which can become a meaningful consideration for tasks involving very long-horizon trajectories, though this cost profile is well understood and benefits from the extensive existing tooling and hardware optimization developed for Transformer training in the language modeling context.
## Practical Implementation Strategies
Selecting an offline RL algorithm family for a given application should be informed by the quality and coverage of the available offline dataset: datasets generated by a single, narrow, near-expert policy (limited action diversity, little coverage of suboptimal or exploratory behavior) tend to favor methods closer to imitation learning or heavily regularized policy-constraint offline RL, since there is little useful signal in the data for the algorithm to use in identifying and stitching together better-than-demonstrated behavior, while broader, more diverse datasets (including substantial exploratory or even random behavior alongside higher-quality trajectories) tend to be more amenable to less conservative offline RL methods capable of extracting genuinely improved policies through value-based stitching across trajectory segments.
Rigorous, purely offline model selection and hyperparameter tuning remains one of the most practically difficult aspects of applying offline RL, since, unlike online RL, there is no ability to simply deploy a candidate policy and observe its actual performance during development; practitioners commonly rely on off-policy evaluation techniques (which themselves carry their own approximation and bias challenges), held-out validation trajectories, and, wherever feasible, limited, carefully controlled live evaluation of a small number of promising candidate policies before broader deployment, given the significant risk of purely offline metrics failing to reliably predict real deployed performance.
For imitation learning specifically, augmenting behavior cloning with even a modest amount of interactive expert feedback, whether through full DAgger-style iterative data aggregation or simpler heuristics for identifying and requesting expert corrections on likely problematic states, substantially improves robustness relative to pure, one-shot behavior cloning trained solely on the original static demonstration dataset, particularly for tasks with long time horizons where compounding error has the most opportunity to accumulate.
## Benchmark Datasets & Evaluation
D4RL (Datasets for Deep Data-Driven Reinforcement Learning) is the most widely used standardized benchmark suite for offline RL research, providing a broad collection of fixed datasets across diverse domains, including simulated locomotion control tasks (such as HalfCheetah, Hopper, and Walker2d from the MuJoCo physics simulator), navigation tasks (AntMaze, requiring long-horizon planning and value-stitching across a fixed maze layout), and Adroit robotic hand manipulation tasks, with datasets deliberately constructed at varying quality levels (random, medium, medium-replay, and expert-level data-generating policies) to allow systematic evaluation of how different offline RL algorithms perform across this range of data quality and coverage conditions.
Standard offline RL evaluation reports normalized average return, typically scaled so that a score of 0 corresponds to a random policy's performance and a score of 100 corresponds to an expert (or a strong, well-tuned online RL) policy's performance on the given task, measured by actually deploying the learned policy in the (typically simulated) environment after training completes purely offline, since true offline evaluation without any environment access at all remains an open and still-imperfect research problem in its own right.
For imitation learning, evaluation is most commonly performed by measuring task success rate or cumulative reward achieved by the learned policy when actually deployed in the environment, compared against the performance of the original expert demonstrator, alongside more specialized evaluation of behavioral similarity to the expert's demonstrated trajectories in settings (such as certain robotics and autonomous driving applications) where closely matching demonstrated style or behavior, not merely achieving a comparable task outcome, is itself an explicit objective.
## Key Challenges & Limitations
Distributional shift and the resulting Q-value overestimation, discussed above as the central theoretical challenge of offline RL, remains only partially solved by current techniques: policy-constraint and value-regularization methods reliably prevent catastrophic overestimation failures but can, if applied too conservatively, unnecessarily limit the algorithm's ability to improve meaningfully beyond the behavior policy present in the dataset, an inherent tension between safety (avoiding overestimation) and performance (extracting genuine improvement) that current algorithms navigate via tunable regularization strength rather than through any fully principled resolution.
Offline policy evaluation, needed both for principled hyperparameter and model selection during development and for any pre-deployment confidence assessment of a learned policy's expected real-world performance, remains a substantially unsolved problem, with existing off-policy evaluation estimators generally exhibiting either high variance (importance-sampling-based estimators, which can produce statistically unreliable estimates particularly when the evaluated policy differs substantially from the data-collection policy) or significant, difficult-to-quantify bias (model-based or value-function-based estimators, whose accuracy depends on the fidelity of a learned model or value function that itself may be poorly calibrated for exactly the out-of-distribution regions where evaluation accuracy matters most).
Both offline RL and imitation learning are fundamentally constrained by the coverage and quality of the fixed dataset used for training: no algorithm can reliably learn good behavior for situations essentially absent from the training data, meaning careful, deliberate data collection strategy (or, more commonly given the retrospective nature of most available logged datasets, careful characterization of a given dataset's actual coverage limitations) remains an important, and frequently underappreciated, practical determinant of what performance level is realistically achievable from a given offline learning problem.
## Hyperparameter Tuning
The strength of the conservatism or regularization term (such as the alpha coefficient in Conservative Q-Learning, or the strength of a policy-constraint penalty in behavior-regularized offline RL methods) is the single most consequential hyperparameter across most offline RL algorithm families, directly trading off protection against Q-value overestimation and unsafe extrapolation against the algorithm's ability to meaningfully improve upon the behavior policy present in the dataset, and this trade-off is generally dataset- and task-dependent, requiring careful tuning informed by whatever offline evaluation signal is available, given the difficulty of direct online validation discussed above.
For Decision-Transformer-style approaches, the target return-to-go value used at evaluation and deployment time functions as an important, somewhat unusual hyperparameter: setting it too close to the average return present in the training data yields conservative, merely data-typical behavior, while setting it substantially higher than any return actually observed in training data risks producing unreliable, poorly-calibrated extrapolated behavior, since the model has no direct training signal for how to act to achieve returns beyond what its training trajectories demonstrated.
For imitation learning methods incorporating interactive expert feedback (such as DAgger), the frequency and volume of new expert-labeled data collected per iteration, and the mixing ratio between newly collected corrective data and the original demonstration dataset, affect both the practical cost of the approach (since interactive expert labeling, particularly from a human expert, is typically far more expensive per-sample than simply reusing a fixed pre-collected demonstration dataset) and how quickly and reliably the learned policy's compounding-error tendencies are corrected over successive training iterations.
## Real-World Applications & Case Studies
Robotic manipulation research has extensively applied offline RL and imitation learning to leverage large logged datasets of prior robot interaction (whether from human teleoperation, scripted controllers, or prior learned policies), since collecting fresh, extensively exploratory physical robot interaction data is comparatively slow, expensive, and carries real risk of hardware damage, making the ability to extract improved policies from already-available logged interaction data particularly valuable in this domain.
Autonomous driving systems have used imitation learning, frequently augmented with DAgger-style iterative expert correction and various forms of data augmentation to improve robustness to distribution shift, to train driving policies from large logged datasets of human driving behavior, given the obvious infeasibility and danger of training a driving policy through naive online trial-and-error exploration directly on public roads.
Healthcare treatment policy research has explored offline RL for tasks such as recommending medication dosing or treatment sequencing strategies from large retrospective datasets of logged clinical records, an application area where the safety-critical, ethically constrained impossibility of live online experimentation makes offline learning from historical data essentially the only viable reinforcement-learning-based approach, though this application area also carries substantial additional challenges around dataset bias, confounding, and the high stakes of any resulting evaluation or deployment errors, requiring particular caution well beyond the standard offline RL considerations discussed above.
## Integration with Other Methods
Offline RL and imitation learning are frequently combined in practice, particularly through techniques that initialize an offline RL policy or value function using behavior cloning on the available dataset before applying more sophisticated offline RL objectives, providing a reasonable, dataset-grounded starting point that helps stabilize the subsequent offline RL optimization, particularly important given offline RL's general sensitivity to the distributional-shift issues discussed throughout this treatment.
The Decision Transformer's sequence-modeling reframing of offline RL directly connects this field to the broader large language model and Transformer architecture research discussed in dedicated treatments of those topics, and follow-on work has explored using pretrained language models, or language-model-style pretraining objectives applied to trajectory data, as a foundation for offline RL and imitation learning policies, extending techniques originally developed for natural language to the sequential decision-making domain.
Ensemble and uncertainty-quantification techniques, used throughout model-based offline RL to penalize policies for operating in poorly-modeled or poorly-covered regions of state-action space, draw directly on broader machine learning research into calibrated uncertainty estimation, and improvements in general-purpose uncertainty quantification techniques have historically translated into improved model-based offline RL algorithms as this connection has been increasingly recognized and exploited by the field.
## Future Research Directions
Developing more reliable offline policy evaluation techniques, reducing the current trade-off between high-variance importance-sampling-based estimators and potentially-biased model-based estimators, remains one of the most practically important open problems in the field, since progress here would directly improve practitioners' ability to confidently select among candidate policies and hyperparameter configurations without risky or expensive live deployment testing.
Better understanding and improving the "stitching" capability of offline RL algorithms, their ability to recombine partial trajectory segments from the offline dataset into novel, higher-quality behavior not directly demonstrated by any single trajectory in the training data, remains an active theoretical and empirical research question, particularly for the newer generation of sequence-modeling-based approaches like the Decision Transformer, whose ability to perform this kind of stitching (as opposed to more standard value-based offline RL methods) is less well established and an active subject of ongoing research and debate.
Improving the safety and reliability of offline-to-online fine-tuning, so that policies trained conservatively offline can be further improved through limited online interaction without a dangerous initial performance collapse or unsafe exploratory behavior during the transition, is an increasingly important research direction as offline RL techniques move from primarily academic benchmark evaluation toward real-world deployment pipelines that plausibly involve at least some capacity for further online refinement after an initial offline training phase.
## Summary & Key Takeaways
Offline reinforcement learning trains policies purely from a fixed, previously collected dataset without further environment interaction, requiring specific techniques (policy constraints, value regularization, or uncertainty-penalized model-based planning) to counteract the distributional-shift-driven Q-value overestimation that causes naive off-policy algorithms to fail badly in this purely offline setting.
Imitation learning learns to reproduce expert-demonstrated behavior, with plain behavior cloning suffering from a compounding-error problem addressed by techniques ranging from interactive expert correction (DAgger) to adversarial imitation methods (GAIL) to inverse reinforcement learning's approach of first recovering an underlying reward function before applying standard RL.
The Decision Transformer illustrates an alternative, sequence-modeling-based framing of offline RL that sidesteps explicit value-function bootstrapping entirely, conditioning autoregressive action prediction on a specified target return, at the cost of being fundamentally bounded by the quality of returns actually represented in its training data.
Reliable purely-offline policy evaluation and hyperparameter selection remain among the field's most significant open practical challenges, and dataset quality and coverage, more than any specific algorithmic choice, remain the dominant determinant of what performance is realistically achievable in any given offline learning problem.
Keywords: offline reinforcement learning, batch reinforcement learning, imitation learning, behavior cloning, distributional shift, Q-value overestimation, extrapolation error, Conservative Q-Learning, CQL, policy constraint methods, Decision Transformer, return-to-go conditioning, DAgger, dataset aggregation, inverse reinforcement learning, GAIL, adversarial imitation learning, D4RL benchmark, offline policy evaluation, offline-to-online fine-tuning
---
## Appendix: Practical Labs
### Lab 1: Demonstrating Q-Value Overestimation from Purely Offline, Narrow Data
import numpy as np
np.random.seed(0)
def true_q_function(state, action):
"""Ground-truth Q-values for a small synthetic MDP: 5 states, 4 actions,
with a known, fixed optimal Q-value surface for evaluation purposes."""
return -np.abs(action - (state % 4)) * 2.0 + state * 0.5
def generate_narrow_offline_dataset(n_states=5, n_actions=4, n_samples=200, seed=0):
"""Simulates a narrow behavior policy that mostly takes only ONE action per state,
a common realistic scenario (e.g., a deterministic scripted controller's logs)."""
rng = np.random.RandomState(seed)
states = rng.randint(0, n_states, size=n_samples)
# Narrow behavior policy: action = state % n_actions almost always, rare exploration
actions = np.where(
rng.rand(n_samples) < 0.9,
states % n_actions,
rng.randint(0, n_actions, size=n_samples),
)
rewards = true_q_function(states, actions) + rng.randn(n_samples) * 0.5
return states, actions, rewards
def fit_naive_q_table(states, actions, rewards, n_states, n_actions, n_epochs=200, lr=0.3):
"""Fits a simple tabular Q-function via supervised regression toward observed
rewards ONLY at (state, action) pairs present in the dataset -- unseen pairs
keep their random initialization, simulating unconstrained function approximation
extrapolating poorly to out-of-distribution actions."""
Q = np.random.randn(n_states, n_actions) * 3.0 # random init -> can be arbitrarily "confident"
counts = np.zeros((n_states, n_actions))
for _ in range(n_epochs):
for s, a, r in zip(states, actions, rewards):
Q[s, a] += lr * (r - Q[s, a]) * 0.1
counts[s, a] += 1
return Q, counts
def test_overestimation_on_unseen_actions():
n_states, n_actions = 5, 4
states, actions, rewards = generate_narrow_offline_dataset(n_states, n_actions, n_samples=300)
Q_learned, counts = fit_naive_q_table(states, actions, rewards, n_states, n_actions)
true_Q = np.array([[true_q_function(s, a) for a in range(n_actions)] for s in range(n_states)])
seen_mask = counts > 0
unseen_mask = counts == 0
seen_error = np.abs(Q_learned[seen_mask] - true_Q[seen_mask]).mean()
unseen_error = np.abs(Q_learned[unseen_mask] - true_Q[unseen_mask]).mean()
print(f"Mean abs Q-value error on WELL-COVERED (state,action) pairs: {seen_error:.3f}")
print(f"Mean abs Q-value error on UNSEEN (state,action) pairs: {unseen_error:.3f}")
print(f"Fraction of (state,action) pairs never observed in dataset: {unseen_mask.mean():.2%}")
assert unseen_mask.sum() > 0, "Narrow dataset should leave some (state,action) pairs unobserved"
assert unseen_error > seen_error, (
"Q-value error should be substantially higher for unseen (state,action) pairs, "
"demonstrating the extrapolation error central to offline RL"
)
print("Offline Q-value overestimation demonstration test passed.")
if __name__ == "__main__":
test_overestimation_on_unseen_actions()### Lab 2: Conservative Q-Learning-Style Regularization Suppressing Out-of-Distribution Q-Values
import numpy as np
np.random.seed(1)
def softmax(x):
shifted = x - np.max(x, axis=-1, keepdims=True)
exp = np.exp(shifted)
return exp / exp.sum(axis=-1, keepdims=True)
def cql_regularized_q_update(Q, states, actions, rewards, n_actions, alpha=1.0, lr=0.1):
"""
Applies one epoch of a simplified CQL-style update:
1. Standard regression toward observed (s, a, r) in the dataset.
2. An explicit penalty pushing DOWN Q-values for actions NOT in the dataset
at each visited state, and pushing UP Q-values for the actually-observed action,
approximating the CQL conservative penalty term.
"""
n_states = Q.shape[0]
dataset_actions_per_state = {s: set() for s in range(n_states)}
for s, a in zip(states, actions):
dataset_actions_per_state[s].add(a)
# Standard TD-style regression toward observed reward
for s, a, r in zip(states, actions, rewards):
Q[s, a] += lr * (r - Q[s, a])
# Conservative penalty: for every visited state, push down Q for ALL actions
# weighted by how "in-distribution" they are (uniform policy proxy here),
# and push up Q for the actually observed dataset action
for s in dataset_actions_per_state:
policy_probs = softmax(Q[s]) # proxy for "current policy" over actions
for a in range(n_actions):
if a in dataset_actions_per_state[s]:
Q[s, a] += lr * alpha * 0.1 # push up in-distribution action value slightly
else:
Q[s, a] -= lr * alpha * policy_probs[a] * 2.0 # push down OOD action value
return Q
def test_cql_penalty_suppresses_ood_actions():
n_states, n_actions = 5, 4
# Narrow dataset: state s always paired with action (s % n_actions)
states = np.repeat(np.arange(n_states), 40)
actions = states % n_actions
rewards = np.ones_like(states, dtype=float) * 2.0
Q_naive = np.random.randn(n_states, n_actions) * 3.0
Q_cql = Q_naive.copy()
# Naive update: no conservative penalty (alpha=0)
for _ in range(50):
Q_naive = cql_regularized_q_update(Q_naive, states, actions, rewards, n_actions, alpha=0.0)
# CQL-style update: with conservative penalty
for _ in range(50):
Q_cql = cql_regularized_q_update(Q_cql, states, actions, rewards, n_actions, alpha=1.0)
# Compare average Q-value assigned to OUT-OF-DISTRIBUTION actions (not state % n_actions)
ood_mask = np.ones((n_states, n_actions), dtype=bool)
for s in range(n_states):
ood_mask[s, s % n_actions] = False
naive_ood_avg = Q_naive[ood_mask].mean()
cql_ood_avg = Q_cql[ood_mask].mean()
print(f"Average Q-value on OOD actions (no conservative penalty): {naive_ood_avg:.3f}")
print(f"Average Q-value on OOD actions (with CQL-style penalty): {cql_ood_avg:.3f}")
assert cql_ood_avg < naive_ood_avg, (
"CQL-style conservative penalty should result in lower Q-values for "
"out-of-distribution actions compared to an unregularized update"
)
print("CQL-style conservative regularization test passed.")
if __name__ == "__main__":
test_cql_penalty_suppresses_ood_actions()### Lab 3: Behavior Cloning Compounding Error vs. DAgger-Style Correction
import numpy as np
np.random.seed(2)
class Simple1DEnv:
"""A toy 1D navigation environment: agent moves along a line, expert policy
always moves toward a fixed target position. The environment terminates
after a fixed horizon."""
def __init__(self, target=10.0, horizon=30):
self.target = target
self.horizon = horizon
def expert_action(self, position):
# Expert always takes a small step directly toward the target
return np.clip(self.target - position, -1.0, 1.0)
def rollout_with_policy(self, policy_fn, start_position=0.0, noise_std=0.0):
position = start_position
positions = [position]
for _ in range(self.horizon):
action = policy_fn(position) + np.random.randn() * noise_std
position += action
positions.append(position)
return np.array(positions)
def fit_linear_behavior_clone(states, actions):
"""Fits a simple linear policy action = w * state + b via least squares."""
A = np.stack([states, np.ones_like(states)], axis=1)
w, b = np.linalg.lstsq(A, actions, rcond=None)[0]
return lambda s: w * s + b
def test_bc_compounding_error_vs_dagger_correction():
env = Simple1DEnv(target=10.0, horizon=30)
# --- Plain behavior cloning: train ONLY on states near the start (0 to 2) ---
narrow_states = np.linspace(0, 2, 30)
narrow_actions = np.array([env.expert_action(s) for s in narrow_states])
bc_policy = fit_linear_behavior_clone(narrow_states, narrow_actions)
bc_trajectory = env.rollout_with_policy(bc_policy, start_position=0.0, noise_std=0.05)
bc_final_error = abs(bc_trajectory[-1] - env.target)
# --- DAgger-style correction: after an initial BC rollout, query the expert
# on the ACTUAL states visited by the learned policy, and retrain including this data ---
all_states = list(narrow_states)
all_actions = list(narrow_actions)
dagger_policy = bc_policy
for dagger_iter in range(4):
rollout = env.rollout_with_policy(dagger_policy, start_position=0.0, noise_std=0.05)
# Query expert on states actually visited (this is the key DAgger step)
visited_states = rollout[:-1]
expert_labels = np.array([env.expert_action(s) for s in visited_states])
all_states.extend(visited_states.tolist())
all_actions.extend(expert_labels.tolist())
dagger_policy = fit_linear_behavior_clone(np.array(all_states), np.array(all_actions))
dagger_trajectory = env.rollout_with_policy(dagger_policy, start_position=0.0, noise_std=0.05)
dagger_final_error = abs(dagger_trajectory[-1] - env.target)
print(f"Plain BC (trained only on states [0,2]) final position error: {bc_final_error:.3f}")
print(f"DAgger-corrected policy final position error: {dagger_final_error:.3f}")
assert dagger_final_error < bc_final_error, (
"DAgger-style iterative correction on actually-visited states should reduce "
"final trajectory error relative to plain behavior cloning trained on a "
"narrow slice of state space"
)
print("Behavior cloning vs. DAgger compounding-error test passed.")
if __name__ == "__main__":
test_bc_compounding_error_vs_dagger_correction()### Lab 4: Return-to-Go Conditioning for Decision-Transformer-Style Action Selection
import numpy as np
np.random.seed(3)
def compute_returns_to_go(rewards):
"""Computes the return-to-go at each timestep: the sum of all rewards from
that timestep until the end of the trajectory."""
returns_to_go = np.zeros_like(rewards, dtype=float)
running_sum = 0.0
for t in reversed(range(len(rewards))):
running_sum += rewards[t]
returns_to_go[t] = running_sum
return returns_to_go
def generate_toy_trajectories(n_trajectories=50, horizon=10, seed=0):
"""Generates trajectories of varying quality: each trajectory has a fixed
'skill level' that determines its typical per-step reward, simulating a
dataset containing a mix of low- and high-quality demonstrated behavior."""
rng = np.random.RandomState(seed)
dataset = []
for _ in range(n_trajectories):
skill_level = rng.uniform(0.0, 1.0) # 0 = poor, 1 = expert
states = rng.randn(horizon)
actions = skill_level * 2.0 + rng.randn(horizon) * (1.0 - skill_level) * 0.5
rewards = skill_level * 1.0 + rng.randn(horizon) * 0.1
returns_to_go = compute_returns_to_go(rewards)
dataset.append({'states': states, 'actions': actions, 'rtg': returns_to_go})
return dataset
def fit_rtg_conditioned_policy(dataset):
"""Fits a simple linear model: action = f(state, return_to_go), approximating
the return-to-go-conditioned action prediction used by the Decision Transformer,
via least squares regression over the pooled dataset."""
all_states, all_rtg, all_actions = [], [], []
for traj in dataset:
all_states.extend(traj['states'])
all_rtg.extend(traj['rtg'])
all_actions.extend(traj['actions'])
X = np.stack([all_states, all_rtg, np.ones(len(all_states))], axis=1)
y = np.array(all_actions)
coeffs, _, _, _ = np.linalg.lstsq(X, y, rcond=None)
w_state, w_rtg, bias = coeffs
def policy(state, target_rtg):
return w_state * state + w_rtg * target_rtg + bias
return policy, w_rtg
def test_higher_target_return_yields_more_expert_like_action():
dataset = generate_toy_trajectories(n_trajectories=80, horizon=10, seed=1)
policy, w_rtg = fit_rtg_conditioned_policy(dataset)
test_state = 0.0
low_target_action = policy(test_state, target_rtg=1.0) # low target return (poor-skill regime)
high_target_action = policy(test_state, target_rtg=9.0) # high target return (expert regime)
print(f"Learned coefficient on return-to-go (w_rtg): {w_rtg:.4f}")
print(f"Action predicted with LOW target return-to-go: {low_target_action:.3f}")
print(f"Action predicted with HIGH target return-to-go: {high_target_action:.3f}")
# Since higher-skill trajectories in the synthetic dataset have both higher
# returns-to-go AND systematically larger action magnitudes, the fitted policy
# should learn a positive association: higher requested return -> more
# "expert-like" (larger magnitude) predicted action.
assert high_target_action > low_target_action, (
"Conditioning on a higher target return-to-go should shift the predicted "
"action toward the higher-skill (higher action magnitude) regime observed "
"in the training trajectories"
)
assert w_rtg > 0, "Return-to-go coefficient should be positive given the synthetic data's construction"
print("Return-to-go conditioning test passed.")
if __name__ == "__main__":
test_higher_target_return_yields_more_expert_like_action()