@elaraai/east-py-datascience 0.0.2-beta.32 → 0.0.2-beta.34

package/dist/lightning/lightning.d.ts

@@ -13,6 +13,33 @@
  */
 import { StructType, VariantType, OptionType, IntegerType, FloatType, BlobType, ArrayType, NullType, BooleanType, FunctionType, StringType } from "@elaraai/east";
 export { VectorType, MatrixType } from "../types.js";
+/**
+ * Return embedding mode for Decision Transformer.
+ */
+export declare const ReturnEmbeddingType: VariantType<{
+    /** Single return value for entire sequence */
+    global: NullType;
+    /** Return-to-go at each timestep */
+    per_timestep: NullType;
+}>;
+/**
+ * Per-head output configuration for multi_head_mixed.
+ */
+export declare const HeadConfigType: StructType<{
+    /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+    head_type: VariantType<{
+        /** Single binary output: 1 logit, sigmoid, BCE loss */
+        binary: NullType;
+        /** Multi-class output: n_classes logits, softmax, CE loss */
+        multiclass: StructType<{
+            n_classes: IntegerType;
+        }>;
+    }>;
+    /** Optional class weights for this head */
+    class_weights: OptionType<ArrayType<FloatType>>;
+    /** Optional: index of head this depends on (loss only computed when that head is 1) */
+    conditional_on: OptionType<IntegerType>;
+}>;
 /**
  * Lightning output mode - determines loss function and output activation.
  */
@@ -40,6 +67,31 @@ export declare const LightningOutputType: VariantType<{
         /** Optional class weights matrix (n_heads, n_classes) */
         class_weights: OptionType<ArrayType<ArrayType<FloatType>>>;
     }>;
+    /**
+     * Mixed output types per head.
+     * For Decision Transformer: combines binary (1 logit) and multiclass (n_classes logits) heads.
+     * Binary heads: 1 logit → sigmoid → BCE loss
+     * Multiclass heads: n_classes logits → softmax → CE loss
+     * Action vectors use one-hot encoding for multiclass heads.
+     */
+    multi_head_mixed: StructType<{
+        /** Array of head configurations */
+        heads: ArrayType<StructType<{
+            /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+            head_type: VariantType<{
+                /** Single binary output: 1 logit, sigmoid, BCE loss */
+                binary: NullType;
+                /** Multi-class output: n_classes logits, softmax, CE loss */
+                multiclass: StructType<{
+                    n_classes: IntegerType;
+                }>;
+            }>;
+            /** Optional class weights for this head */
+            class_weights: OptionType<ArrayType<FloatType>>;
+            /** Optional: index of head this depends on (loss only computed when that head is 1) */
+            conditional_on: OptionType<IntegerType>;
+        }>>;
+    }>;
 }>;
 /**
  * Cell type for sequential architectures.
@@ -122,6 +174,36 @@ export declare const LightningArchitectureType: VariantType<{
         /** Optional condition dimension for conditional generation */
         condition_dim: OptionType<IntegerType>;
     }>;
+    /**
+     * Decision Transformer: return-conditioned sequence generation.
+     * Token layout: [R, s_0, a_0, s_1, a_1, ..., s_{T-1}, a_{T-1}]
+     * Predicts actions conditioned on desired return and state history.
+     */
+    decision_transformer: StructType<{
+        /** Sequence length (timesteps) */
+        sequence_length: IntegerType;
+        /** State dimension per timestep */
+        state_dim: IntegerType;
+        /** Action dimension per timestep */
+        action_dim: IntegerType;
+        /** Model dimension (transformer hidden size) */
+        d_model: IntegerType;
+        /** Number of attention heads */
+        n_attention_heads: IntegerType;
+        /** Number of transformer layers */
+        n_layers: IntegerType;
+        /** Feedforward dimension (default: 4 * d_model) */
+        d_ff: OptionType<IntegerType>;
+        /** Dropout rate */
+        dropout: OptionType<FloatType>;
+        /** Whether return is per-timestep or global */
+        return_embedding: VariantType<{
+            /** Single return value for entire sequence */
+            global: NullType;
+            /** Return-to-go at each timestep */
+            per_timestep: NullType;
+        }>;
+    }>;
 }>;
 /**
  * Epoch callback function type: (epoch, train_loss, val_loss) -> void
@@ -203,6 +285,36 @@ export declare const LightningConfigType: StructType<{
             /** Optional condition dimension for conditional generation */
             condition_dim: OptionType<IntegerType>;
         }>;
+        /**
+         * Decision Transformer: return-conditioned sequence generation.
+         * Token layout: [R, s_0, a_0, s_1, a_1, ..., s_{T-1}, a_{T-1}]
+         * Predicts actions conditioned on desired return and state history.
+         */
+        decision_transformer: StructType<{
+            /** Sequence length (timesteps) */
+            sequence_length: IntegerType;
+            /** State dimension per timestep */
+            state_dim: IntegerType;
+            /** Action dimension per timestep */
+            action_dim: IntegerType;
+            /** Model dimension (transformer hidden size) */
+            d_model: IntegerType;
+            /** Number of attention heads */
+            n_attention_heads: IntegerType;
+            /** Number of transformer layers */
+            n_layers: IntegerType;
+            /** Feedforward dimension (default: 4 * d_model) */
+            d_ff: OptionType<IntegerType>;
+            /** Dropout rate */
+            dropout: OptionType<FloatType>;
+            /** Whether return is per-timestep or global */
+            return_embedding: VariantType<{
+                /** Single return value for entire sequence */
+                global: NullType;
+                /** Return-to-go at each timestep */
+                per_timestep: NullType;
+            }>;
+        }>;
     }>;
     /** Output mode (determines loss function) */
     output: VariantType<{
@@ -229,6 +341,31 @@ export declare const LightningConfigType: StructType<{
             /** Optional class weights matrix (n_heads, n_classes) */
             class_weights: OptionType<ArrayType<ArrayType<FloatType>>>;
         }>;
+        /**
+         * Mixed output types per head.
+         * For Decision Transformer: combines binary (1 logit) and multiclass (n_classes logits) heads.
+         * Binary heads: 1 logit → sigmoid → BCE loss
+         * Multiclass heads: n_classes logits → softmax → CE loss
+         * Action vectors use one-hot encoding for multiclass heads.
+         */
+        multi_head_mixed: StructType<{
+            /** Array of head configurations */
+            heads: ArrayType<StructType<{
+                /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+                head_type: VariantType<{
+                    /** Single binary output: 1 logit, sigmoid, BCE loss */
+                    binary: NullType;
+                    /** Multi-class output: n_classes logits, softmax, CE loss */
+                    multiclass: StructType<{
+                        n_classes: IntegerType;
+                    }>;
+                }>;
+                /** Optional class weights for this head */
+                class_weights: OptionType<ArrayType<FloatType>>;
+                /** Optional: index of head this depends on (loss only computed when that head is 1) */
+                conditional_on: OptionType<IntegerType>;
+            }>>;
+        }>;
     }>;
     /** Learning rate (default: 1e-3) */
     learning_rate: OptionType<FloatType>;
@@ -401,6 +538,36 @@ export declare const lightning_train: import("@elaraai/east").PlatformDefinition
             /** Optional condition dimension for conditional generation */
             condition_dim: OptionType<IntegerType>;
         }>;
+        /**
+         * Decision Transformer: return-conditioned sequence generation.
+         * Token layout: [R, s_0, a_0, s_1, a_1, ..., s_{T-1}, a_{T-1}]
+         * Predicts actions conditioned on desired return and state history.
+         */
+        decision_transformer: StructType<{
+            /** Sequence length (timesteps) */
+            sequence_length: IntegerType;
+            /** State dimension per timestep */
+            state_dim: IntegerType;
+            /** Action dimension per timestep */
+            action_dim: IntegerType;
+            /** Model dimension (transformer hidden size) */
+            d_model: IntegerType;
+            /** Number of attention heads */
+            n_attention_heads: IntegerType;
+            /** Number of transformer layers */
+            n_layers: IntegerType;
+            /** Feedforward dimension (default: 4 * d_model) */
+            d_ff: OptionType<IntegerType>;
+            /** Dropout rate */
+            dropout: OptionType<FloatType>;
+            /** Whether return is per-timestep or global */
+            return_embedding: VariantType<{
+                /** Single return value for entire sequence */
+                global: NullType;
+                /** Return-to-go at each timestep */
+                per_timestep: NullType;
+            }>;
+        }>;
     }>;
     /** Output mode (determines loss function) */
     output: VariantType<{
@@ -427,6 +594,31 @@ export declare const lightning_train: import("@elaraai/east").PlatformDefinition
             /** Optional class weights matrix (n_heads, n_classes) */
             class_weights: OptionType<ArrayType<ArrayType<FloatType>>>;
         }>;
+        /**
+         * Mixed output types per head.
+         * For Decision Transformer: combines binary (1 logit) and multiclass (n_classes logits) heads.
+         * Binary heads: 1 logit → sigmoid → BCE loss
+         * Multiclass heads: n_classes logits → softmax → CE loss
+         * Action vectors use one-hot encoding for multiclass heads.
+         */
+        multi_head_mixed: StructType<{
+            /** Array of head configurations */
+            heads: ArrayType<StructType<{
+                /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+                head_type: VariantType<{
+                    /** Single binary output: 1 logit, sigmoid, BCE loss */
+                    binary: NullType;
+                    /** Multi-class output: n_classes logits, softmax, CE loss */
+                    multiclass: StructType<{
+                        n_classes: IntegerType;
+                    }>;
+                }>;
+                /** Optional class weights for this head */
+                class_weights: OptionType<ArrayType<FloatType>>;
+                /** Optional: index of head this depends on (loss only computed when that head is 1) */
+                conditional_on: OptionType<IntegerType>;
+            }>>;
+        }>;
     }>;
     /** Learning rate (default: 1e-3) */
     learning_rate: OptionType<FloatType>;
@@ -624,6 +816,305 @@ export declare const lightning_generate_sequence: import("@elaraai/east").Platfo
     /** If true, return probabilities. If false, return samples. */
     return_probs: BooleanType;
 }>], ArrayType<ArrayType<FloatType>>>;
+/**
+ * Configuration for Decision Transformer trajectory generation.
+ */
+export declare const TrajectoryGenerateConfigType: StructType<{
+    /** Sampling temperature (0.0 = argmax, > 0 = stochastic) */
+    temperature: FloatType;
+    /** Whether to return probabilities or samples */
+    return_probs: BooleanType;
+    /** Optional constraint mask: (seq_len, action_dim) - FALSE disables action */
+    action_constraints: OptionType<ArrayType<ArrayType<FloatType>>>;
+    /** Optional temporal mask: (seq_len,) - FALSE marks invalid timesteps */
+    temporal_mask: OptionType<ArrayType<FloatType>>;
+    /** Optional head configs for multi_head_mixed output (enables proper multiclass sampling) */
+    head_configs: OptionType<ArrayType<StructType<{
+        /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+        head_type: VariantType<{
+            /** Single binary output: 1 logit, sigmoid, BCE loss */
+            binary: NullType;
+            /** Multi-class output: n_classes logits, softmax, CE loss */
+            multiclass: StructType<{
+                n_classes: IntegerType;
+            }>;
+        }>;
+        /** Optional class weights for this head */
+        class_weights: OptionType<ArrayType<FloatType>>;
+        /** Optional: index of head this depends on (loss only computed when that head is 1) */
+        conditional_on: OptionType<IntegerType>;
+    }>>>;
+    /** Optional action prefix: (seq_len, action_dim) - known actions for timesteps 0..start_timestep-1 */
+    action_prefix: OptionType<ArrayType<ArrayType<FloatType>>>;
+    /** Timestep to start generation from (0 = generate all, 5 = use prefix for 0-4, generate 5+) */
+    start_timestep: OptionType<IntegerType>;
+}>;
+/**
+ * Train with trajectory data for return-conditioned sequence generation.
+ *
+ * Use with decision_transformer architecture.
+ *
+ * @param returns - Return per sample (n_samples,) - actual outcome achieved
+ * @param states - State matrices: n_samples × (seq_len, state_dim)
+ * @param actions - Action matrices: n_samples × (seq_len, action_dim)
+ * @param masks - Temporal masks: n_samples × (seq_len,) - valid timesteps
+ * @param config - Training configuration with decision_transformer architecture
+ * @returns Training result with model blob and metrics
+ */
+export declare const lightning_train_trajectory: import("@elaraai/east").PlatformDefinition<[ArrayType<FloatType>, ArrayType<ArrayType<ArrayType<FloatType>>>, ArrayType<ArrayType<ArrayType<FloatType>>>, ArrayType<ArrayType<FloatType>>, StructType<{
+    /** Model architecture */
+    architecture: VariantType<{
+        /** Simple MLP: input → hidden → output */
+        mlp: StructType<{
+            /** Hidden layer sizes */
+            hidden_layers: ArrayType<IntegerType>;
+        }>;
+        /** Autoencoder: input → encoder → latent → decoder → output */
+        autoencoder: StructType<{
+            /** Encoder hidden layer sizes */
+            encoder_layers: ArrayType<IntegerType>;
+            /** Latent dimension (bottleneck) */
+            latent_dim: IntegerType;
+            /** Decoder hidden layer sizes */
+            decoder_layers: ArrayType<IntegerType>;
+        }>;
+        /** Conv1D: 1D convolutional autoencoder for temporal patterns */
+        conv1d: StructType<{
+            /** Number of channels (e.g., additive types) */
+            n_channels: IntegerType;
+            /** Sequence length (e.g., days) */
+            sequence_length: IntegerType;
+            /** Conv layer channel sizes */
+            conv_channels: ArrayType<IntegerType>;
+            /** Kernel size for convolutions (must be odd) */
+            kernel_size: IntegerType;
+            /** Latent dimension after flattening */
+            latent_dim: IntegerType;
+            /** Optional condition dimension for conditional generation */
+            condition_dim: OptionType<IntegerType>;
+        }>;
+        /** Sequential: LSTM/GRU autoencoder for long-range dependencies */
+        sequential: StructType<{
+            /** Number of channels (e.g., additive types) */
+            n_channels: IntegerType;
+            /** Sequence length (e.g., days) */
+            sequence_length: IntegerType;
+            /** RNN hidden size */
+            hidden_size: IntegerType;
+            /** Number of RNN layers */
+            n_layers: IntegerType;
+            /** Cell type: lstm or gru */
+            cell_type: VariantType<{
+                lstm: NullType;
+                gru: NullType;
+            }>;
+            /** Latent dimension (from final hidden state) */
+            latent_dim: IntegerType;
+            /** Bidirectional encoder (decoder is always unidirectional) */
+            bidirectional: BooleanType;
+            /** Optional condition dimension for conditional generation */
+            condition_dim: OptionType<IntegerType>;
+        }>;
+        /** Transformer: attention-based autoencoder for complex patterns */
+        transformer: StructType<{
+            /** Number of channels (e.g., additive types) */
+            n_channels: IntegerType;
+            /** Sequence length (e.g., days) */
+            sequence_length: IntegerType;
+            /** Model dimension */
+            d_model: IntegerType;
+            /** Number of attention heads (must divide d_model evenly) */
+            n_attention_heads: IntegerType;
+            /** Number of transformer layers */
+            n_layers: IntegerType;
+            /** Feedforward dimension (default: 4 * d_model) */
+            d_ff: OptionType<IntegerType>;
+            /** Latent dimension (mean pooled output) */
+            latent_dim: IntegerType;
+            /** Optional condition dimension for conditional generation */
+            condition_dim: OptionType<IntegerType>;
+        }>;
+        /**
+         * Decision Transformer: return-conditioned sequence generation.
+         * Token layout: [R, s_0, a_0, s_1, a_1, ..., s_{T-1}, a_{T-1}]
+         * Predicts actions conditioned on desired return and state history.
+         */
+        decision_transformer: StructType<{
+            /** Sequence length (timesteps) */
+            sequence_length: IntegerType;
+            /** State dimension per timestep */
+            state_dim: IntegerType;
+            /** Action dimension per timestep */
+            action_dim: IntegerType;
+            /** Model dimension (transformer hidden size) */
+            d_model: IntegerType;
+            /** Number of attention heads */
+            n_attention_heads: IntegerType;
+            /** Number of transformer layers */
+            n_layers: IntegerType;
+            /** Feedforward dimension (default: 4 * d_model) */
+            d_ff: OptionType<IntegerType>;
+            /** Dropout rate */
+            dropout: OptionType<FloatType>;
+            /** Whether return is per-timestep or global */
+            return_embedding: VariantType<{
+                /** Single return value for entire sequence */
+                global: NullType;
+                /** Return-to-go at each timestep */
+                per_timestep: NullType;
+            }>;
+        }>;
+    }>;
+    /** Output mode (determines loss function) */
+    output: VariantType<{
+        /** Regression: MSE loss, no activation */
+        regression: NullType;
+        /** Binary: BCE loss, sigmoid activation */
+        binary: StructType<{
+            /** Optional per-position pos_weights for class imbalance [output_dim] */
+            pos_weight: OptionType<ArrayType<FloatType>>;
+        }>;
+        /** Multiclass: CrossEntropy loss, softmax activation */
+        multiclass: StructType<{
+            /** Number of classes */
+            n_classes: IntegerType;
+            /** Optional per-class weights */
+            class_weights: OptionType<ArrayType<FloatType>>;
+        }>;
+        /** Multi-head categorical: N independent CrossEntropy heads */
+        multi_head: StructType<{
+            /** Number of heads (e.g., 84 time slots) */
+            n_heads: IntegerType;
+            /** Classes per head (e.g., 4 bins) */
+            n_classes_per_head: IntegerType;
+            /** Optional class weights matrix (n_heads, n_classes) */
+            class_weights: OptionType<ArrayType<ArrayType<FloatType>>>;
+        }>;
+        /**
+         * Mixed output types per head.
+         * For Decision Transformer: combines binary (1 logit) and multiclass (n_classes logits) heads.
+         * Binary heads: 1 logit → sigmoid → BCE loss
+         * Multiclass heads: n_classes logits → softmax → CE loss
+         * Action vectors use one-hot encoding for multiclass heads.
+         */
+        multi_head_mixed: StructType<{
+            /** Array of head configurations */
+            heads: ArrayType<StructType<{
+                /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+                head_type: VariantType<{
+                    /** Single binary output: 1 logit, sigmoid, BCE loss */
+                    binary: NullType;
+                    /** Multi-class output: n_classes logits, softmax, CE loss */
+                    multiclass: StructType<{
+                        n_classes: IntegerType;
+                    }>;
+                }>;
+                /** Optional class weights for this head */
+                class_weights: OptionType<ArrayType<FloatType>>;
+                /** Optional: index of head this depends on (loss only computed when that head is 1) */
+                conditional_on: OptionType<IntegerType>;
+            }>>;
+        }>;
+    }>;
+    /** Learning rate (default: 1e-3) */
+    learning_rate: OptionType<FloatType>;
+    /** Maximum epochs (default: 100) */
+    max_epochs: OptionType<IntegerType>;
+    /** Early stopping patience (default: 10) */
+    patience: OptionType<IntegerType>;
+    /** Batch size (default: 32) */
+    batch_size: OptionType<IntegerType>;
+    /** Dropout rate (default: 0.1) */
+    dropout: OptionType<FloatType>;
+    /** Gradient clipping value (default: 1.0) */
+    gradient_clip: OptionType<FloatType>;
+    /** L2 regularization weight decay (default: 0) */
+    weight_decay: OptionType<FloatType>;
+    /** Random seed for reproducibility */
+    random_state: OptionType<IntegerType>;
+    /** Optional callback called each epoch */
+    epoch_callback: OptionType<FunctionType<[IntegerType, FloatType, FloatType], NullType>>;
+}>], StructType<{
+    /** Trained model blob */
+    model: VariantType<{
+        lightning: StructType<{
+            /** Serialized model data (state_dict + hparams) */
+            data: BlobType;
+            /** Input dimension */
+            n_features: IntegerType;
+            /** Output dimension */
+            output_dim: IntegerType;
+            /** Architecture type */
+            architecture_type: StringType;
+            /** Output type */
+            output_type: StringType;
+            /** Latent dimension (autoencoder only) */
+            latent_dim: OptionType<IntegerType>;
+        }>;
+    }>;
+    /** Final training loss */
+    train_loss: FloatType;
+    /** Final validation loss */
+    val_loss: FloatType;
+    /** Best epoch (for early stopping) */
+    best_epoch: IntegerType;
+}>>;
+/**
+ * Generate action sequences autoregressively from trajectory model.
+ *
+ * Use with models trained via trainTrajectory.
+ *
+ * @param model - Trained model from trainTrajectory
+ * @param states - State matrices: n_samples × (seq_len, state_dim)
+ * @param target_returns - Target returns: (n_samples,)
+ * @param config - Generation configuration
+ * @returns Generated actions: n_samples × (seq_len, action_dim)
+ */
+export declare const lightning_generate_trajectory: import("@elaraai/east").PlatformDefinition<[VariantType<{
+    lightning: StructType<{
+        /** Serialized model data (state_dict + hparams) */
+        data: BlobType;
+        /** Input dimension */
+        n_features: IntegerType;
+        /** Output dimension */
+        output_dim: IntegerType;
+        /** Architecture type */
+        architecture_type: StringType;
+        /** Output type */
+        output_type: StringType;
+        /** Latent dimension (autoencoder only) */
+        latent_dim: OptionType<IntegerType>;
+    }>;
+}>, ArrayType<ArrayType<ArrayType<FloatType>>>, ArrayType<FloatType>, StructType<{
+    /** Sampling temperature (0.0 = argmax, > 0 = stochastic) */
+    temperature: FloatType;
+    /** Whether to return probabilities or samples */
+    return_probs: BooleanType;
+    /** Optional constraint mask: (seq_len, action_dim) - FALSE disables action */
+    action_constraints: OptionType<ArrayType<ArrayType<FloatType>>>;
+    /** Optional temporal mask: (seq_len,) - FALSE marks invalid timesteps */
+    temporal_mask: OptionType<ArrayType<FloatType>>;
+    /** Optional head configs for multi_head_mixed output (enables proper multiclass sampling) */
+    head_configs: OptionType<ArrayType<StructType<{
+        /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+        head_type: VariantType<{
+            /** Single binary output: 1 logit, sigmoid, BCE loss */
+            binary: NullType;
+            /** Multi-class output: n_classes logits, softmax, CE loss */
+            multiclass: StructType<{
+                n_classes: IntegerType;
+            }>;
+        }>;
+        /** Optional class weights for this head */
+        class_weights: OptionType<ArrayType<FloatType>>;
+        /** Optional: index of head this depends on (loss only computed when that head is 1) */
+        conditional_on: OptionType<IntegerType>;
+    }>>>;
+    /** Optional action prefix: (seq_len, action_dim) - known actions for timesteps 0..start_timestep-1 */
+    action_prefix: OptionType<ArrayType<ArrayType<FloatType>>>;
+    /** Timestep to start generation from (0 = generate all, 5 = use prefix for 0-4, generate 5+) */
+    start_timestep: OptionType<IntegerType>;
+}>], ArrayType<ArrayType<ArrayType<FloatType>>>>;
 /**
  * Lightning types namespace.
  */
@@ -652,6 +1143,31 @@ export declare const LightningTypes: {
             /** Optional class weights matrix (n_heads, n_classes) */
             class_weights: OptionType<ArrayType<ArrayType<FloatType>>>;
         }>;
+        /**
+         * Mixed output types per head.
+         * For Decision Transformer: combines binary (1 logit) and multiclass (n_classes logits) heads.
+         * Binary heads: 1 logit → sigmoid → BCE loss
+         * Multiclass heads: n_classes logits → softmax → CE loss
+         * Action vectors use one-hot encoding for multiclass heads.
+         */
+        multi_head_mixed: StructType<{
+            /** Array of head configurations */
+            heads: ArrayType<StructType<{
+                /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+                head_type: VariantType<{
+                    /** Single binary output: 1 logit, sigmoid, BCE loss */
+                    binary: NullType;
+                    /** Multi-class output: n_classes logits, softmax, CE loss */
+                    multiclass: StructType<{
+                        n_classes: IntegerType;
+                    }>;
+                }>;
+                /** Optional class weights for this head */
+                class_weights: OptionType<ArrayType<FloatType>>;
+                /** Optional: index of head this depends on (loss only computed when that head is 1) */
+                conditional_on: OptionType<IntegerType>;
+            }>>;
+        }>;
     }>;
     readonly ArchitectureType: VariantType<{
         /** Simple MLP: input → hidden → output */
@@ -724,6 +1240,36 @@ export declare const LightningTypes: {
             /** Optional condition dimension for conditional generation */
             condition_dim: OptionType<IntegerType>;
         }>;
+        /**
+         * Decision Transformer: return-conditioned sequence generation.
+         * Token layout: [R, s_0, a_0, s_1, a_1, ..., s_{T-1}, a_{T-1}]
+         * Predicts actions conditioned on desired return and state history.
+         */
+        decision_transformer: StructType<{
+            /** Sequence length (timesteps) */
+            sequence_length: IntegerType;
+            /** State dimension per timestep */
+            state_dim: IntegerType;
+            /** Action dimension per timestep */
+            action_dim: IntegerType;
+            /** Model dimension (transformer hidden size) */
+            d_model: IntegerType;
+            /** Number of attention heads */
+            n_attention_heads: IntegerType;
+            /** Number of transformer layers */
+            n_layers: IntegerType;
+            /** Feedforward dimension (default: 4 * d_model) */
+            d_ff: OptionType<IntegerType>;
+            /** Dropout rate */
+            dropout: OptionType<FloatType>;
+            /** Whether return is per-timestep or global */
+            return_embedding: VariantType<{
+                /** Single return value for entire sequence */
+                global: NullType;
+                /** Return-to-go at each timestep */
+                per_timestep: NullType;
+            }>;
+        }>;
     }>;
     readonly CellType: VariantType<{
         lstm: NullType;
@@ -803,6 +1349,36 @@ export declare const LightningTypes: {
                 /** Optional condition dimension for conditional generation */
                 condition_dim: OptionType<IntegerType>;
             }>;
+            /**
+             * Decision Transformer: return-conditioned sequence generation.
+             * Token layout: [R, s_0, a_0, s_1, a_1, ..., s_{T-1}, a_{T-1}]
+             * Predicts actions conditioned on desired return and state history.
+             */
+            decision_transformer: StructType<{
+                /** Sequence length (timesteps) */
+                sequence_length: IntegerType;
+                /** State dimension per timestep */
+                state_dim: IntegerType;
+                /** Action dimension per timestep */
+                action_dim: IntegerType;
+                /** Model dimension (transformer hidden size) */
+                d_model: IntegerType;
+                /** Number of attention heads */
+                n_attention_heads: IntegerType;
+                /** Number of transformer layers */
+                n_layers: IntegerType;
+                /** Feedforward dimension (default: 4 * d_model) */
+                d_ff: OptionType<IntegerType>;
+                /** Dropout rate */
+                dropout: OptionType<FloatType>;
+                /** Whether return is per-timestep or global */
+                return_embedding: VariantType<{
+                    /** Single return value for entire sequence */
+                    global: NullType;
+                    /** Return-to-go at each timestep */
+                    per_timestep: NullType;
+                }>;
+            }>;
         }>;
         /** Output mode (determines loss function) */
         output: VariantType<{
@@ -829,6 +1405,31 @@ export declare const LightningTypes: {
                 /** Optional class weights matrix (n_heads, n_classes) */
                 class_weights: OptionType<ArrayType<ArrayType<FloatType>>>;
             }>;
+            /**
+             * Mixed output types per head.
+             * For Decision Transformer: combines binary (1 logit) and multiclass (n_classes logits) heads.
+             * Binary heads: 1 logit → sigmoid → BCE loss
+             * Multiclass heads: n_classes logits → softmax → CE loss
+             * Action vectors use one-hot encoding for multiclass heads.
+             */
+            multi_head_mixed: StructType<{
+                /** Array of head configurations */
+                heads: ArrayType<StructType<{
+                    /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+                    head_type: VariantType<{
+                        /** Single binary output: 1 logit, sigmoid, BCE loss */
+                        binary: NullType;
+                        /** Multi-class output: n_classes logits, softmax, CE loss */
+                        multiclass: StructType<{
+                            n_classes: IntegerType;
+                        }>;
+                    }>;
+                    /** Optional class weights for this head */
+                    class_weights: OptionType<ArrayType<FloatType>>;
+                    /** Optional: index of head this depends on (loss only computed when that head is 1) */
+                    conditional_on: OptionType<IntegerType>;
+                }>>;
+            }>;
         }>;
         /** Learning rate (default: 1e-3) */
         learning_rate: OptionType<FloatType>;
@@ -910,6 +1511,57 @@ export declare const LightningTypes: {
         /** If true, return probabilities. If false, return samples. */
         return_probs: BooleanType;
     }>;
+    readonly ReturnEmbeddingType: VariantType<{
+        /** Single return value for entire sequence */
+        global: NullType;
+        /** Return-to-go at each timestep */
+        per_timestep: NullType;
+    }>;
+    readonly HeadConfigType: StructType<{
+        /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+        head_type: VariantType<{
+            /** Single binary output: 1 logit, sigmoid, BCE loss */
+            binary: NullType;
+            /** Multi-class output: n_classes logits, softmax, CE loss */
+            multiclass: StructType<{
+                n_classes: IntegerType;
+            }>;
+        }>;
+        /** Optional class weights for this head */
+        class_weights: OptionType<ArrayType<FloatType>>;
+        /** Optional: index of head this depends on (loss only computed when that head is 1) */
+        conditional_on: OptionType<IntegerType>;
+    }>;
+    readonly TrajectoryGenerateConfigType: StructType<{
+        /** Sampling temperature (0.0 = argmax, > 0 = stochastic) */
+        temperature: FloatType;
+        /** Whether to return probabilities or samples */
+        return_probs: BooleanType;
+        /** Optional constraint mask: (seq_len, action_dim) - FALSE disables action */
+        action_constraints: OptionType<ArrayType<ArrayType<FloatType>>>;
+        /** Optional temporal mask: (seq_len,) - FALSE marks invalid timesteps */
+        temporal_mask: OptionType<ArrayType<FloatType>>;
+        /** Optional head configs for multi_head_mixed output (enables proper multiclass sampling) */
+        head_configs: OptionType<ArrayType<StructType<{
+            /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+            head_type: VariantType<{
+                /** Single binary output: 1 logit, sigmoid, BCE loss */
+                binary: NullType;
+                /** Multi-class output: n_classes logits, softmax, CE loss */
+                multiclass: StructType<{
+                    n_classes: IntegerType;
+                }>;
+            }>;
+            /** Optional class weights for this head */
+            class_weights: OptionType<ArrayType<FloatType>>;
+            /** Optional: index of head this depends on (loss only computed when that head is 1) */
+            conditional_on: OptionType<IntegerType>;
+        }>>>;
+        /** Optional action prefix: (seq_len, action_dim) - known actions for timesteps 0..start_timestep-1 */
+        action_prefix: OptionType<ArrayType<ArrayType<FloatType>>>;
+        /** Timestep to start generation from (0 = generate all, 5 = use prefix for 0-4, generate 5+) */
+        start_timestep: OptionType<IntegerType>;
+    }>;
 };
 /**
  * Lightning platform functions namespace.
@@ -1015,6 +1667,36 @@ export declare const Lightning: {
                 /** Optional condition dimension for conditional generation */
                 condition_dim: OptionType<IntegerType>;
             }>;
+            /**
+             * Decision Transformer: return-conditioned sequence generation.
+             * Token layout: [R, s_0, a_0, s_1, a_1, ..., s_{T-1}, a_{T-1}]
+             * Predicts actions conditioned on desired return and state history.
+             */
+            decision_transformer: StructType<{
+                /** Sequence length (timesteps) */
+                sequence_length: IntegerType;
+                /** State dimension per timestep */
+                state_dim: IntegerType;
+                /** Action dimension per timestep */
+                action_dim: IntegerType;
+                /** Model dimension (transformer hidden size) */
+                d_model: IntegerType;
+                /** Number of attention heads */
+                n_attention_heads: IntegerType;
+                /** Number of transformer layers */
+                n_layers: IntegerType;
+                /** Feedforward dimension (default: 4 * d_model) */
+                d_ff: OptionType<IntegerType>;
+                /** Dropout rate */
+                dropout: OptionType<FloatType>;
+                /** Whether return is per-timestep or global */
+                return_embedding: VariantType<{
+                    /** Single return value for entire sequence */
+                    global: NullType;
+                    /** Return-to-go at each timestep */
+                    per_timestep: NullType;
+                }>;
+            }>;
         }>;
         /** Output mode (determines loss function) */
         output: VariantType<{
@@ -1041,6 +1723,31 @@ export declare const Lightning: {
                 /** Optional class weights matrix (n_heads, n_classes) */
                 class_weights: OptionType<ArrayType<ArrayType<FloatType>>>;
             }>;
+            /**
+             * Mixed output types per head.
+             * For Decision Transformer: combines binary (1 logit) and multiclass (n_classes logits) heads.
+             * Binary heads: 1 logit → sigmoid → BCE loss
+             * Multiclass heads: n_classes logits → softmax → CE loss
+             * Action vectors use one-hot encoding for multiclass heads.
+             */
+            multi_head_mixed: StructType<{
+                /** Array of head configurations */
+                heads: ArrayType<StructType<{
+                    /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+                    head_type: VariantType<{
+                        /** Single binary output: 1 logit, sigmoid, BCE loss */
+                        binary: NullType;
+                        /** Multi-class output: n_classes logits, softmax, CE loss */
+                        multiclass: StructType<{
+                            n_classes: IntegerType;
+                        }>;
+                    }>;
+                    /** Optional class weights for this head */
+                    class_weights: OptionType<ArrayType<FloatType>>;
+                    /** Optional: index of head this depends on (loss only computed when that head is 1) */
+                    conditional_on: OptionType<IntegerType>;
+                }>>;
+            }>;
         }>;
         /** Learning rate (default: 1e-3) */
         learning_rate: OptionType<FloatType>;
@@ -1210,6 +1917,283 @@ export declare const Lightning: {
         /** If true, return probabilities. If false, return samples. */
         return_probs: BooleanType;
     }>], ArrayType<ArrayType<FloatType>>>;
+    /**
+     * Train a Decision Transformer with trajectory data.
+     *
+     * Trains a return-conditioned sequence generation model that learns
+     * to predict actions given states and desired returns.
+     *
+     * @example
+     * ```typescript
+     * const result = Lightning.trainTrajectory(
+     *     returns, states, actions, masks,
+     *     {
+     *         architecture: variant("decision_transformer", {
+     *             sequence_length: 14n,
+     *             state_dim: 8n,
+     *             action_dim: 11n,
+     *             d_model: 64n,
+     *             n_attention_heads: 4n,
+     *             n_layers: 3n,
+     *             d_ff: variant("none", null),
+     *             dropout: variant("some", 0.1),
+     *             return_embedding: variant("global", null),
+     *         }),
+     *         output: variant("multi_head_mixed", { heads: [...] }),
+     *         ...
+     *     }
+     * );
+     * ```
+     */
+    readonly trainTrajectory: import("@elaraai/east").PlatformDefinition<[ArrayType<FloatType>, ArrayType<ArrayType<ArrayType<FloatType>>>, ArrayType<ArrayType<ArrayType<FloatType>>>, ArrayType<ArrayType<FloatType>>, StructType<{
+        /** Model architecture */
+        architecture: VariantType<{
+            /** Simple MLP: input → hidden → output */
+            mlp: StructType<{
+                /** Hidden layer sizes */
+                hidden_layers: ArrayType<IntegerType>;
+            }>;
+            /** Autoencoder: input → encoder → latent → decoder → output */
+            autoencoder: StructType<{
+                /** Encoder hidden layer sizes */
+                encoder_layers: ArrayType<IntegerType>;
+                /** Latent dimension (bottleneck) */
+                latent_dim: IntegerType;
+                /** Decoder hidden layer sizes */
+                decoder_layers: ArrayType<IntegerType>;
+            }>;
+            /** Conv1D: 1D convolutional autoencoder for temporal patterns */
+            conv1d: StructType<{
+                /** Number of channels (e.g., additive types) */
+                n_channels: IntegerType;
+                /** Sequence length (e.g., days) */
+                sequence_length: IntegerType;
+                /** Conv layer channel sizes */
+                conv_channels: ArrayType<IntegerType>;
+                /** Kernel size for convolutions (must be odd) */
+                kernel_size: IntegerType;
+                /** Latent dimension after flattening */
+                latent_dim: IntegerType;
+                /** Optional condition dimension for conditional generation */
+                condition_dim: OptionType<IntegerType>;
+            }>;
+            /** Sequential: LSTM/GRU autoencoder for long-range dependencies */
+            sequential: StructType<{
+                /** Number of channels (e.g., additive types) */
+                n_channels: IntegerType;
+                /** Sequence length (e.g., days) */
+                sequence_length: IntegerType;
+                /** RNN hidden size */
+                hidden_size: IntegerType;
+                /** Number of RNN layers */
+                n_layers: IntegerType;
+                /** Cell type: lstm or gru */
+                cell_type: VariantType<{
+                    lstm: NullType;
+                    gru: NullType;
+                }>;
+                /** Latent dimension (from final hidden state) */
+                latent_dim: IntegerType;
+                /** Bidirectional encoder (decoder is always unidirectional) */
+                bidirectional: BooleanType;
+                /** Optional condition dimension for conditional generation */
+                condition_dim: OptionType<IntegerType>;
+            }>;
+            /** Transformer: attention-based autoencoder for complex patterns */
+            transformer: StructType<{
+                /** Number of channels (e.g., additive types) */
+                n_channels: IntegerType;
+                /** Sequence length (e.g., days) */
+                sequence_length: IntegerType;
+                /** Model dimension */
+                d_model: IntegerType;
+                /** Number of attention heads (must divide d_model evenly) */
+                n_attention_heads: IntegerType;
+                /** Number of transformer layers */
+                n_layers: IntegerType;
+                /** Feedforward dimension (default: 4 * d_model) */
+                d_ff: OptionType<IntegerType>;
+                /** Latent dimension (mean pooled output) */
+                latent_dim: IntegerType;
+                /** Optional condition dimension for conditional generation */
+                condition_dim: OptionType<IntegerType>;
+            }>;
+            /**
+             * Decision Transformer: return-conditioned sequence generation.
+             * Token layout: [R, s_0, a_0, s_1, a_1, ..., s_{T-1}, a_{T-1}]
+             * Predicts actions conditioned on desired return and state history.
+             */
+            decision_transformer: StructType<{
+                /** Sequence length (timesteps) */
+                sequence_length: IntegerType;
+                /** State dimension per timestep */
+                state_dim: IntegerType;
+                /** Action dimension per timestep */
+                action_dim: IntegerType;
+                /** Model dimension (transformer hidden size) */
+                d_model: IntegerType;
+                /** Number of attention heads */
+                n_attention_heads: IntegerType;
+                /** Number of transformer layers */
+                n_layers: IntegerType;
+                /** Feedforward dimension (default: 4 * d_model) */
+                d_ff: OptionType<IntegerType>;
+                /** Dropout rate */
+                dropout: OptionType<FloatType>;
+                /** Whether return is per-timestep or global */
+                return_embedding: VariantType<{
+                    /** Single return value for entire sequence */
+                    global: NullType;
+                    /** Return-to-go at each timestep */
+                    per_timestep: NullType;
+                }>;
+            }>;
+        }>;
+        /** Output mode (determines loss function) */
+        output: VariantType<{
+            /** Regression: MSE loss, no activation */
+            regression: NullType;
+            /** Binary: BCE loss, sigmoid activation */
+            binary: StructType<{
+                /** Optional per-position pos_weights for class imbalance [output_dim] */
+                pos_weight: OptionType<ArrayType<FloatType>>;
+            }>;
+            /** Multiclass: CrossEntropy loss, softmax activation */
+            multiclass: StructType<{
+                /** Number of classes */
+                n_classes: IntegerType;
+                /** Optional per-class weights */
+                class_weights: OptionType<ArrayType<FloatType>>;
+            }>;
+            /** Multi-head categorical: N independent CrossEntropy heads */
+            multi_head: StructType<{
+                /** Number of heads (e.g., 84 time slots) */
+                n_heads: IntegerType;
+                /** Classes per head (e.g., 4 bins) */
+                n_classes_per_head: IntegerType;
+                /** Optional class weights matrix (n_heads, n_classes) */
+                class_weights: OptionType<ArrayType<ArrayType<FloatType>>>;
+            }>;
+            /**
+             * Mixed output types per head.
+             * For Decision Transformer: combines binary (1 logit) and multiclass (n_classes logits) heads.
+             * Binary heads: 1 logit → sigmoid → BCE loss
+             * Multiclass heads: n_classes logits → softmax → CE loss
+             * Action vectors use one-hot encoding for multiclass heads.
+             */
+            multi_head_mixed: StructType<{
+                /** Array of head configurations */
+                heads: ArrayType<StructType<{
+                    /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+                    head_type: VariantType<{
+                        /** Single binary output: 1 logit, sigmoid, BCE loss */
+                        binary: NullType;
+                        /** Multi-class output: n_classes logits, softmax, CE loss */
+                        multiclass: StructType<{
+                            n_classes: IntegerType;
+                        }>;
+                    }>;
+                    /** Optional class weights for this head */
+                    class_weights: OptionType<ArrayType<FloatType>>;
+                    /** Optional: index of head this depends on (loss only computed when that head is 1) */
+                    conditional_on: OptionType<IntegerType>;
+                }>>;
+            }>;
+        }>;
+        /** Learning rate (default: 1e-3) */
+        learning_rate: OptionType<FloatType>;
+        /** Maximum epochs (default: 100) */
+        max_epochs: OptionType<IntegerType>;
+        /** Early stopping patience (default: 10) */
+        patience: OptionType<IntegerType>;
+        /** Batch size (default: 32) */
+        batch_size: OptionType<IntegerType>;
+        /** Dropout rate (default: 0.1) */
+        dropout: OptionType<FloatType>;
+        /** Gradient clipping value (default: 1.0) */
+        gradient_clip: OptionType<FloatType>;
+        /** L2 regularization weight decay (default: 0) */
+        weight_decay: OptionType<FloatType>;
+        /** Random seed for reproducibility */
+        random_state: OptionType<IntegerType>;
+        /** Optional callback called each epoch */
+        epoch_callback: OptionType<FunctionType<[IntegerType, FloatType, FloatType], NullType>>;
+    }>], StructType<{
+        /** Trained model blob */
+        model: VariantType<{
+            lightning: StructType<{
+                /** Serialized model data (state_dict + hparams) */
+                data: BlobType;
+                /** Input dimension */
+                n_features: IntegerType;
+                /** Output dimension */
+                output_dim: IntegerType;
+                /** Architecture type */
+                architecture_type: StringType;
+                /** Output type */
+                output_type: StringType;
+                /** Latent dimension (autoencoder only) */
+                latent_dim: OptionType<IntegerType>;
+            }>;
+        }>;
+        /** Final training loss */
+        train_loss: FloatType;
+        /** Final validation loss */
+        val_loss: FloatType;
+        /** Best epoch (for early stopping) */
+        best_epoch: IntegerType;
+    }>>;
+    /**
+     * Generate action sequences from a Decision Transformer.
+     *
+     * Autoregressively generates actions conditioned on target returns
+     * and state sequences.
+     */
+    readonly generateTrajectory: import("@elaraai/east").PlatformDefinition<[VariantType<{
+        lightning: StructType<{
+            /** Serialized model data (state_dict + hparams) */
+            data: BlobType;
+            /** Input dimension */
+            n_features: IntegerType;
+            /** Output dimension */
+            output_dim: IntegerType;
+            /** Architecture type */
+            architecture_type: StringType;
+            /** Output type */
+            output_type: StringType;
+            /** Latent dimension (autoencoder only) */
+            latent_dim: OptionType<IntegerType>;
+        }>;
+    }>, ArrayType<ArrayType<ArrayType<FloatType>>>, ArrayType<FloatType>, StructType<{
+        /** Sampling temperature (0.0 = argmax, > 0 = stochastic) */
+        temperature: FloatType;
+        /** Whether to return probabilities or samples */
+        return_probs: BooleanType;
+        /** Optional constraint mask: (seq_len, action_dim) - FALSE disables action */
+        action_constraints: OptionType<ArrayType<ArrayType<FloatType>>>;
+        /** Optional temporal mask: (seq_len,) - FALSE marks invalid timesteps */
+        temporal_mask: OptionType<ArrayType<FloatType>>;
+        /** Optional head configs for multi_head_mixed output (enables proper multiclass sampling) */
+        head_configs: OptionType<ArrayType<StructType<{
+            /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+            head_type: VariantType<{
+                /** Single binary output: 1 logit, sigmoid, BCE loss */
+                binary: NullType;
+                /** Multi-class output: n_classes logits, softmax, CE loss */
+                multiclass: StructType<{
+                    n_classes: IntegerType;
+                }>;
+            }>;
+            /** Optional class weights for this head */
+            class_weights: OptionType<ArrayType<FloatType>>;
+            /** Optional: index of head this depends on (loss only computed when that head is 1) */
+            conditional_on: OptionType<IntegerType>;
+        }>>>;
+        /** Optional action prefix: (seq_len, action_dim) - known actions for timesteps 0..start_timestep-1 */
+        action_prefix: OptionType<ArrayType<ArrayType<FloatType>>>;
+        /** Timestep to start generation from (0 = generate all, 5 = use prefix for 0-4, generate 5+) */
+        start_timestep: OptionType<IntegerType>;
+    }>], ArrayType<ArrayType<ArrayType<FloatType>>>>;
     /**
      * Type definitions for Lightning functions.
      */
@@ -1238,6 +2222,31 @@ export declare const Lightning: {
                 /** Optional class weights matrix (n_heads, n_classes) */
                 class_weights: OptionType<ArrayType<ArrayType<FloatType>>>;
             }>;
+            /**
+             * Mixed output types per head.
+             * For Decision Transformer: combines binary (1 logit) and multiclass (n_classes logits) heads.
+             * Binary heads: 1 logit → sigmoid → BCE loss
+             * Multiclass heads: n_classes logits → softmax → CE loss
+             * Action vectors use one-hot encoding for multiclass heads.
+             */
+            multi_head_mixed: StructType<{
+                /** Array of head configurations */
+                heads: ArrayType<StructType<{
+                    /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+                    head_type: VariantType<{
+                        /** Single binary output: 1 logit, sigmoid, BCE loss */
+                        binary: NullType;
+                        /** Multi-class output: n_classes logits, softmax, CE loss */
+                        multiclass: StructType<{
+                            n_classes: IntegerType;
+                        }>;
+                    }>;
+                    /** Optional class weights for this head */
+                    class_weights: OptionType<ArrayType<FloatType>>;
+                    /** Optional: index of head this depends on (loss only computed when that head is 1) */
+                    conditional_on: OptionType<IntegerType>;
+                }>>;
+            }>;
         }>;
         readonly ArchitectureType: VariantType<{
             /** Simple MLP: input → hidden → output */
@@ -1310,6 +2319,36 @@ export declare const Lightning: {
                 /** Optional condition dimension for conditional generation */
                 condition_dim: OptionType<IntegerType>;
             }>;
+            /**
+             * Decision Transformer: return-conditioned sequence generation.
+             * Token layout: [R, s_0, a_0, s_1, a_1, ..., s_{T-1}, a_{T-1}]
+             * Predicts actions conditioned on desired return and state history.
+             */
+            decision_transformer: StructType<{
+                /** Sequence length (timesteps) */
+                sequence_length: IntegerType;
+                /** State dimension per timestep */
+                state_dim: IntegerType;
+                /** Action dimension per timestep */
+                action_dim: IntegerType;
+                /** Model dimension (transformer hidden size) */
+                d_model: IntegerType;
+                /** Number of attention heads */
+                n_attention_heads: IntegerType;
+                /** Number of transformer layers */
+                n_layers: IntegerType;
+                /** Feedforward dimension (default: 4 * d_model) */
+                d_ff: OptionType<IntegerType>;
+                /** Dropout rate */
+                dropout: OptionType<FloatType>;
+                /** Whether return is per-timestep or global */
+                return_embedding: VariantType<{
+                    /** Single return value for entire sequence */
+                    global: NullType;
+                    /** Return-to-go at each timestep */
+                    per_timestep: NullType;
+                }>;
+            }>;
         }>;
         readonly CellType: VariantType<{
             lstm: NullType;
@@ -1389,6 +2428,36 @@ export declare const Lightning: {
                     /** Optional condition dimension for conditional generation */
                     condition_dim: OptionType<IntegerType>;
                 }>;
+                /**
+                 * Decision Transformer: return-conditioned sequence generation.
+                 * Token layout: [R, s_0, a_0, s_1, a_1, ..., s_{T-1}, a_{T-1}]
+                 * Predicts actions conditioned on desired return and state history.
+                 */
+                decision_transformer: StructType<{
+                    /** Sequence length (timesteps) */
+                    sequence_length: IntegerType;
+                    /** State dimension per timestep */
+                    state_dim: IntegerType;
+                    /** Action dimension per timestep */
+                    action_dim: IntegerType;
+                    /** Model dimension (transformer hidden size) */
+                    d_model: IntegerType;
+                    /** Number of attention heads */
+                    n_attention_heads: IntegerType;
+                    /** Number of transformer layers */
+                    n_layers: IntegerType;
+                    /** Feedforward dimension (default: 4 * d_model) */
+                    d_ff: OptionType<IntegerType>;
+                    /** Dropout rate */
+                    dropout: OptionType<FloatType>;
+                    /** Whether return is per-timestep or global */
+                    return_embedding: VariantType<{
+                        /** Single return value for entire sequence */
+                        global: NullType;
+                        /** Return-to-go at each timestep */
+                        per_timestep: NullType;
+                    }>;
+                }>;
             }>;
             /** Output mode (determines loss function) */
             output: VariantType<{
@@ -1415,6 +2484,31 @@ export declare const Lightning: {
                     /** Optional class weights matrix (n_heads, n_classes) */
                     class_weights: OptionType<ArrayType<ArrayType<FloatType>>>;
                 }>;
+                /**
+                 * Mixed output types per head.
+                 * For Decision Transformer: combines binary (1 logit) and multiclass (n_classes logits) heads.
+                 * Binary heads: 1 logit → sigmoid → BCE loss
+                 * Multiclass heads: n_classes logits → softmax → CE loss
+                 * Action vectors use one-hot encoding for multiclass heads.
+                 */
+                multi_head_mixed: StructType<{
+                    /** Array of head configurations */
+                    heads: ArrayType<StructType<{
+                        /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+                        head_type: VariantType<{
+                            /** Single binary output: 1 logit, sigmoid, BCE loss */
+                            binary: NullType;
+                            /** Multi-class output: n_classes logits, softmax, CE loss */
+                            multiclass: StructType<{
+                                n_classes: IntegerType;
+                            }>;
+                        }>;
+                        /** Optional class weights for this head */
+                        class_weights: OptionType<ArrayType<FloatType>>;
+                        /** Optional: index of head this depends on (loss only computed when that head is 1) */
+                        conditional_on: OptionType<IntegerType>;
+                    }>>;
+                }>;
             }>;
             /** Learning rate (default: 1e-3) */
             learning_rate: OptionType<FloatType>;
@@ -1496,6 +2590,57 @@ export declare const Lightning: {
             /** If true, return probabilities. If false, return samples. */
             return_probs: BooleanType;
         }>;
+        readonly ReturnEmbeddingType: VariantType<{
+            /** Single return value for entire sequence */
+            global: NullType;
+            /** Return-to-go at each timestep */
+            per_timestep: NullType;
+        }>;
+        readonly HeadConfigType: StructType<{
+            /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+            head_type: VariantType<{
+                /** Single binary output: 1 logit, sigmoid, BCE loss */
+                binary: NullType;
+                /** Multi-class output: n_classes logits, softmax, CE loss */
+                multiclass: StructType<{
+                    n_classes: IntegerType;
+                }>;
+            }>;
+            /** Optional class weights for this head */
+            class_weights: OptionType<ArrayType<FloatType>>;
+            /** Optional: index of head this depends on (loss only computed when that head is 1) */
+            conditional_on: OptionType<IntegerType>;
+        }>;
+        readonly TrajectoryGenerateConfigType: StructType<{
+            /** Sampling temperature (0.0 = argmax, > 0 = stochastic) */
+            temperature: FloatType;
+            /** Whether to return probabilities or samples */
+            return_probs: BooleanType;
+            /** Optional constraint mask: (seq_len, action_dim) - FALSE disables action */
+            action_constraints: OptionType<ArrayType<ArrayType<FloatType>>>;
+            /** Optional temporal mask: (seq_len,) - FALSE marks invalid timesteps */
+            temporal_mask: OptionType<ArrayType<FloatType>>;
+            /** Optional head configs for multi_head_mixed output (enables proper multiclass sampling) */
+            head_configs: OptionType<ArrayType<StructType<{
+                /** Output type: binary (1 logit, sigmoid, BCE) or multiclass (n_classes logits, softmax, CE) */
+                head_type: VariantType<{
+                    /** Single binary output: 1 logit, sigmoid, BCE loss */
+                    binary: NullType;
+                    /** Multi-class output: n_classes logits, softmax, CE loss */
+                    multiclass: StructType<{
+                        n_classes: IntegerType;
+                    }>;
+                }>;
+                /** Optional class weights for this head */
+                class_weights: OptionType<ArrayType<FloatType>>;
+                /** Optional: index of head this depends on (loss only computed when that head is 1) */
+                conditional_on: OptionType<IntegerType>;
+            }>>>;
+            /** Optional action prefix: (seq_len, action_dim) - known actions for timesteps 0..start_timestep-1 */
+            action_prefix: OptionType<ArrayType<ArrayType<FloatType>>>;
+            /** Timestep to start generation from (0 = generate all, 5 = use prefix for 0-4, generate 5+) */
+            start_timestep: OptionType<IntegerType>;
+        }>;
     };
 };
 //# sourceMappingURL=lightning.d.ts.map