Skip to content

Commit a2878a8

Browse files
speedstorm1copybara-github
speedstorm1
authored and
copybara-github
committed
feat: Support Reinforcement Tuning in GenAI SDK
PiperOrigin-RevId: 919226756
1 parent 8d2a892 commit a2878a8

8 files changed

Lines changed: 771 additions & 5 deletions

‎src/main/java/com/google/genai/Tunings.java‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,42 @@ ObjectNode createTuningJobConfigToMldev(
247247
+ " in Gemini Developer API mode.");
248248
}
249249

250+
if (!Common.isZero(Common.getValueByPath(fromObject, new String[] {"rewardConfig"}))) {
251+
throw new IllegalArgumentException(
252+
"rewardConfig parameter is only supported in Gemini Enterprise Agent Platform mode, not"
253+
+ " in Gemini Developer API mode.");
254+
}
255+
256+
if (!Common.isZero(Common.getValueByPath(fromObject, new String[] {"compositeRewardConfig"}))) {
257+
throw new IllegalArgumentException(
258+
"compositeRewardConfig parameter is only supported in Gemini Enterprise Agent Platform"
259+
+ " mode, not in Gemini Developer API mode.");
260+
}
261+
262+
if (!Common.isZero(Common.getValueByPath(fromObject, new String[] {"samplesPerPrompt"}))) {
263+
throw new IllegalArgumentException(
264+
"samplesPerPrompt parameter is only supported in Gemini Enterprise Agent Platform mode,"
265+
+ " not in Gemini Developer API mode.");
266+
}
267+
268+
if (!Common.isZero(Common.getValueByPath(fromObject, new String[] {"evaluateInterval"}))) {
269+
throw new IllegalArgumentException(
270+
"evaluateInterval parameter is only supported in Gemini Enterprise Agent Platform mode,"
271+
+ " not in Gemini Developer API mode.");
272+
}
273+
274+
if (!Common.isZero(Common.getValueByPath(fromObject, new String[] {"checkpointInterval"}))) {
275+
throw new IllegalArgumentException(
276+
"checkpointInterval parameter is only supported in Gemini Enterprise Agent Platform mode,"
277+
+ " not in Gemini Developer API mode.");
278+
}
279+
280+
if (!Common.isZero(Common.getValueByPath(fromObject, new String[] {"maxOutputTokens"}))) {
281+
throw new IllegalArgumentException(
282+
"maxOutputTokens parameter is only supported in Gemini Enterprise Agent Platform mode,"
283+
+ " not in Gemini Developer API mode.");
284+
}
285+
250286
return toObject;
251287
}
252288

@@ -294,6 +330,17 @@ ObjectNode createTuningJobConfigToVertex(
294330
toObject,
295331
rootObject));
296332
}
333+
} else if (discriminatorValueValidationDataset.equals("REINFORCEMENT_TUNING")) {
334+
if (Common.getValueByPath(fromObject, new String[] {"validationDataset"}) != null) {
335+
Common.setValueByPath(
336+
parentObject,
337+
new String[] {"reinforcementTuningSpec"},
338+
tuningValidationDatasetToVertex(
339+
JsonSerializable.toJsonNode(
340+
Common.getValueByPath(fromObject, new String[] {"validationDataset"})),
341+
toObject,
342+
rootObject));
343+
}
297344
}
298345
if (Common.getValueByPath(fromObject, new String[] {"tunedModelDisplayName"}) != null) {
299346
Common.setValueByPath(
@@ -336,6 +383,13 @@ ObjectNode createTuningJobConfigToVertex(
336383
new String[] {"distillationSpec", "hyperParameters", "epochCount"},
337384
Common.getValueByPath(fromObject, new String[] {"epochCount"}));
338385
}
386+
} else if (discriminatorValueEpochCount.equals("REINFORCEMENT_TUNING")) {
387+
if (Common.getValueByPath(fromObject, new String[] {"epochCount"}) != null) {
388+
Common.setValueByPath(
389+
parentObject,
390+
new String[] {"reinforcementTuningSpec", "hyperParameters", "epochCount"},
391+
Common.getValueByPath(fromObject, new String[] {"epochCount"}));
392+
}
339393
}
340394

341395
JsonNode discriminatorLearningRateMultiplier =
@@ -367,6 +421,13 @@ ObjectNode createTuningJobConfigToVertex(
367421
new String[] {"distillationSpec", "hyperParameters", "learningRateMultiplier"},
368422
Common.getValueByPath(fromObject, new String[] {"learningRateMultiplier"}));
369423
}
424+
} else if (discriminatorValueLearningRateMultiplier.equals("REINFORCEMENT_TUNING")) {
425+
if (Common.getValueByPath(fromObject, new String[] {"learningRateMultiplier"}) != null) {
426+
Common.setValueByPath(
427+
parentObject,
428+
new String[] {"reinforcementTuningSpec", "hyperParameters", "learningRateMultiplier"},
429+
Common.getValueByPath(fromObject, new String[] {"learningRateMultiplier"}));
430+
}
370431
}
371432

372433
JsonNode discriminatorExportLastCheckpointOnly =
@@ -425,6 +486,13 @@ ObjectNode createTuningJobConfigToVertex(
425486
new String[] {"distillationSpec", "hyperParameters", "adapterSize"},
426487
Common.getValueByPath(fromObject, new String[] {"adapterSize"}));
427488
}
489+
} else if (discriminatorValueAdapterSize.equals("REINFORCEMENT_TUNING")) {
490+
if (Common.getValueByPath(fromObject, new String[] {"adapterSize"}) != null) {
491+
Common.setValueByPath(
492+
parentObject,
493+
new String[] {"reinforcementTuningSpec", "hyperParameters", "adapterSize"},
494+
Common.getValueByPath(fromObject, new String[] {"adapterSize"}));
495+
}
428496
}
429497

430498
JsonNode discriminatorTuningMode =
@@ -473,6 +541,13 @@ ObjectNode createTuningJobConfigToVertex(
473541
new String[] {"distillationSpec", "hyperParameters", "batchSize"},
474542
Common.getValueByPath(fromObject, new String[] {"batchSize"}));
475543
}
544+
} else if (discriminatorValueBatchSize.equals("REINFORCEMENT_TUNING")) {
545+
if (Common.getValueByPath(fromObject, new String[] {"batchSize"}) != null) {
546+
Common.setValueByPath(
547+
parentObject,
548+
new String[] {"reinforcementTuningSpec", "hyperParameters", "batchSize"},
549+
Common.getValueByPath(fromObject, new String[] {"batchSize"}));
550+
}
476551
}
477552

478553
JsonNode discriminatorLearningRate =
@@ -586,6 +661,48 @@ ObjectNode createTuningJobConfigToVertex(
586661
Common.getValueByPath(fromObject, new String[] {"encryptionSpec"}));
587662
}
588663

664+
if (Common.getValueByPath(fromObject, new String[] {"rewardConfig"}) != null) {
665+
Common.setValueByPath(
666+
parentObject,
667+
new String[] {"reinforcementTuningSpec", "singleRewardConfig"},
668+
Common.getValueByPath(fromObject, new String[] {"rewardConfig"}));
669+
}
670+
671+
if (Common.getValueByPath(fromObject, new String[] {"compositeRewardConfig"}) != null) {
672+
Common.setValueByPath(
673+
parentObject,
674+
new String[] {"reinforcementTuningSpec", "compositeRewardConfig"},
675+
Common.getValueByPath(fromObject, new String[] {"compositeRewardConfig"}));
676+
}
677+
678+
if (Common.getValueByPath(fromObject, new String[] {"samplesPerPrompt"}) != null) {
679+
Common.setValueByPath(
680+
parentObject,
681+
new String[] {"reinforcementTuningSpec", "hyperParameters", "samplesPerPrompt"},
682+
Common.getValueByPath(fromObject, new String[] {"samplesPerPrompt"}));
683+
}
684+
685+
if (Common.getValueByPath(fromObject, new String[] {"evaluateInterval"}) != null) {
686+
Common.setValueByPath(
687+
parentObject,
688+
new String[] {"reinforcementTuningSpec", "hyperParameters", "evaluateInterval"},
689+
Common.getValueByPath(fromObject, new String[] {"evaluateInterval"}));
690+
}
691+
692+
if (Common.getValueByPath(fromObject, new String[] {"checkpointInterval"}) != null) {
693+
Common.setValueByPath(
694+
parentObject,
695+
new String[] {"reinforcementTuningSpec", "hyperParameters", "checkpointInterval"},
696+
Common.getValueByPath(fromObject, new String[] {"checkpointInterval"}));
697+
}
698+
699+
if (Common.getValueByPath(fromObject, new String[] {"maxOutputTokens"}) != null) {
700+
Common.setValueByPath(
701+
parentObject,
702+
new String[] {"reinforcementTuningSpec", "hyperParameters", "maxOutputTokens"},
703+
Common.getValueByPath(fromObject, new String[] {"maxOutputTokens"}));
704+
}
705+
589706
return toObject;
590707
}
591708

@@ -1430,6 +1547,13 @@ ObjectNode tuningDatasetToVertex(
14301547
new String[] {"distillationSpec", "promptDatasetUri"},
14311548
Common.getValueByPath(fromObject, new String[] {"gcsUri"}));
14321549
}
1550+
} else if (discriminatorValueGcsUri.equals("REINFORCEMENT_TUNING")) {
1551+
if (Common.getValueByPath(fromObject, new String[] {"gcsUri"}) != null) {
1552+
Common.setValueByPath(
1553+
parentObject,
1554+
new String[] {"reinforcementTuningSpec", "trainingDatasetUri"},
1555+
Common.getValueByPath(fromObject, new String[] {"gcsUri"}));
1556+
}
14331557
}
14341558

14351559
JsonNode discriminatorVertexDatasetResource =
@@ -1459,6 +1583,13 @@ ObjectNode tuningDatasetToVertex(
14591583
new String[] {"distillationSpec", "promptDatasetUri"},
14601584
Common.getValueByPath(fromObject, new String[] {"vertexDatasetResource"}));
14611585
}
1586+
} else if (discriminatorValueVertexDatasetResource.equals("REINFORCEMENT_TUNING")) {
1587+
if (Common.getValueByPath(fromObject, new String[] {"vertexDatasetResource"}) != null) {
1588+
Common.setValueByPath(
1589+
parentObject,
1590+
new String[] {"reinforcementTuningSpec", "trainingDatasetUri"},
1591+
Common.getValueByPath(fromObject, new String[] {"vertexDatasetResource"}));
1592+
}
14621593
}
14631594
if (!Common.isZero(Common.getValueByPath(fromObject, new String[] {"examples"}))) {
14641595
throw new IllegalArgumentException(

‎src/main/java/com/google/genai/types/AutoraterConfig.java‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,17 @@ public abstract class AutoraterConfig extends JsonSerializable {
5252
@JsonProperty("generationConfig")
5353
public abstract Optional<GenerationConfig> generationConfig();
5454

55+
/**
56+
* Optional. Default is true. Whether to flip the candidate and baseline responses. This is only
57+
* applicable to the pairwise metric. If enabled, also provide
58+
* PairwiseMetricSpec.candidate_response_field_name and
59+
* PairwiseMetricSpec.baseline_response_field_name. When rendering
60+
* PairwiseMetricSpec.metric_prompt_template, the candidate and baseline fields will be flipped
61+
* for half of the samples to reduce bias.
62+
*/
63+
@JsonProperty("flipEnabled")
64+
public abstract Optional<Boolean> flipEnabled();
65+
5566
/** Instantiates a builder for AutoraterConfig. */
5667
@ExcludeFromGeneratedCoverageReport
5768
public static Builder builder() {
@@ -142,6 +153,29 @@ public Builder clearGenerationConfig() {
142153
return generationConfig(Optional.empty());
143154
}
144155

156+
/**
157+
* Setter for flipEnabled.
158+
*
159+
* <p>flipEnabled: Optional. Default is true. Whether to flip the candidate and baseline
160+
* responses. This is only applicable to the pairwise metric. If enabled, also provide
161+
* PairwiseMetricSpec.candidate_response_field_name and
162+
* PairwiseMetricSpec.baseline_response_field_name. When rendering
163+
* PairwiseMetricSpec.metric_prompt_template, the candidate and baseline fields will be flipped
164+
* for half of the samples to reduce bias.
165+
*/
166+
@JsonProperty("flipEnabled")
167+
public abstract Builder flipEnabled(boolean flipEnabled);
168+
169+
@ExcludeFromGeneratedCoverageReport
170+
abstract Builder flipEnabled(Optional<Boolean> flipEnabled);
171+
172+
/** Clears the value of flipEnabled field. */
173+
@ExcludeFromGeneratedCoverageReport
174+
@CanIgnoreReturnValue
175+
public Builder clearFlipEnabled() {
176+
return flipEnabled(Optional.empty());
177+
}
178+
145179
public abstract AutoraterConfig build();
146180
}
147181

‎src/main/java/com/google/genai/types/CompositeReinforcementTuningRewardConfig.java‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
// Auto-generated code. Do not edit.
18+
19+
package com.google.genai.types;
20+
21+
import static com.google.common.collect.ImmutableList.toImmutableList;
22+
23+
import com.fasterxml.jackson.annotation.JsonCreator;
24+
import com.fasterxml.jackson.annotation.JsonProperty;
25+
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
26+
import com.google.auto.value.AutoValue;
27+
import com.google.errorprone.annotations.CanIgnoreReturnValue;
28+
import com.google.genai.JsonSerializable;
29+
import java.util.Arrays;
30+
import java.util.List;
31+
import java.util.Optional;
32+
33+
/** Composite reinforcement tuning reward config. */
34+
@AutoValue
35+
@JsonDeserialize(builder = CompositeReinforcementTuningRewardConfig.Builder.class)
36+
public abstract class CompositeReinforcementTuningRewardConfig extends JsonSerializable {
37+
/** */
38+
@JsonProperty("weightedRewardConfigs")
39+
public abstract Optional<List<CompositeReinforcementTuningRewardConfigWeightedRewardConfig>>
40+
weightedRewardConfigs();
41+
42+
/** Instantiates a builder for CompositeReinforcementTuningRewardConfig. */
43+
@ExcludeFromGeneratedCoverageReport
44+
public static Builder builder() {
45+
return new AutoValue_CompositeReinforcementTuningRewardConfig.Builder();
46+
}
47+
48+
/** Creates a builder with the same values as this instance. */
49+
public abstract Builder toBuilder();
50+
51+
/** Builder for CompositeReinforcementTuningRewardConfig. */
52+
@AutoValue.Builder
53+
public abstract static class Builder {
54+
/**
55+
* For internal usage. Please use `CompositeReinforcementTuningRewardConfig.builder()` for
56+
* instantiation.
57+
*/
58+
@JsonCreator
59+
private static Builder create() {
60+
return new AutoValue_CompositeReinforcementTuningRewardConfig.Builder();
61+
}
62+
63+
/**
64+
* Setter for weightedRewardConfigs.
65+
*
66+
* <p>weightedRewardConfigs:
67+
*/
68+
@JsonProperty("weightedRewardConfigs")
69+
public abstract Builder weightedRewardConfigs(
70+
List<CompositeReinforcementTuningRewardConfigWeightedRewardConfig> weightedRewardConfigs);
71+
72+
/**
73+
* Setter for weightedRewardConfigs.
74+
*
75+
* <p>weightedRewardConfigs:
76+
*/
77+
@CanIgnoreReturnValue
78+
public Builder weightedRewardConfigs(
79+
CompositeReinforcementTuningRewardConfigWeightedRewardConfig... weightedRewardConfigs) {
80+
return weightedRewardConfigs(Arrays.asList(weightedRewardConfigs));
81+
}
82+
83+
/**
84+
* Setter for weightedRewardConfigs builder.
85+
*
86+
* <p>weightedRewardConfigs:
87+
*/
88+
@CanIgnoreReturnValue
89+
public Builder weightedRewardConfigs(
90+
CompositeReinforcementTuningRewardConfigWeightedRewardConfig.Builder...
91+
weightedRewardConfigsBuilders) {
92+
return weightedRewardConfigs(
93+
Arrays.asList(weightedRewardConfigsBuilders).stream()
94+
.map(CompositeReinforcementTuningRewardConfigWeightedRewardConfig.Builder::build)
95+
.collect(toImmutableList()));
96+
}
97+
98+
@ExcludeFromGeneratedCoverageReport
99+
abstract Builder weightedRewardConfigs(
100+
Optional<List<CompositeReinforcementTuningRewardConfigWeightedRewardConfig>>
101+
weightedRewardConfigs);
102+
103+
/** Clears the value of weightedRewardConfigs field. */
104+
@ExcludeFromGeneratedCoverageReport
105+
@CanIgnoreReturnValue
106+
public Builder clearWeightedRewardConfigs() {
107+
return weightedRewardConfigs(Optional.empty());
108+
}
109+
110+
public abstract CompositeReinforcementTuningRewardConfig build();
111+
}
112+
113+
/** Deserializes a JSON string to a CompositeReinforcementTuningRewardConfig object. */
114+
@ExcludeFromGeneratedCoverageReport
115+
public static CompositeReinforcementTuningRewardConfig fromJson(String jsonString) {
116+
return JsonSerializable.fromJsonString(
117+
jsonString, CompositeReinforcementTuningRewardConfig.class);
118+
}
119+
}

0 commit comments

Comments
 (0)